internode.cu 106 KB
Newer Older
lijian6's avatar
lijian6 committed
1
#include "hip/hip_runtime.h"
Chenggang Zhao's avatar
Chenggang Zhao committed
2
#include "buffer.cuh"
lijian6's avatar
lijian6 committed
3
#include "configs.cuh"
Chenggang Zhao's avatar
Chenggang Zhao committed
4
5
#include "launch.cuh"
#include "utils.cuh"
6
#include "shmem_wrapper.cuh"
lijian6's avatar
lijian6 committed
7
8
9

#ifndef DISABLE_ROCSHMEM

Chenggang Zhao's avatar
Chenggang Zhao committed
10
11
12
13
namespace deep_ep {

namespace internode {

14
extern shmem_team_t cpu_rdma_team;
Chenggang Zhao's avatar
Chenggang Zhao committed
15
16

struct SourceMeta {
root's avatar
root committed
17
    int src_rdma_rank, is_token_in_nvl_rank_bits;
Chenggang Zhao's avatar
Chenggang Zhao committed
18
19
20
21
22
23

    EP_STATIC_ASSERT(NUM_MAX_NVL_PEERS == 8, "Invalid number of maximum NVL peers");

    __forceinline__ SourceMeta() = default;

    // TODO: faster encoding
lijian6's avatar
lijian6 committed
24
25
    __device__ __forceinline__ SourceMeta(int rdma_rank, const bool *is_token_in_nvl_ranks) {
        src_rdma_rank             = rdma_rank;
Chenggang Zhao's avatar
Chenggang Zhao committed
26
        is_token_in_nvl_rank_bits = is_token_in_nvl_ranks[0];
lijian6's avatar
lijian6 committed
27
28
#pragma unroll
        for (int i = 1; i < NUM_MAX_NVL_PEERS; ++i)
Chenggang Zhao's avatar
Chenggang Zhao committed
29
30
31
32
33
34
35
36
37
38
39
40
            is_token_in_nvl_rank_bits |= is_token_in_nvl_ranks[i] << i;
    }

    __device__ __forceinline__ bool is_token_in_nvl_rank(int nvl_rank) const {
        return (is_token_in_nvl_rank_bits >> nvl_rank) & 1;
    }
};

int get_source_meta_bytes() {
    return sizeof(SourceMeta);
}

lijian6's avatar
lijian6 committed
41
42
43
44
45
__host__ __device__ __forceinline__ int get_num_bytes_per_rdma_token(int hidden_int4,
                                                                     int num_scales,
                                                                     int num_topk_idx,
                                                                     int num_topk_weights) {
    return static_cast<int>(ALIGN(hidden_int4 * sizeof(int4) + sizeof(SourceMeta) +
46
47
                                  num_scales * sizeof(float) + num_topk_idx * sizeof(int) +
                                  num_topk_weights * sizeof(float), sizeof(int4)));
Chenggang Zhao's avatar
Chenggang Zhao committed
48
49
}

lijian6's avatar
lijian6 committed
50
51
__host__ __device__ __forceinline__ std::pair<int, int>
get_rdma_clean_meta(int hidden_int4, int num_scales, int num_topk_idx, int num_topk_weights,
root's avatar
root committed
52
                    int num_rdma_ranks, int num_rdma_recv_buffer_tokens, int num_sms) {
lijian6's avatar
lijian6 committed
53
    // Return `int32_t` offset and count to clean
lijian6's avatar
lijian6 committed
54
    return {(get_num_bytes_per_rdma_token(hidden_int4, num_scales, num_topk_idx, num_topk_weights) *
root's avatar
root committed
55
56
           num_rdma_recv_buffer_tokens * num_rdma_ranks * 2 * num_sms) / sizeof(int),
           (NUM_MAX_NVL_PEERS * 2 + 4) * num_rdma_ranks * 2 * num_sms};
Chenggang Zhao's avatar
Chenggang Zhao committed
57
}
lijian6's avatar
lijian6 committed
58
59
60
__host__ __device__ __forceinline__ std::pair<int, int>
get_nvl_clean_meta(int hidden_int4, int num_scales, int num_topk_idx, int num_topk_weights,
                   int num_rdma_ranks, int num_nvl_ranks, int num_nvl_recv_buffer_tokens,
root's avatar
root committed
61
                   int num_sms) {
Chenggang Zhao's avatar
Chenggang Zhao committed
62
    // Return `int32_t` offset and to clean
lijian6's avatar
lijian6 committed
63
64
    EP_STATIC_ASSERT(sizeof(SourceMeta) % sizeof(int) == 0,
                              "Invalid size of `SourceMeta`");
Chenggang Zhao's avatar
Chenggang Zhao committed
65
    return {
lijian6's avatar
lijian6 committed
66
        (num_nvl_recv_buffer_tokens *
67
68
        (hidden_int4 * sizeof(int4) + num_scales * sizeof(float) + num_topk_idx * sizeof(int) +
        num_topk_weights * sizeof(float) + sizeof(SourceMeta)) *
root's avatar
root committed
69
70
        num_nvl_ranks * num_sms) / sizeof(int),
        num_nvl_ranks * (2 * num_rdma_ranks + 2) * num_sms,
Chenggang Zhao's avatar
Chenggang Zhao committed
71
72
73
74
    };
}

template <bool kLowLatencyMode>
lijian6's avatar
lijian6 committed
75
76
__forceinline__ __device__ int translate_dst_rdma_rank(const int dst_rdma_rank,
                                                       const int nvl_rank) {
Chenggang Zhao's avatar
Chenggang Zhao committed
77
78
79
80
    return kLowLatencyMode ? (dst_rdma_rank * NUM_MAX_NVL_PEERS + nvl_rank) : dst_rdma_rank;
}

template <bool kLowLatencyMode>
lijian6's avatar
lijian6 committed
81
__forceinline__ __device__ void
root's avatar
root committed
82
shmem_sync_with_same_gpu_idx(const shmem_team_t &rdma_team) {
lijian6's avatar
lijian6 committed
83
84
    // NOTE: shmem_device_barrier_all() might be an issue as
    // it doesn't follow OpenSHMEM specification on ROCm
root's avatar
root committed
85
    // kLowLatencyMode ? shmem_device_sync(rdma_team) : shmem_device_sync_all();
86
    kLowLatencyMode ? shmem_barrier(rdma_team) : shmem_device_barrier_all();
Chenggang Zhao's avatar
Chenggang Zhao committed
87
88
89
90
}

template <bool kLowLatencyMode, int kNumRDMARanks>
__global__ void
lijian6's avatar
lijian6 committed
91
92
93
94
95
notify_dispatch(const int *num_tokens_per_rank, int *moe_recv_counter_mapped, int num_ranks,
                const int *num_tokens_per_rdma_rank, int *moe_recv_rdma_counter_mapped,
                const int *num_tokens_per_expert, int *moe_recv_expert_counter_mapped,
                int num_experts, const bool *is_token_in_rank, int num_tokens, int num_channels,
                int expert_alignment, const int rdma_clean_offset, const int rdma_num_int_clean,
Chenggang Zhao's avatar
Chenggang Zhao committed
96
                const int nvl_clean_offset, const int nvl_num_int_clean,
lijian6's avatar
lijian6 committed
97
98
99
                int *rdma_channel_prefix_matrix, int *recv_rdma_rank_prefix_sum,
                int *gbl_channel_prefix_matrix, int *recv_gbl_rank_prefix_sum,
                void *rdma_buffer_ptr, void **buffer_ptrs, int **barrier_signal_ptrs, int rank,
100
                const shmem_team_t rdma_team) {
lijian6's avatar
lijian6 committed
101
102
103
104
    auto sm_id     = static_cast<int>(blockIdx.x);
    auto thread_id = static_cast<int>(threadIdx.x), warp_id = thread_id / kWarpSize,
         lane_id     = get_lane_id();
    auto num_threads = static_cast<int>(blockDim.x), num_warps = num_threads / kWarpSize;
Chenggang Zhao's avatar
Chenggang Zhao committed
105
106

    auto rdma_rank = rank / NUM_MAX_NVL_PEERS, nvl_rank = rank % NUM_MAX_NVL_PEERS;
lijian6's avatar
lijian6 committed
107
108
    auto num_rdma_experts = num_experts / kNumRDMARanks,
         num_nvl_experts  = num_rdma_experts / NUM_MAX_NVL_PEERS;
Chenggang Zhao's avatar
Chenggang Zhao committed
109
110
111

    if (sm_id == 0) {
        // Communication with others
lijian6's avatar
lijian6 committed
112
113
        // Global barrier: the first warp do intra-node sync, the second warp do internode sync
        if (thread_id == kWarpSize)
root's avatar
root committed
114
            shmem_sync_with_same_gpu_idx<kLowLatencyMode>(rdma_team);
115

lijian6's avatar
lijian6 committed
116
        barrier_block<NUM_MAX_NVL_PEERS>(barrier_signal_ptrs, nvl_rank);
117

Chenggang Zhao's avatar
Chenggang Zhao committed
118
        // Send numbers of tokens per rank/expert to RDMA ranks
lijian6's avatar
lijian6 committed
119
120
121
        auto rdma_buffer_ptr_int        = reinterpret_cast<int *>(rdma_buffer_ptr);
        auto rdma_recv_num_tokens_mixed = SymBuffer<int>(
            rdma_buffer_ptr, NUM_MAX_NVL_PEERS + num_rdma_experts + 1, kNumRDMARanks);
Chenggang Zhao's avatar
Chenggang Zhao committed
122
123

        // Clean up for later data dispatch
lijian6's avatar
lijian6 committed
124
125
        EP_DEVICE_ASSERT(rdma_recv_num_tokens_mixed.total_bytes <=
                                  rdma_clean_offset * sizeof(int));
Chenggang Zhao's avatar
Chenggang Zhao committed
126
127
128
129
130
        for (int i = thread_id; i < rdma_num_int_clean; i += num_threads)
            rdma_buffer_ptr_int[rdma_clean_offset + i] = 0;

        // Copy to send buffer
        for (int i = thread_id; i < num_ranks; i += num_threads)
lijian6's avatar
lijian6 committed
131
132
            rdma_recv_num_tokens_mixed.send_buffer(i / NUM_MAX_NVL_PEERS)[i % NUM_MAX_NVL_PEERS] =
                num_tokens_per_rank[i];
Chenggang Zhao's avatar
Chenggang Zhao committed
133
        for (int i = thread_id; i < num_experts; i += num_threads)
lijian6's avatar
lijian6 committed
134
135
136
            rdma_recv_num_tokens_mixed.send_buffer(
                i / num_rdma_experts)[NUM_MAX_NVL_PEERS + i % num_rdma_experts] =
                num_tokens_per_expert[i];
Chenggang Zhao's avatar
Chenggang Zhao committed
137
        if (thread_id < kNumRDMARanks)
lijian6's avatar
lijian6 committed
138
139
140
            rdma_recv_num_tokens_mixed.send_buffer(
                thread_id)[NUM_MAX_NVL_PEERS + num_rdma_experts] =
                num_tokens_per_rdma_rank[thread_id];
Chenggang Zhao's avatar
Chenggang Zhao committed
141
142
143
144
145
        __syncthreads();

        // Issue send
        // TODO: more light fence or barrier or signaling
        // TODO: overlap EP barrier and NVL cleaning
lishen's avatar
lishen committed
146
147
148
        for (int i = warp_id; i < kNumRDMARanks; i += num_warps) {
            if (i != rdma_rank) {
                shmemx_int_put_nbi_warp(
lijian6's avatar
lijian6 committed
149
                rdma_recv_num_tokens_mixed.recv_buffer(rdma_rank),
lishen's avatar
lishen committed
150
                rdma_recv_num_tokens_mixed.send_buffer(i),
lijian6's avatar
lijian6 committed
151
                NUM_MAX_NVL_PEERS + num_rdma_experts + 1,
lishen's avatar
lishen committed
152
153
154
155
156
157
158
159
160
161
                translate_dst_rdma_rank<kLowLatencyMode>(i, nvl_rank));
            } else {
                UNROLLED_WARP_COPY(1,
                                   lane_id,
                                   NUM_MAX_NVL_PEERS + num_rdma_experts + 1,
                                   rdma_recv_num_tokens_mixed.recv_buffer(rdma_rank),
                                   rdma_recv_num_tokens_mixed.send_buffer(i),
                                   ld_volatile_global,
                                   st_na_global);
            }
Chenggang Zhao's avatar
Chenggang Zhao committed
162
        }
alpha-baby's avatar
alpha-baby committed
163
        __syncthreads();
lishen's avatar
lishen committed
164

Chenggang Zhao's avatar
Chenggang Zhao committed
165
        if (thread_id == 0)
root's avatar
root committed
166
            shmem_sync_with_same_gpu_idx<kLowLatencyMode>(rdma_team);
lijian6's avatar
lijian6 committed
167

Chenggang Zhao's avatar
Chenggang Zhao committed
168
169
170
171
172
        __syncthreads();

        // NVL buffers
        auto nvl_send_buffer = thread_id < NUM_MAX_NVL_PEERS ? buffer_ptrs[thread_id] : nullptr;
        auto nvl_recv_buffer = buffer_ptrs[nvl_rank];
lijian6's avatar
lijian6 committed
173
174
175
176
177
178
179
180
181
182
        auto nvl_reduced_num_tokens_per_expert =
            Buffer<int>(nvl_recv_buffer, num_rdma_experts).advance_also(nvl_send_buffer);
        auto nvl_send_num_tokens_per_rank =
            AsymBuffer<int>(nvl_send_buffer, kNumRDMARanks, NUM_MAX_NVL_PEERS);
        auto nvl_send_num_tokens_per_expert =
            AsymBuffer<int>(nvl_send_buffer, num_nvl_experts, NUM_MAX_NVL_PEERS);
        auto nvl_recv_num_tokens_per_rank =
            AsymBuffer<int>(nvl_recv_buffer, kNumRDMARanks, NUM_MAX_NVL_PEERS);
        auto nvl_recv_num_tokens_per_expert =
            AsymBuffer<int>(nvl_recv_buffer, num_nvl_experts, NUM_MAX_NVL_PEERS);
Chenggang Zhao's avatar
Chenggang Zhao committed
183
184

        // Clean up for later data dispatch
lijian6's avatar
lijian6 committed
185
186
187
188
189
        auto nvl_buffer_ptr_int = reinterpret_cast<int *>(buffer_ptrs[nvl_rank]);
        EP_DEVICE_ASSERT(nvl_reduced_num_tokens_per_expert.total_bytes +
                                      nvl_send_num_tokens_per_rank.total_bytes +
                                      nvl_send_num_tokens_per_expert.total_bytes <=
                                  nvl_clean_offset * sizeof(int));
Chenggang Zhao's avatar
Chenggang Zhao committed
190
191
192
193
        for (int i = thread_id; i < nvl_num_int_clean; i += num_threads)
            nvl_buffer_ptr_int[nvl_clean_offset + i] = 0;

        // Reduce number of tokens per expert into the NVL send buffer
lijian6's avatar
lijian6 committed
194
        // TODO: may use DUSHMEM reduction
Chenggang Zhao's avatar
Chenggang Zhao committed
195
196
197
        EP_DEVICE_ASSERT(num_rdma_experts <= num_threads);
        if (thread_id < num_rdma_experts) {
            int sum = 0;
lijian6's avatar
lijian6 committed
198
199
#pragma unroll
            for (int i = 0; i < kNumRDMARanks; ++i)
Chenggang Zhao's avatar
Chenggang Zhao committed
200
201
202
203
204
205
206
207
                sum += rdma_recv_num_tokens_mixed.recv_buffer(i)[NUM_MAX_NVL_PEERS + thread_id];
            nvl_reduced_num_tokens_per_expert[thread_id] = sum;
        }
        __syncthreads();

        // Reduce RDMA received tokens
        if (thread_id == 0) {
            int sum = 0;
lijian6's avatar
lijian6 committed
208
209
210
211
#pragma unroll
            for (int i = 0; i < kNumRDMARanks; ++i) {
                sum +=
                    rdma_recv_num_tokens_mixed.recv_buffer(i)[NUM_MAX_NVL_PEERS + num_rdma_experts];
Chenggang Zhao's avatar
Chenggang Zhao committed
212
213
                recv_rdma_rank_prefix_sum[i] = sum;
            }
lijian6's avatar
lijian6 committed
214
215
            while (ld_volatile_global(moe_recv_rdma_counter_mapped) != -1)
                ;
Chenggang Zhao's avatar
Chenggang Zhao committed
216
217
218
219
220
            *moe_recv_rdma_counter_mapped = sum;
        }

        // Send numbers of tokens per rank/expert to NVL ranks
        if (thread_id < NUM_MAX_NVL_PEERS) {
lijian6's avatar
lijian6 committed
221
222
223
224
225
226
227
#pragma unroll
            for (int i = 0; i < kNumRDMARanks; ++i)
                nvl_send_num_tokens_per_rank.buffer(nvl_rank)[i] =
                    rdma_recv_num_tokens_mixed.recv_buffer(i)[thread_id];
            for (int i = 0; i < num_nvl_experts; ++i)
                nvl_send_num_tokens_per_expert.buffer(nvl_rank)[i] =
                    nvl_reduced_num_tokens_per_expert[thread_id * num_nvl_experts + i];
Chenggang Zhao's avatar
Chenggang Zhao committed
228
        }
229
        barrier_block<NUM_MAX_NVL_PEERS>(barrier_signal_ptrs, nvl_rank);
Chenggang Zhao's avatar
Chenggang Zhao committed
230

lijian6's avatar
lijian6 committed
231
        // Reduce number of tokens per rank/expert
Chenggang Zhao's avatar
Chenggang Zhao committed
232
233
234
        EP_DEVICE_ASSERT(num_nvl_experts <= num_threads);
        if (thread_id == 0) {
            int sum = 0;
lijian6's avatar
lijian6 committed
235
            for (int i = 0; i < num_ranks; ++i) {
Chenggang Zhao's avatar
Chenggang Zhao committed
236
237
238
239
                int src_rdma_rank = i / NUM_MAX_NVL_PEERS, src_nvl_rank = i % NUM_MAX_NVL_PEERS;
                sum += nvl_recv_num_tokens_per_rank.buffer(src_nvl_rank)[src_rdma_rank];
                recv_gbl_rank_prefix_sum[i] = sum;
            }
lijian6's avatar
lijian6 committed
240
241
            while (ld_volatile_global(moe_recv_counter_mapped) != -1)
                ;
Chenggang Zhao's avatar
Chenggang Zhao committed
242
243
244
245
            *moe_recv_counter_mapped = sum;
        }
        if (thread_id < num_nvl_experts) {
            int sum = 0;
lijian6's avatar
lijian6 committed
246
247
#pragma unroll
            for (int i = 0; i < NUM_MAX_NVL_PEERS; ++i)
Chenggang Zhao's avatar
Chenggang Zhao committed
248
249
                sum += nvl_recv_num_tokens_per_expert.buffer(i)[thread_id];
            sum = (sum + expert_alignment - 1) / expert_alignment * expert_alignment;
lijian6's avatar
lijian6 committed
250
251
            while (ld_volatile_global(moe_recv_expert_counter_mapped + thread_id) != -1)
                ;
Chenggang Zhao's avatar
Chenggang Zhao committed
252
253
254
255
            moe_recv_expert_counter_mapped[thread_id] = sum;
        }

        // Finally barrier
lijian6's avatar
lijian6 committed
256
        if (thread_id == kWarpSize)
root's avatar
root committed
257
            shmem_sync_with_same_gpu_idx<kLowLatencyMode>(rdma_team);
lijian6's avatar
lijian6 committed
258

259
        barrier_block<NUM_MAX_NVL_PEERS>(barrier_signal_ptrs, nvl_rank);
Chenggang Zhao's avatar
Chenggang Zhao committed
260
261
262
263
264
    } else {
        // Calculate meta data
        int dst_rdma_rank = sm_id - 1;
        for (int channel_id = warp_id; channel_id < num_channels; channel_id += num_warps) {
            int token_start_idx, token_end_idx;
lijian6's avatar
lijian6 committed
265
266
            get_channel_task_range(num_tokens, num_channels, channel_id, token_start_idx,
                                   token_end_idx);
Chenggang Zhao's avatar
Chenggang Zhao committed
267
268
269

            // Iterate over tokens
            int total_count = 0, per_nvl_rank_count[NUM_MAX_NVL_PEERS] = {0};
lijian6's avatar
lijian6 committed
270
271
272
273
274
275
276
277
278
            for (int64_t i = token_start_idx + lane_id; i < token_end_idx; i += kWarpSize) {
                EP_STATIC_ASSERT(NUM_MAX_NVL_PEERS * sizeof(bool) == sizeof(uint64_t),
                                          "Invalid number of NVL peers");
                auto is_token_in_rank_uint64 = *reinterpret_cast<const uint64_t *>(
                    is_token_in_rank + i * num_ranks + dst_rdma_rank * NUM_MAX_NVL_PEERS);
                auto is_token_in_rank_values =
                    reinterpret_cast<const bool *>(&is_token_in_rank_uint64);
#pragma unroll
                for (int j = 0; j < NUM_MAX_NVL_PEERS; ++j)
Chenggang Zhao's avatar
Chenggang Zhao committed
279
280
281
282
283
284
                    per_nvl_rank_count[j] += is_token_in_rank_values[j];
                total_count += (is_token_in_rank_uint64 != 0);
            }

            // Warp reduce
            total_count = warp_reduce_sum(total_count);
lijian6's avatar
lijian6 committed
285
286
#pragma unroll
            for (int i = 0; i < NUM_MAX_NVL_PEERS; ++i)
Chenggang Zhao's avatar
Chenggang Zhao committed
287
288
289
                per_nvl_rank_count[i] = warp_reduce_sum(per_nvl_rank_count[i]);

            // Write into channel matrix
lijian6's avatar
lijian6 committed
290
291
292
293
294
295
            if (lane_id == 0) {
#pragma unroll
                for (int i = 0; i < NUM_MAX_NVL_PEERS; ++i)
                    gbl_channel_prefix_matrix[(dst_rdma_rank * NUM_MAX_NVL_PEERS + i) *
                                                  num_channels +
                                              channel_id] = per_nvl_rank_count[i];
Chenggang Zhao's avatar
Chenggang Zhao committed
296
297
298
299
300
301
                rdma_channel_prefix_matrix[dst_rdma_rank * num_channels + channel_id] = total_count;
            }
        }

        // Calculate prefix sum
        __syncthreads();
Chenggang Zhao's avatar
Chenggang Zhao committed
302
        if (thread_id == 0) {
Chenggang Zhao's avatar
Chenggang Zhao committed
303
            auto prefix_row = rdma_channel_prefix_matrix + dst_rdma_rank * num_channels;
lijian6's avatar
lijian6 committed
304
            for (int i = 1; i < num_channels; ++i)
Chenggang Zhao's avatar
Chenggang Zhao committed
305
306
307
                prefix_row[i] += prefix_row[i - 1];
        }

lijian6's avatar
lijian6 committed
308
        EP_STATIC_ASSERT(NUM_MAX_NVL_PEERS <= kWarpSize, "Invalid number of NVL peers");
Chenggang Zhao's avatar
Chenggang Zhao committed
309
        if (thread_id < NUM_MAX_NVL_PEERS) {
lijian6's avatar
lijian6 committed
310
311
312
            auto prefix_row = gbl_channel_prefix_matrix +
                              (dst_rdma_rank * NUM_MAX_NVL_PEERS + thread_id) * num_channels;
            for (int i = 1; i < num_channels; ++i)
Chenggang Zhao's avatar
Chenggang Zhao committed
313
314
315
316
317
                prefix_row[i] += prefix_row[i - 1];
        }
    }
}

lijian6's avatar
lijian6 committed
318
319
320
321
322
323
324
325
326
327
328
void notify_dispatch(const int *num_tokens_per_rank, int *moe_recv_counter_mapped, int num_ranks,
                     const int *num_tokens_per_rdma_rank, int *moe_recv_rdma_counter_mapped,
                     const int *num_tokens_per_expert, int *moe_recv_expert_counter_mapped,
                     int num_experts, const bool *is_token_in_rank, int num_tokens,
                     int num_channels, int hidden_int4, int num_scales, int num_topk,
                     int expert_alignment, int *rdma_channel_prefix_matrix,
                     int *recv_rdma_rank_prefix_sum, int *gbl_channel_prefix_matrix,
                     int *recv_gbl_rank_prefix_sum, void *rdma_buffer_ptr,
                     int num_max_rdma_chunked_recv_tokens, void **buffer_ptrs,
                     int num_max_nvl_chunked_recv_tokens, int **barrier_signal_ptrs, int rank,
                     hipStream_t stream, int64_t num_rdma_bytes, int64_t num_nvl_bytes,
Chenggang Zhao's avatar
Chenggang Zhao committed
329
                     bool low_latency_mode) {
lijian6's avatar
lijian6 committed
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
#define NOTIFY_DISPATCH_LAUNCH_CASE(num_rdma_ranks)                                                \
    {                                                                                              \
        auto notify_dispatch_func = low_latency_mode ? notify_dispatch<true, num_rdma_ranks>       \
                                                     : notify_dispatch<false, num_rdma_ranks>;     \
        LAUNCH_KERNEL_NON_COOPERATIVE(                                                             \
            &cfg, notify_dispatch_func, num_tokens_per_rank, moe_recv_counter_mapped, num_ranks,   \
            num_tokens_per_rdma_rank, moe_recv_rdma_counter_mapped, num_tokens_per_expert,         \
            moe_recv_expert_counter_mapped, num_experts, is_token_in_rank, num_tokens,             \
            num_channels, expert_alignment, rdma_clean_meta.first, rdma_clean_meta.second,         \
            nvl_clean_meta.first, nvl_clean_meta.second, rdma_channel_prefix_matrix,               \
            recv_rdma_rank_prefix_sum, gbl_channel_prefix_matrix, recv_gbl_rank_prefix_sum,        \
            rdma_buffer_ptr, buffer_ptrs, barrier_signal_ptrs, rank, cpu_rdma_team);               \
    }                                                                                              \
    break

    constexpr int kNumThreads    = 256;
    const auto    num_rdma_ranks = num_ranks / NUM_MAX_NVL_PEERS;
Chenggang Zhao's avatar
Chenggang Zhao committed
347
348

    // Get clean meta
lijian6's avatar
lijian6 committed
349
350
351
352
353
354
    auto rdma_clean_meta =
        get_rdma_clean_meta(hidden_int4, num_scales, num_topk, num_topk, num_rdma_ranks,
                            num_max_rdma_chunked_recv_tokens, num_channels);
    auto nvl_clean_meta =
        get_nvl_clean_meta(hidden_int4, num_scales, num_topk, num_topk, num_rdma_ranks,
                           NUM_MAX_NVL_PEERS, num_max_nvl_chunked_recv_tokens, num_channels);
lishen's avatar
lishen committed
355
356
    EP_HOST_ASSERT((rdma_clean_meta.first + rdma_clean_meta.second) * sizeof(int) <= num_rdma_bytes);
    EP_HOST_ASSERT((nvl_clean_meta.first + nvl_clean_meta.second) * sizeof(int) <= num_nvl_bytes);
Chenggang Zhao's avatar
Chenggang Zhao committed
357
358
    EP_HOST_ASSERT(num_rdma_bytes < std::numeric_limits<int>::max());
    EP_HOST_ASSERT(num_nvl_bytes < std::numeric_limits<int>::max());
lishen's avatar
lishen committed
359
360
361
    // add assert origin kernel
    EP_HOST_ASSERT(num_rdma_ranks <= kNumThreads);
    EP_STATIC_ASSERT(NUM_MAX_NVL_PEERS <= kNumThreads, "Assert NUM_MAX_NVL_PEERS <= kNumThreads");
Chenggang Zhao's avatar
Chenggang Zhao committed
362
363
364
365
366
367
368
369
370
371
372
373

    // Launch kernel
    SETUP_LAUNCH_CONFIG(1 + num_rdma_ranks, kNumThreads, stream);
    SWITCH_RDMA_RANKS(NOTIFY_DISPATCH_LAUNCH_CASE);
#undef NOTIFY_DISPATCH_LAUNCH_CASE
}

// At most 8 RDMA ranks to be sent
constexpr int get_num_topk_rdma_ranks(int num_rdma_ranks) {
    return num_rdma_ranks < 8 ? num_rdma_ranks : 8;
}

root's avatar
root committed
374
template <bool kLowLatencyMode, int kNumRDMARanks, bool kCachedMode,
lijian6's avatar
lijian6 committed
375
376
          int kNumDispatchRDMASenderWarps,
          int kNumTopkRDMARanks = get_num_topk_rdma_ranks(kNumRDMARanks)>
root's avatar
root committed
377
378
379
380
381
382
383
384
385
386
387
388
389
390
__global__ void
__launch_bounds__(((kNumDispatchRDMASenderWarps + 1 + NUM_MAX_NVL_PEERS) * kWarpSize), 1)
    dispatch(int4 *recv_x, float *recv_x_scales, int64_t *recv_topk_idx, float *recv_topk_weights,
             SourceMeta *recv_src_meta, const int4 *x, const float *x_scales,
             const int64_t *topk_idx, const float *topk_weights, int *send_rdma_head,
             int *send_nvl_head, int *recv_rdma_channel_prefix_matrix,
             int *recv_gbl_channel_prefix_matrix, const int *rdma_channel_prefix_matrix,
             const int *recv_rdma_rank_prefix_sum, const int *gbl_channel_prefix_matrix,
             const int *recv_gbl_rank_prefix_sum, const bool *is_token_in_rank, int num_tokens,
             int hidden_int4, int num_scales, int num_topk, int num_experts, int scale_token_stride,
             int scale_hidden_stride, void *rdma_buffer_ptr, int num_max_rdma_chunked_send_tokens,
             int num_max_rdma_chunked_recv_tokens, void **buffer_ptrs,
             int num_max_nvl_chunked_send_tokens, int num_max_nvl_chunked_recv_tokens, int rank,
             int num_ranks) {
Chenggang Zhao's avatar
Chenggang Zhao committed
391
    enum class WarpRole {
root's avatar
root committed
392
393
394
395
396
        kRDMASender,
        kRDMASenderCoordinator,
        kRDMAAndNVLForwarder,
        kForwarderCoordinator,
        kNVLReceivers
Chenggang Zhao's avatar
Chenggang Zhao committed
397
    };
root's avatar
root committed
398
399
400

    __shared__ rocshmem::rocshmem_ctx_t ctx;
    rocshmem::rocshmem_wg_ctx_create(0, &ctx);
lijian6's avatar
lijian6 committed
401
402
403

    const auto sm_id       = static_cast<int>(blockIdx.x);
    const auto num_threads = static_cast<int>(blockDim.x), num_warps = num_threads / kWarpSize;
root's avatar
root committed
404
405
406
407
    const auto thread_id = static_cast<int>(threadIdx.x), warp_id = thread_id / kWarpSize,
               lane_id      = get_lane_id();
    const auto num_channels = static_cast<int>(gridDim.x) / 2, channel_id = sm_id / 2;
    const bool is_forwarder = sm_id % 2 == 0;
Chenggang Zhao's avatar
Chenggang Zhao committed
408
409
410
    const auto rdma_rank = rank / NUM_MAX_NVL_PEERS, nvl_rank = rank % NUM_MAX_NVL_PEERS;

    const auto role_meta = [=]() -> std::pair<WarpRole, int> {
root's avatar
root committed
411
412
        if (is_forwarder) {
            if (warp_id < NUM_MAX_NVL_PEERS) {
Chenggang Zhao's avatar
Chenggang Zhao committed
413
414
415
416
                return {WarpRole::kRDMAAndNVLForwarder, (warp_id + channel_id) % NUM_MAX_NVL_PEERS};
            } else {
                return {WarpRole::kForwarderCoordinator, warp_id - NUM_MAX_NVL_PEERS};
            }
root's avatar
root committed
417
418
419
420
        } else if (warp_id < kNumDispatchRDMASenderWarps) {
            return {WarpRole::kRDMASender, -1};
        } else if (warp_id == kNumDispatchRDMASenderWarps) {
            return {WarpRole::kRDMASenderCoordinator, -1};
Chenggang Zhao's avatar
Chenggang Zhao committed
421
        } else {
root's avatar
root committed
422
423
            return {WarpRole::kNVLReceivers,
                    (warp_id + channel_id - kNumDispatchRDMASenderWarps) % NUM_MAX_NVL_PEERS};
Chenggang Zhao's avatar
Chenggang Zhao committed
424
425
        }
    }();
lijian6's avatar
lijian6 committed
426

root's avatar
root committed
427
    auto warp_role   = role_meta.first;
Chenggang Zhao's avatar
Chenggang Zhao committed
428
    auto target_rank = role_meta.second; // Not applicable for RDMA senders
root's avatar
root committed
429
430
431
432
    EP_DEVICE_ASSERT(num_warps == kNumDispatchRDMASenderWarps + 1 + NUM_MAX_NVL_PEERS);

    // Data checks
    EP_DEVICE_ASSERT(num_topk <= kWarpSize);
Chenggang Zhao's avatar
Chenggang Zhao committed
433
434

    // RDMA symmetric layout
root's avatar
root committed
435
436
437
438
439
440
441
442
443
444
445
446
447
448
    EP_STATIC_ASSERT(NUM_MAX_NVL_PEERS * sizeof(bool) == sizeof(uint64_t),
                              "Invalid number of NVL peers");
    auto hidden_bytes = hidden_int4 * sizeof(int4);
    auto num_bytes_per_rdma_token =
        get_num_bytes_per_rdma_token(hidden_int4, num_scales, num_topk, num_topk);
    auto rdma_channel_data = SymBuffer<int8_t>(
        rdma_buffer_ptr, num_max_rdma_chunked_recv_tokens * num_bytes_per_rdma_token, kNumRDMARanks,
        channel_id, num_channels);
    auto rdma_channel_meta = SymBuffer<int>(rdma_buffer_ptr, NUM_MAX_NVL_PEERS * 2 + 2,
                                            kNumRDMARanks, channel_id, num_channels);
    auto rdma_channel_head =
        SymBuffer<uint64_t, false>(rdma_buffer_ptr, 1, kNumRDMARanks, channel_id, num_channels);
    auto rdma_channel_tail =
        SymBuffer<uint64_t, false>(rdma_buffer_ptr, 1, kNumRDMARanks, channel_id, num_channels);
Chenggang Zhao's avatar
Chenggang Zhao committed
449
450

    // NVL buffer layouts
root's avatar
root committed
451
452
    // NOTES: `rs_wr_buffer_ptr` means "Read for Senders, Write for Receivers", `ws_rr_buffer_ptr`
    // means "Write for Senders, Read for Receivers"
Chenggang Zhao's avatar
Chenggang Zhao committed
453
    void *rs_wr_buffer_ptr = nullptr, *ws_rr_buffer_ptr = nullptr;
root's avatar
root committed
454
    int   rs_wr_rank = 0, ws_rr_rank = 0;
Chenggang Zhao's avatar
Chenggang Zhao committed
455
    if (warp_role == WarpRole::kRDMAAndNVLForwarder)
root's avatar
root committed
456
457
        rs_wr_buffer_ptr = buffer_ptrs[nvl_rank], ws_rr_buffer_ptr = buffer_ptrs[target_rank],
        rs_wr_rank = nvl_rank, ws_rr_rank = target_rank;
Chenggang Zhao's avatar
Chenggang Zhao committed
458
    if (warp_role == WarpRole::kNVLReceivers)
root's avatar
root committed
459
460
        rs_wr_buffer_ptr = buffer_ptrs[target_rank], ws_rr_buffer_ptr = buffer_ptrs[nvl_rank],
        rs_wr_rank = target_rank, ws_rr_rank = nvl_rank;
Chenggang Zhao's avatar
Chenggang Zhao committed
461
462

    // Allocate buffers
root's avatar
root committed
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
    auto nvl_channel_x =
        AsymBuffer<int4>(ws_rr_buffer_ptr, num_max_nvl_chunked_recv_tokens * hidden_int4,
                         NUM_MAX_NVL_PEERS, channel_id, num_channels, rs_wr_rank)
            .advance_also(rs_wr_buffer_ptr);
    auto nvl_channel_src_meta =
        AsymBuffer<SourceMeta>(ws_rr_buffer_ptr, num_max_nvl_chunked_recv_tokens, NUM_MAX_NVL_PEERS,
                               channel_id, num_channels, rs_wr_rank)
            .advance_also(rs_wr_buffer_ptr);
    auto nvl_channel_x_scales =
        AsymBuffer<float>(ws_rr_buffer_ptr, num_max_nvl_chunked_recv_tokens * num_scales,
                          NUM_MAX_NVL_PEERS, channel_id, num_channels, rs_wr_rank)
            .advance_also(rs_wr_buffer_ptr);
    auto nvl_channel_topk_idx =
        AsymBuffer<int>(ws_rr_buffer_ptr, num_max_nvl_chunked_recv_tokens * num_topk,
                        NUM_MAX_NVL_PEERS, channel_id, num_channels, rs_wr_rank)
            .advance_also(rs_wr_buffer_ptr);
    auto nvl_channel_topk_weights =
        AsymBuffer<float>(ws_rr_buffer_ptr, num_max_nvl_chunked_recv_tokens * num_topk,
                          NUM_MAX_NVL_PEERS, channel_id, num_channels, rs_wr_rank)
            .advance_also(rs_wr_buffer_ptr);
    auto nvl_channel_prefix_start =
        AsymBuffer<int>(ws_rr_buffer_ptr, kNumRDMARanks, NUM_MAX_NVL_PEERS, channel_id,
                        num_channels, rs_wr_rank)
            .advance_also(rs_wr_buffer_ptr);
    auto nvl_channel_prefix_end =
        AsymBuffer<int>(ws_rr_buffer_ptr, kNumRDMARanks, NUM_MAX_NVL_PEERS, channel_id,
                        num_channels, rs_wr_rank)
            .advance_also(rs_wr_buffer_ptr);
    auto nvl_channel_head = AsymBuffer<int>(rs_wr_buffer_ptr, 1, NUM_MAX_NVL_PEERS, channel_id,
                                            num_channels, ws_rr_rank)
                                .advance_also(ws_rr_buffer_ptr);
    auto nvl_channel_tail = AsymBuffer<int>(ws_rr_buffer_ptr, 1, NUM_MAX_NVL_PEERS, channel_id,
                                            num_channels, rs_wr_rank)
                                .advance_also(rs_wr_buffer_ptr);
Chenggang Zhao's avatar
Chenggang Zhao committed
497
498

    // RDMA sender warp synchronization
lijian6's avatar
lijian6 committed
499
500
501
    __shared__ volatile int rdma_send_next_token_idx;
    __shared__ volatile int rdma_send_channel_tail[kNumRDMARanks];
    __shared__ volatile int rdma_send_channel_next_tail[kNumRDMARanks];
root's avatar
root committed
502
503
504
505
506
507
508
    __shared__ volatile int rdma_sender_counter[1];
    __shared__ volatile int rdma_forwarder_counter[1];
    if (threadIdx.x == 0) {
        rdma_sender_counter[0]    = 0;
        rdma_forwarder_counter[0] = 0;
    }
    __syncthreads();
509

root's avatar
root committed
510
511
512
513
514
515
516
517
518
519
    auto sync_rdma_sender_smem = [&]() {
        if (lane_id == 0) {
            volatile int ret = __hip_atomic_fetch_add(&rdma_sender_counter[0], 1, __ATOMIC_RELAXED,
                                                     __HIP_MEMORY_SCOPE_WORKGROUP);
	    // volatile int ret = atomicAdd((int*)&rdma_sender_counter[0], 1);
        }
        syncwarp();
        while (rdma_sender_counter[0] < (kNumDispatchRDMASenderWarps + 1)) {
        }
    };
lijian6's avatar
lijian6 committed
520

root's avatar
root committed
521
522
523
524
525
526
527
528
529
    // Forward warp synchronization
    __shared__ volatile int  forward_channel_head[NUM_MAX_NVL_PEERS][kNumRDMARanks];
    __shared__ volatile bool forward_channel_retired[NUM_MAX_NVL_PEERS];
    // NOTE: Not sure that __syncthreads() is a suitable replacement
    auto sync_forwarder_smem = [&]() {
        if (lane_id == 0) {
            volatile int ret = __hip_atomic_fetch_add(
                 &rdma_forwarder_counter[0], 1, __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_WORKGROUP);
	    // volatile int ret = atomicAdd((int*)&rdma_forwarder_counter[0], 1);
lijian6's avatar
lijian6 committed
530
        }
root's avatar
root committed
531
532
        syncwarp();
        while (rdma_forwarder_counter[0] < (NUM_MAX_NVL_PEERS + 1)) {
lijian6's avatar
lijian6 committed
533
        }
root's avatar
root committed
534
    };
Chenggang Zhao's avatar
Chenggang Zhao committed
535

root's avatar
root committed
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
    if (warp_role == WarpRole::kRDMASender) {
        // Get tasks
        int token_start_idx, token_end_idx;
        get_channel_task_range(num_tokens, num_channels, channel_id, token_start_idx,
                               token_end_idx);

        // Clean shared memory
        EP_STATIC_ASSERT(kNumRDMARanks <= kWarpSize, "Invalid number of RDMA ranks");
        (warp_id == 0 and lane_id == 0) ? (rdma_send_next_token_idx = token_start_idx) : 0;
        (warp_id == 0 and lane_id < kNumRDMARanks) ? (rdma_send_channel_tail[lane_id] = 0) : 0;
        (warp_id == 0 and lane_id < kNumRDMARanks) ? (rdma_send_channel_next_tail[lane_id] = 0) : 0;

        // Send number of tokens in this channel by `-value - 1`
        EP_STATIC_ASSERT(NUM_MAX_NVL_PEERS * 2 + 2 <= kWarpSize,
                                  "Invalid number of NVL peers");
        for (int dst_rdma_rank = warp_id; dst_rdma_rank < kNumRDMARanks;
             dst_rdma_rank += kNumDispatchRDMASenderWarps) {
Chenggang Zhao's avatar
Chenggang Zhao committed
553
            if (lane_id < NUM_MAX_NVL_PEERS) {
root's avatar
root committed
554
555
556
557
558
559
560
                rdma_channel_meta.send_buffer(dst_rdma_rank)[lane_id] =
                    -(channel_id == 0
                          ? 0
                          : gbl_channel_prefix_matrix
                                [(dst_rdma_rank * NUM_MAX_NVL_PEERS + lane_id) * num_channels +
                                 channel_id - 1]) -
                    1;
Chenggang Zhao's avatar
Chenggang Zhao committed
561
            } else if (lane_id < NUM_MAX_NVL_PEERS * 2) {
root's avatar
root committed
562
563
564
565
566
567
                rdma_channel_meta.send_buffer(dst_rdma_rank)[lane_id] =
                    -gbl_channel_prefix_matrix[(dst_rdma_rank * NUM_MAX_NVL_PEERS + lane_id -
                                                NUM_MAX_NVL_PEERS) *
                                                   num_channels +
                                               channel_id] -
                    1;
Chenggang Zhao's avatar
Chenggang Zhao committed
568
            } else if (lane_id == NUM_MAX_NVL_PEERS * 2) {
root's avatar
root committed
569
570
571
572
573
                rdma_channel_meta.send_buffer(dst_rdma_rank)[lane_id] =
                    -(channel_id == 0 ? 0
                                      : rdma_channel_prefix_matrix[dst_rdma_rank * num_channels +
                                                                   channel_id - 1]) -
                    1;
Chenggang Zhao's avatar
Chenggang Zhao committed
574
            } else if (lane_id == NUM_MAX_NVL_PEERS * 2 + 1) {
root's avatar
root committed
575
576
                rdma_channel_meta.send_buffer(dst_rdma_rank)[lane_id] =
                    -rdma_channel_prefix_matrix[dst_rdma_rank * num_channels + channel_id] - 1;
Chenggang Zhao's avatar
Chenggang Zhao committed
577
            }
lijian6's avatar
lijian6 committed
578

root's avatar
root committed
579
580
            rocshmem::rocshmem_ctx_int_put_nbi_wave(
                ctx, rdma_channel_meta.recv_buffer(rdma_rank),
lijian6's avatar
lijian6 committed
581
582
                rdma_channel_meta.send_buffer(dst_rdma_rank), NUM_MAX_NVL_PEERS * 2 + 2,
                translate_dst_rdma_rank<kLowLatencyMode>(dst_rdma_rank, nvl_rank));
Chenggang Zhao's avatar
Chenggang Zhao committed
583
        }
584

root's avatar
root committed
585
586
        rocshmem::rocshmem_ctx_quiet(ctx);
        sync_rdma_sender_smem();
587

root's avatar
root committed
588
        // Iterate over tokens and copy into buffer
Chenggang Zhao's avatar
Chenggang Zhao committed
589
        int64_t token_idx;
root's avatar
root committed
590
591
592
593
594
595
        int     cached_rdma_channel_head = 0, last_rdma_tail_idx = -1;
        auto    send_buffer = lane_id == rdma_rank ? rdma_channel_data.recv_buffer(lane_id)
                                                   : rdma_channel_data.send_buffer(lane_id);
        for (token_idx = token_start_idx + warp_id; token_idx < token_end_idx;
             token_idx += kNumDispatchRDMASenderWarps) {
            // Read RDMA rank existence
Chenggang Zhao's avatar
Chenggang Zhao committed
596
            uint64_t is_token_in_rank_uint64 = 0;
root's avatar
root committed
597
598
599
            if (lane_id < kNumRDMARanks)
                is_token_in_rank_uint64 = *reinterpret_cast<const uint64_t *>(
                    is_token_in_rank + token_idx * num_ranks + lane_id * NUM_MAX_NVL_PEERS);
Chenggang Zhao's avatar
Chenggang Zhao committed
600

root's avatar
root committed
601
602
603
            // Acquire sequential lock
            while (lane_id == 0 and rdma_send_next_token_idx != token_idx)
                ;
lijian6's avatar
lijian6 committed
604
            syncwarp();
605

root's avatar
root committed
606
            // Acquire next tail
lijian6's avatar
lijian6 committed
607
            int rdma_tail_idx = -1;
root's avatar
root committed
608
            if (is_token_in_rank_uint64 != 0) {
lijian6's avatar
lijian6 committed
609
                rdma_tail_idx = rdma_send_channel_next_tail[lane_id]++;
root's avatar
root committed
610
611
612
                while (rdma_tail_idx - cached_rdma_channel_head >= num_max_rdma_chunked_recv_tokens)
                    cached_rdma_channel_head =
                        static_cast<int>(ld_volatile_global(rdma_channel_head.buffer(lane_id)));
Chenggang Zhao's avatar
Chenggang Zhao committed
613
            }
lijian6's avatar
lijian6 committed
614
            syncwarp();
Chenggang Zhao's avatar
Chenggang Zhao committed
615

root's avatar
root committed
616
617
            // Store RDMA head for combine
            if (lane_id < kNumRDMARanks and not kCachedMode)
Chenggang Zhao's avatar
Chenggang Zhao committed
618
619
                send_rdma_head[token_idx * kNumRDMARanks + lane_id] = rdma_tail_idx;

root's avatar
root committed
620
621
622
623
            // Update last token tail
            if (last_rdma_tail_idx >= 0)
                st_release_cta(const_cast<const int *>(rdma_send_channel_tail + lane_id),
                               last_rdma_tail_idx + 1);
lijian6's avatar
lijian6 committed
624
625
            last_rdma_tail_idx = rdma_tail_idx;

root's avatar
root committed
626
627
            // Release sequential lock
            lane_id == 0 ? (rdma_send_next_token_idx += 1) : 0;
lijian6's avatar
lijian6 committed
628

root's avatar
root committed
629
            // Broadcast tails
Chenggang Zhao's avatar
Chenggang Zhao committed
630
            SourceMeta src_meta;
root's avatar
root committed
631
632
633
634
635
636
637
            int        num_topk_ranks = 0, topk_ranks[kNumTopkRDMARanks];
            void      *dst_send_buffers[kNumTopkRDMARanks];
#pragma unroll
            for (int i = 0, slot_idx; i < kNumRDMARanks; ++i)
                if ((slot_idx = shfl_sync(rdma_tail_idx, i)) >= 0) {
                    slot_idx                          = slot_idx % num_max_rdma_chunked_recv_tokens;
                    topk_ranks[num_topk_ranks]        = i;
lijian6's avatar
lijian6 committed
638
                    auto recv_is_token_in_rank_uint64 = broadcast(is_token_in_rank_uint64, i);
root's avatar
root committed
639
640
641
                    auto recv_is_token_in_rank_values =
                        reinterpret_cast<const bool *>(&recv_is_token_in_rank_uint64);
                    if (lane_id == num_topk_ranks)
lijian6's avatar
lijian6 committed
642
                        src_meta = SourceMeta(rdma_rank, recv_is_token_in_rank_values);
root's avatar
root committed
643
644
645
                    dst_send_buffers[num_topk_ranks++] =
                        reinterpret_cast<uint8_t *>(broadcast(send_buffer, i)) +
                        slot_idx * num_bytes_per_rdma_token;
lijian6's avatar
lijian6 committed
646
                }
Chenggang Zhao's avatar
Chenggang Zhao committed
647
648
            EP_DEVICE_ASSERT(num_topk_ranks <= kNumTopkRDMARanks);

root's avatar
root committed
649
650
651
652
653
654
655
656
657
            // Copy `x` into symmetric send buffer
            auto st_broadcast = [=](const int key, const int4 &value) {
                for (int j = 0; j < num_topk_ranks; ++j)
                    st_na_global(reinterpret_cast<int4 *>(dst_send_buffers[j]) + key, value);
            };
            UNROLLED_WARP_COPY(5, lane_id, hidden_int4, 0, x + token_idx * hidden_int4,
                               ld_nc_global, st_broadcast);
            for (int i = 0; i < num_topk_ranks; ++i)
                dst_send_buffers[i] = reinterpret_cast<int4 *>(dst_send_buffers[i]) + hidden_int4;
lishen's avatar
lishen committed
658

root's avatar
root committed
659
660
661
            // Copy source metadata into symmetric send buffer
            if (lane_id < num_topk_ranks)
                st_na_global(reinterpret_cast<SourceMeta *>(dst_send_buffers[lane_id]), src_meta);
lishen's avatar
lishen committed
662

root's avatar
root committed
663
664
            for (int i = 0; i < num_topk_ranks; ++i)
                dst_send_buffers[i] = reinterpret_cast<SourceMeta *>(dst_send_buffers[i]) + 1;
lijian6's avatar
lijian6 committed
665

root's avatar
root committed
666
667
            // Copy `x_scales` into symmetric send buffer
            for (int i = lane_id; i < num_scales; i += kWarpSize) {
lijian6's avatar
lijian6 committed
668
                auto value = ld_nc_global(x_scales + token_idx * num_scales + i);
root's avatar
root committed
669
670
                for (int j = 0; j < num_topk_ranks; ++j)
                    st_na_global(reinterpret_cast<float *>(dst_send_buffers[j]) + i, value);
lijian6's avatar
lijian6 committed
671
            }
root's avatar
root committed
672
673
            for (int i = 0; i < num_topk_ranks; ++i)
                dst_send_buffers[i] = reinterpret_cast<float *>(dst_send_buffers[i]) + num_scales;
674

root's avatar
root committed
675
676
            // Copy `topk_idx` and `topk_weights` into symmetric send buffer
            for (int i = lane_id; i < num_topk * num_topk_ranks; i += kWarpSize) {
Chenggang Zhao's avatar
Chenggang Zhao committed
677
                auto rank_idx = i / num_topk, copy_idx = i % num_topk;
root's avatar
root committed
678
679
                auto idx_value =
                    static_cast<int>(ld_nc_global(topk_idx + token_idx * num_topk + copy_idx));
Chenggang Zhao's avatar
Chenggang Zhao committed
680
                auto weight_value = ld_nc_global(topk_weights + token_idx * num_topk + copy_idx);
root's avatar
root committed
681
682
683
684
685
                st_na_global(reinterpret_cast<int *>(dst_send_buffers[rank_idx]) + copy_idx,
                             idx_value);
                st_na_global(reinterpret_cast<float *>(dst_send_buffers[rank_idx]) + num_topk +
                                 copy_idx,
                             weight_value);
Chenggang Zhao's avatar
Chenggang Zhao committed
686
            }
lijian6's avatar
lijian6 committed
687
        }
Chenggang Zhao's avatar
Chenggang Zhao committed
688

root's avatar
root committed
689
690
691
692
        // Epilogue
        // Acquire sequential lock
        while (lane_id == 0 and rdma_send_next_token_idx != token_idx)
            ;
lijian6's avatar
lijian6 committed
693
        syncwarp();
Chenggang Zhao's avatar
Chenggang Zhao committed
694

root's avatar
root committed
695
696
697
698
        // Update last token tail
        if (last_rdma_tail_idx >= 0)
            st_release_cta(const_cast<const int *>(rdma_send_channel_tail + lane_id),
                           last_rdma_tail_idx + 1);
Chenggang Zhao's avatar
Chenggang Zhao committed
699

root's avatar
root committed
700
701
702
703
704
705
        // Release sequential lock
        lane_id == 0 ? (rdma_send_next_token_idx += 1) : 0;
    } else if (warp_role == WarpRole::kRDMASenderCoordinator) {
        // NOTES: in case of splitting the issued put at the end of the buffer
        EP_DEVICE_ASSERT(
            num_max_rdma_chunked_recv_tokens % num_max_rdma_chunked_send_tokens == 0);
706

root's avatar
root committed
707
708
        // Synchronize shared memory
        sync_rdma_sender_smem();
Chenggang Zhao's avatar
Chenggang Zhao committed
709

root's avatar
root committed
710
        // Get number of tokens to send for each RDMA rank
Chenggang Zhao's avatar
Chenggang Zhao committed
711
        int num_tokens_to_send = 0;
root's avatar
root committed
712
        if (lane_id < kNumRDMARanks) {
Chenggang Zhao's avatar
Chenggang Zhao committed
713
            num_tokens_to_send = rdma_channel_prefix_matrix[lane_id * num_channels + channel_id];
root's avatar
root committed
714
715
716
            if (channel_id > 0)
                num_tokens_to_send -=
                    rdma_channel_prefix_matrix[lane_id * num_channels + channel_id - 1];
Chenggang Zhao's avatar
Chenggang Zhao committed
717
718
        }

root's avatar
root committed
719
        // Iterate all RDMA ranks
Chenggang Zhao's avatar
Chenggang Zhao committed
720
        int last_issued_tail = 0;
root's avatar
root committed
721
722
723
        while (__any_sync(kFullWarpMask, num_tokens_to_send > 0)) {
            for (int i = 0, synced_num_tokens_to_send; i < kNumRDMARanks; ++i) {
                int dst_rdma_rank         = (i + channel_id) % kNumRDMARanks;
lijian6's avatar
lijian6 committed
724
                synced_num_tokens_to_send = shfl_sync(num_tokens_to_send, dst_rdma_rank);
root's avatar
root committed
725
726
                if (synced_num_tokens_to_send == 0)
                    continue;
Chenggang Zhao's avatar
Chenggang Zhao committed
727

root's avatar
root committed
728
                // Read progress
lijian6's avatar
lijian6 committed
729
                auto synced_last_issued_tail = shfl_sync(last_issued_tail, dst_rdma_rank);
root's avatar
root committed
730
731
732
733
734
                auto processed_tail =
                    ld_acquire_cta(const_cast<const int *>(rdma_send_channel_tail + dst_rdma_rank));
                auto num_tokens_processed = processed_tail - synced_last_issued_tail;
                if (num_tokens_processed != synced_num_tokens_to_send and
                    num_tokens_processed < num_max_rdma_chunked_send_tokens)
Chenggang Zhao's avatar
Chenggang Zhao committed
735
736
                    continue;

root's avatar
root committed
737
738
739
740
741
742
                // Issue RDMA send
                auto num_tokens_to_issue =
                    min(num_tokens_processed, num_max_rdma_chunked_send_tokens);
                EP_DEVICE_ASSERT(num_tokens_to_issue >= 0 and
                                          num_tokens_to_issue <= synced_num_tokens_to_send);
                if (dst_rdma_rank != rdma_rank) {
Chenggang Zhao's avatar
Chenggang Zhao committed
743
                    auto dst_slot_idx = synced_last_issued_tail % num_max_rdma_chunked_recv_tokens;
root's avatar
root committed
744
745
746
747
                    EP_DEVICE_ASSERT(dst_slot_idx + num_tokens_to_issue <=
                                              num_max_rdma_chunked_recv_tokens);
                    rocshmem::rocshmem_ctx_schar_put_nbi_wave(
                        ctx,
lijian6's avatar
lijian6 committed
748
749
750
751
752
753
                        rdma_channel_data.recv_buffer(rdma_rank) +
                            dst_slot_idx * num_bytes_per_rdma_token,
                        rdma_channel_data.send_buffer(dst_rdma_rank) +
                            dst_slot_idx * num_bytes_per_rdma_token,
                        num_bytes_per_rdma_token * num_tokens_to_issue,
                        translate_dst_rdma_rank<kLowLatencyMode>(dst_rdma_rank, nvl_rank));
root's avatar
root committed
754
                    rocshmem::rocshmem_ctx_quiet(ctx);
Chenggang Zhao's avatar
Chenggang Zhao committed
755
                } else {
root's avatar
root committed
756
                    // Lighter fence for local RDMA rank
Chenggang Zhao's avatar
Chenggang Zhao committed
757
758
759
                    memory_fence();
                }

root's avatar
root committed
760
                // Update tails
lijian6's avatar
lijian6 committed
761
                syncwarp();
root's avatar
root committed
762
                if (lane_id == dst_rdma_rank) {
Chenggang Zhao's avatar
Chenggang Zhao committed
763
764
                    last_issued_tail += num_tokens_to_issue;
                    num_tokens_to_send -= num_tokens_to_issue;
root's avatar
root committed
765
766
767

                    rocshmem::rocshmem_ctx_ulong_atomic_add(
                        ctx, rdma_channel_tail.buffer(rdma_rank), num_tokens_to_issue,
lijian6's avatar
lijian6 committed
768
                        translate_dst_rdma_rank<kLowLatencyMode>(dst_rdma_rank, nvl_rank));
Chenggang Zhao's avatar
Chenggang Zhao committed
769
770
                }
            }
root's avatar
root committed
771
772
773
774
775
776
777
778
779
        }
    } else if (warp_role == WarpRole::kRDMAAndNVLForwarder) {
        // RDMA consumers and NVL producers
        const auto dst_nvl_rank          = target_rank;
        const auto dst_rank              = rdma_rank * NUM_MAX_NVL_PEERS + dst_nvl_rank;
        const auto dst_rank_expert_begin = dst_rank * (num_experts / num_ranks);
        const auto dst_rank_expert_end   = dst_rank_expert_begin + (num_experts / num_ranks);

        // Wait counters to arrive
Chenggang Zhao's avatar
Chenggang Zhao committed
780
        int num_tokens_to_recv_from_rdma = 0, src_rdma_channel_prefix = 0;
lijian6's avatar
lijian6 committed
781
782
        EP_DEVICE_ASSERT(kNumRDMARanks <= kWarpSize);
        auto start_time = wall_clock64();
root's avatar
root committed
783
784
785
786
787
788
789
790
791
792
793
794
        if (lane_id < kNumRDMARanks) {
            while (true) {
                auto meta_0 =
                    ld_volatile_global(rdma_channel_meta.recv_buffer(lane_id) + dst_nvl_rank);
                auto meta_1 = ld_volatile_global(rdma_channel_meta.recv_buffer(lane_id) +
                                                 NUM_MAX_NVL_PEERS + dst_nvl_rank);
                auto meta_2 = ld_volatile_global(rdma_channel_meta.recv_buffer(lane_id) +
                                                 NUM_MAX_NVL_PEERS * 2);
                auto meta_3 = ld_volatile_global(rdma_channel_meta.recv_buffer(lane_id) +
                                                 NUM_MAX_NVL_PEERS * 2 + 1);
                if (meta_0 < 0 and meta_1 < 0 and meta_2 < 0 and meta_3 < 0) {
                    // Notify NVL ranks
Chenggang Zhao's avatar
Chenggang Zhao committed
795
                    int start_sum = -meta_0 - 1, end_sum = -meta_1 - 1;
root's avatar
root committed
796
797
798
799
                    EP_DEVICE_ASSERT(start_sum >= 0 and end_sum >= 0 and
                                              end_sum >= start_sum);
                    st_relaxed_sys_global(nvl_channel_prefix_start.buffer() + lane_id,
                                          -start_sum - 1);
Chenggang Zhao's avatar
Chenggang Zhao committed
800
801
                    st_relaxed_sys_global(nvl_channel_prefix_end.buffer() + lane_id, -end_sum - 1);

root's avatar
root committed
802
803
                    // Save RDMA channel received token count
                    src_rdma_channel_prefix        = -meta_2 - 1;
Chenggang Zhao's avatar
Chenggang Zhao committed
804
                    auto src_rdma_channel_prefix_1 = -meta_3 - 1;
root's avatar
root committed
805
806
807
808
809
810
811
                    num_tokens_to_recv_from_rdma =
                        src_rdma_channel_prefix_1 - src_rdma_channel_prefix;
                    if (not kCachedMode)
                        recv_rdma_channel_prefix_matrix[lane_id * num_channels + channel_id] =
                            src_rdma_channel_prefix_1;
                    src_rdma_channel_prefix +=
                        lane_id == 0 ? 0 : recv_rdma_rank_prefix_sum[lane_id - 1];
Chenggang Zhao's avatar
Chenggang Zhao committed
812
813
814
815
                    EP_DEVICE_ASSERT(num_tokens_to_recv_from_rdma >= 0);
                    break;
                }

root's avatar
root committed
816
817
818
819
820
821
822
823
                // Timeout check
                long long int elapsed_time =
                    wall_clock64() > start_time ? wall_clock64() - start_time : 0;
                if (elapsed_time > NUM_TIMEOUT_CYCLES) {
                    printf("DeepEP dispatch forwarder timeout (RDMA meta), channel: %d, RDMA: %d, "
                           "nvl: %d, src RDMA lane: %d, dst NVL: %d, meta: %d, %d, %d, %d\n",
                           channel_id, rdma_rank, nvl_rank, lane_id, dst_nvl_rank, meta_0, meta_1,
                           meta_2, meta_3);
Chenggang Zhao's avatar
Chenggang Zhao committed
824
825
826
827
                    trap();
                }
            }
        }
lijian6's avatar
lijian6 committed
828
        syncwarp();
root's avatar
root committed
829
        // Shift cached head
Chenggang Zhao's avatar
Chenggang Zhao committed
830
831
        send_nvl_head += src_rdma_channel_prefix * NUM_MAX_NVL_PEERS + dst_nvl_rank;

root's avatar
root committed
832
833
        // Wait shared memory to be cleaned
        sync_forwarder_smem();
Chenggang Zhao's avatar
Chenggang Zhao committed
834

root's avatar
root committed
835
836
837
        // Forward tokens from RDMA buffer
        // NOTES: always start from the local rank
        int src_rdma_rank            = sm_id % kNumRDMARanks;
Chenggang Zhao's avatar
Chenggang Zhao committed
838
839
        int cached_rdma_channel_head = 0, cached_rdma_channel_tail = 0;
        int cached_nvl_channel_head = 0, cached_nvl_channel_tail = 0, rdma_nvl_token_idx = 0;
root's avatar
root committed
840
841
        while (__any_sync(kFullWarpMask, num_tokens_to_recv_from_rdma > 0)) {
            // Check destination queue emptiness, or wait a buffer to be released
lijian6's avatar
lijian6 committed
842
            start_time = wall_clock64();
root's avatar
root committed
843
            while (lane_id == 0) {
lijian6's avatar
lijian6 committed
844
                int num_used_slots = cached_nvl_channel_tail - cached_nvl_channel_head;
root's avatar
root committed
845
846
                if (num_max_nvl_chunked_recv_tokens - num_used_slots >=
                    num_max_nvl_chunked_send_tokens)
Chenggang Zhao's avatar
Chenggang Zhao committed
847
                    break;
lijian6's avatar
lijian6 committed
848
                cached_nvl_channel_head = ld_volatile_global(nvl_channel_head.buffer());
Chenggang Zhao's avatar
Chenggang Zhao committed
849

root's avatar
root committed
850
851
852
853
854
855
856
857
                // Timeout check
                long long int elapsed_time =
                    wall_clock64() > start_time ? wall_clock64() - start_time : 0;
                if (elapsed_time > NUM_TIMEOUT_CYCLES) {
                    printf("DeepEP dispatch forwarder timeout (NVL check), channel: %d, RDMA: %d, "
                           "nvl: %d, dst NVL: %d, head: %d, tail: %d\n",
                           channel_id, rdma_rank, nvl_rank, dst_nvl_rank,
                           ld_volatile_global(nvl_channel_head.buffer()), cached_nvl_channel_tail);
Chenggang Zhao's avatar
Chenggang Zhao committed
858
859
860
                    trap();
                }
            }
lijian6's avatar
lijian6 committed
861
            syncwarp();
Chenggang Zhao's avatar
Chenggang Zhao committed
862

root's avatar
root committed
863
            // Find next source RDMA rank (round-robin)
lijian6's avatar
lijian6 committed
864
            start_time = wall_clock64();
root's avatar
root committed
865
            while (true) {
Chenggang Zhao's avatar
Chenggang Zhao committed
866
                src_rdma_rank = (src_rdma_rank + 1) % kNumRDMARanks;
root's avatar
root committed
867
868
869
870
871
872
873
                if (shfl_sync(num_tokens_to_recv_from_rdma, src_rdma_rank) > 0) {
                    if (lane_id == src_rdma_rank and
                        cached_rdma_channel_head == cached_rdma_channel_tail)
                        cached_rdma_channel_tail = static_cast<int>(
                            ld_relaxed_sys_global(rdma_channel_tail.buffer(src_rdma_rank)));
                    if (shfl_sync(cached_rdma_channel_tail > cached_rdma_channel_head,
                                  src_rdma_rank))
Chenggang Zhao's avatar
Chenggang Zhao committed
874
875
876
                        break;
                }

root's avatar
root committed
877
878
879
880
881
882
883
884
885
886
                // Timeout check
                long long int elapsed_time =
                    wall_clock64() > start_time ? wall_clock64() - start_time : 0;
                if (elapsed_time > NUM_TIMEOUT_CYCLES and lane_id < kNumRDMARanks) {
                    printf("DeepEP dispatch forwarder timeout (RDMA check), channel: %d, RDMA: %d, "
                           "nvl: %d, dst NVL: %d, src RDMA lane: %d, head: %d, tail: %d, expected: "
                           "%d\n",
                           channel_id, rdma_rank, nvl_rank, dst_nvl_rank, lane_id,
                           cached_rdma_channel_head, cached_rdma_channel_tail,
                           num_tokens_to_recv_from_rdma);
Chenggang Zhao's avatar
Chenggang Zhao committed
887
888
889
                    trap();
                }
            }
lijian6's avatar
lijian6 committed
890
891
            auto src_rdma_head = shfl_sync(cached_rdma_channel_head, src_rdma_rank);
            auto src_rdma_tail = shfl_sync(cached_rdma_channel_tail, src_rdma_rank);
Chenggang Zhao's avatar
Chenggang Zhao committed
892

root's avatar
root committed
893
894
895
896
897
898
899
900
            // Iterate over every token from the RDMA buffer
            for (int i = src_rdma_head, num_tokens_sent = 0; i < src_rdma_tail; ++i) {
                auto  rdma_slot_idx = i % num_max_rdma_chunked_recv_tokens;
                void *shifted       = rdma_channel_data.recv_buffer(src_rdma_rank) +
                                rdma_slot_idx * num_bytes_per_rdma_token;
                auto src_meta = ld_nc_global(reinterpret_cast<SourceMeta *>(
                    reinterpret_cast<int8_t *>(shifted) + hidden_bytes));
                lane_id == src_rdma_rank ? (num_tokens_to_recv_from_rdma -= 1) : 0;
Chenggang Zhao's avatar
Chenggang Zhao committed
901
                bool is_in_dst_nvl_rank = src_meta.is_token_in_nvl_rank(dst_nvl_rank);
root's avatar
root committed
902
                if (lane_id == src_rdma_rank) {
Chenggang Zhao's avatar
Chenggang Zhao committed
903
904
                    auto cached_head = is_in_dst_nvl_rank ? rdma_nvl_token_idx : -1;
                    rdma_nvl_token_idx += is_in_dst_nvl_rank;
root's avatar
root committed
905
                    if (not kCachedMode)
Chenggang Zhao's avatar
Chenggang Zhao committed
906
907
                        send_nvl_head[i * NUM_MAX_NVL_PEERS] = cached_head;
                }
root's avatar
root committed
908
                if (not is_in_dst_nvl_rank)
Chenggang Zhao's avatar
Chenggang Zhao committed
909
910
                    continue;

root's avatar
root committed
911
                // Get an empty slot
lijian6's avatar
lijian6 committed
912
                int dst_slot_idx = (cached_nvl_channel_tail++) % num_max_nvl_chunked_recv_tokens;
Chenggang Zhao's avatar
Chenggang Zhao committed
913

root's avatar
root committed
914
                // Copy data
lishen's avatar
lishen committed
915
                UNROLLED_WARP_COPY(5, lane_id, hidden_int4,
root's avatar
root committed
916
917
918
919
920
921
922
923
                                   nvl_channel_x.buffer() + dst_slot_idx * hidden_int4,
                                   reinterpret_cast<int4 *>(shifted), ld_nc_global, st_na_global);
                shifted = reinterpret_cast<int4 *>(shifted) + hidden_int4;

                // Copy source meta
                if (lane_id == 0)
                    st_na_global(nvl_channel_src_meta.buffer() + dst_slot_idx, src_meta);
                shifted = reinterpret_cast<SourceMeta *>(shifted) + 1;
lijian6's avatar
lijian6 committed
924

root's avatar
root committed
925
                // Copy `x_scales`
lijian6's avatar
lijian6 committed
926
                UNROLLED_WARP_COPY(1, lane_id, num_scales,
root's avatar
root committed
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
                                   nvl_channel_x_scales.buffer() + dst_slot_idx * num_scales,
                                   reinterpret_cast<float *>(shifted), ld_nc_global, st_na_global);
                shifted = reinterpret_cast<float *>(shifted) + num_scales;

                // Copy `topk_idx` and `topk_weights`
                // NOTES: do not use `shifted` after this `if`, because only several lanes are
                // shifted
                if (lane_id < num_topk) {
                    // Read
                    auto idx_value    = ld_nc_global(reinterpret_cast<int *>(shifted) + lane_id);
                    shifted           = reinterpret_cast<int *>(shifted) + num_topk;
                    auto weight_value = ld_nc_global(reinterpret_cast<float *>(shifted) + lane_id);

                    // Transform and write
                    idx_value =
                        (idx_value >= dst_rank_expert_begin and idx_value < dst_rank_expert_end)
                            ? idx_value - dst_rank_expert_begin
                            : -1;
                    st_na_global(nvl_channel_topk_idx.buffer() + dst_slot_idx * num_topk + lane_id,
                                 idx_value);
                    weight_value = idx_value >= 0 ? weight_value : 0.0f;
                    st_na_global(nvl_channel_topk_weights.buffer() + dst_slot_idx * num_topk +
                                     lane_id,
                                 weight_value);
Chenggang Zhao's avatar
Chenggang Zhao committed
951
952
                }

root's avatar
root committed
953
954
                // In case of insufficient NVL buffers, early stopping
                if ((++num_tokens_sent) == num_max_nvl_chunked_send_tokens)
Chenggang Zhao's avatar
Chenggang Zhao committed
955
956
957
                    src_rdma_tail = i + 1;
            }

root's avatar
root committed
958
959
960
961
            // Sync head index
            if (lane_id == src_rdma_rank)
                forward_channel_head[dst_nvl_rank][src_rdma_rank] =
                    (cached_rdma_channel_head = src_rdma_tail);
Chenggang Zhao's avatar
Chenggang Zhao committed
962

root's avatar
root committed
963
            // Move tail index
lijian6's avatar
lijian6 committed
964
            syncwarp();
root's avatar
root committed
965
966
            if (lane_id == 0)
                st_relaxed_sys_global(nvl_channel_tail.buffer(), cached_nvl_channel_tail);
Chenggang Zhao's avatar
Chenggang Zhao committed
967
968
969
        }

        // Retired
lijian6's avatar
lijian6 committed
970
        syncwarp();
root's avatar
root committed
971
        if (lane_id == 0)
Chenggang Zhao's avatar
Chenggang Zhao committed
972
            forward_channel_retired[dst_nvl_rank] = true;
root's avatar
root committed
973
    } else if (warp_role == WarpRole::kForwarderCoordinator) {
Chenggang Zhao's avatar
Chenggang Zhao committed
974
        // Extra warps for forwarder coordinator should exit directly
root's avatar
root committed
975
        if (target_rank > 0)
Chenggang Zhao's avatar
Chenggang Zhao committed
976
977
            return;

root's avatar
root committed
978
979
980
981
982
983
        // Forward warp coordinator
        EP_STATIC_ASSERT(kNumRDMARanks <= kWarpSize, "Invalid number of RDMA peers");

        // Clean shared memory
        EP_STATIC_ASSERT(NUM_MAX_NVL_PEERS <= kWarpSize, "Invalid number of NVL peers");
        for (int i = lane_id; i < kNumRDMARanks * NUM_MAX_NVL_PEERS; i += kWarpSize)
Chenggang Zhao's avatar
Chenggang Zhao committed
984
            forward_channel_head[i % NUM_MAX_NVL_PEERS][i / NUM_MAX_NVL_PEERS] = 0;
root's avatar
root committed
985
        if (lane_id < NUM_MAX_NVL_PEERS)
Chenggang Zhao's avatar
Chenggang Zhao committed
986
            forward_channel_retired[lane_id] = false;
root's avatar
root committed
987
        sync_forwarder_smem();
Chenggang Zhao's avatar
Chenggang Zhao committed
988
989

        int last_head = 0, target_rdma = lane_id < kNumRDMARanks ? lane_id : 0;
root's avatar
root committed
990
991
        while (true) {
            // Find minimum head
Chenggang Zhao's avatar
Chenggang Zhao committed
992
            int min_head = std::numeric_limits<int>::max();
lijian6's avatar
lijian6 committed
993
#pragma unroll
root's avatar
root committed
994
995
            for (int i = 0; i < NUM_MAX_NVL_PEERS; ++i)
                if (not forward_channel_retired[i])
lijian6's avatar
lijian6 committed
996
                    min_head = min(min_head, forward_channel_head[i][target_rdma]);
root's avatar
root committed
997
            if (__all_sync(kFullWarpMask, min_head == std::numeric_limits<int>::max()))
Chenggang Zhao's avatar
Chenggang Zhao committed
998
999
                break;

root's avatar
root committed
1000
1001
1002
1003
1004
1005
            // Update remote head
            if (min_head != std::numeric_limits<int>::max() and
                min_head >= last_head + num_max_rdma_chunked_send_tokens and
                lane_id < kNumRDMARanks) {
                rocshmem::rocshmem_ctx_ulong_atomic_add(
                    ctx, rdma_channel_head.buffer(rdma_rank), min_head - last_head,
lijian6's avatar
lijian6 committed
1006
                    translate_dst_rdma_rank<kLowLatencyMode>(lane_id, nvl_rank));
1007
1008
                last_head = min_head;
            }
Chenggang Zhao's avatar
Chenggang Zhao committed
1009

root's avatar
root committed
1010
            // Nanosleep and let other warps work
lijian6's avatar
lijian6 committed
1011
            __builtin_amdgcn_s_sleep(NUM_WAIT_CYCLES_TIMES_64);
Chenggang Zhao's avatar
Chenggang Zhao committed
1012
        }
root's avatar
root committed
1013
1014
1015
    } else {
        // NVL consumers
        // Retrieve rank offset from barrier results (each lane's register stores an RDMA rank)
Chenggang Zhao's avatar
Chenggang Zhao committed
1016
        int src_nvl_rank = target_rank, total_offset = 0;
root's avatar
root committed
1017
1018
        EP_STATIC_ASSERT(kNumRDMARanks <= kWarpSize, "Invalid number of RDMA peers");
        if (lane_id < kNumRDMARanks and lane_id * NUM_MAX_NVL_PEERS + src_nvl_rank > 0)
Chenggang Zhao's avatar
Chenggang Zhao committed
1019
1020
            total_offset = recv_gbl_rank_prefix_sum[lane_id * NUM_MAX_NVL_PEERS + src_nvl_rank - 1];

root's avatar
root committed
1021
1022
        // Receive channel offsets
        int  start_offset = 0, end_offset = 0, num_tokens_to_recv;
lijian6's avatar
lijian6 committed
1023
        auto start_time = wall_clock64();
root's avatar
root committed
1024
        while (lane_id < kNumRDMARanks) {
Chenggang Zhao's avatar
Chenggang Zhao committed
1025
            start_offset = ld_volatile_global(nvl_channel_prefix_start.buffer() + lane_id);
lijian6's avatar
lijian6 committed
1026
            end_offset   = ld_volatile_global(nvl_channel_prefix_end.buffer() + lane_id);
root's avatar
root committed
1027
            if (start_offset < 0 and end_offset < 0) {
Chenggang Zhao's avatar
Chenggang Zhao committed
1028
1029
1030
1031
                start_offset = -start_offset - 1, end_offset = -end_offset - 1;
                total_offset += start_offset;
                break;
            }
root's avatar
root committed
1032
1033
1034
1035
1036
1037
1038
1039
1040

            // Timeout check
            long long int elapsed_time =
                wall_clock64() > start_time ? wall_clock64() - start_time : 0;
            if (elapsed_time > NUM_TIMEOUT_CYCLES) {
                printf("DeepEP dispatch NVL receiver timeout, channel: %d, RDMA: %d, nvl: %d, src "
                       "RDMA: %d, src nvl: %d, start: %d, end: %d\n",
                       channel_id, rdma_rank, nvl_rank, lane_id, src_nvl_rank, start_offset,
                       end_offset);
Chenggang Zhao's avatar
Chenggang Zhao committed
1041
1042
1043
1044
1045
                trap();
            }
        }
        num_tokens_to_recv = warp_reduce_sum(end_offset - start_offset);

root's avatar
root committed
1046
1047
1048
1049
1050
        // Save for combine usage
        if (lane_id < kNumRDMARanks and not kCachedMode)
            recv_gbl_channel_prefix_matrix[(lane_id * NUM_MAX_NVL_PEERS + src_nvl_rank) *
                                               num_channels +
                                           channel_id] = total_offset;
lijian6's avatar
lijian6 committed
1051
        syncwarp();
Chenggang Zhao's avatar
Chenggang Zhao committed
1052
1053

        int cached_channel_head_idx = 0, cached_channel_tail_idx = 0;
root's avatar
root committed
1054
1055
        while (num_tokens_to_recv > 0) {
            // Check channel status by lane 0
lijian6's avatar
lijian6 committed
1056
            start_time = wall_clock64();
root's avatar
root committed
1057
1058
1059
            while (lane_id == 0) {
                // Ready to copy
                if (cached_channel_head_idx != cached_channel_tail_idx)
Chenggang Zhao's avatar
Chenggang Zhao committed
1060
                    break;
root's avatar
root committed
1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
                cached_channel_tail_idx = ld_relaxed_sys_global(nvl_channel_tail.buffer());

                // Timeout check
                long long int elapsed_time =
                    wall_clock64() > start_time ? wall_clock64() - start_time : 0;
                if (elapsed_time > NUM_TIMEOUT_CYCLES) {
                    printf("DeepEP dispatch NVL receiver timeout, channel: %d, RDMA: %d, nvl: %d, "
                           "src NVL: %d, head: %d, tail: %d\n",
                           channel_id, rdma_rank, nvl_rank, src_nvl_rank, cached_channel_head_idx,
                           cached_channel_tail_idx);
Chenggang Zhao's avatar
Chenggang Zhao committed
1071
1072
1073
1074
                    trap();
                }
            }

root's avatar
root committed
1075
            // Sync queue tail
lijian6's avatar
lijian6 committed
1076
1077
            cached_channel_tail_idx = shfl_sync(cached_channel_tail_idx, 0);

root's avatar
root committed
1078
            // Copy data
Chenggang Zhao's avatar
Chenggang Zhao committed
1079
            int num_recv_tokens = cached_channel_tail_idx - cached_channel_head_idx;
root's avatar
root committed
1080
1081
1082
1083
1084
1085
            for (int chunk_idx = 0; chunk_idx < num_recv_tokens;
                 ++chunk_idx, --num_tokens_to_recv) {
                int token_idx_in_buffer =
                    (cached_channel_head_idx++) % num_max_nvl_chunked_recv_tokens;
                auto    meta = ld_nc_global(nvl_channel_src_meta.buffer() + token_idx_in_buffer);
                int64_t recv_token_idx = shfl_sync(total_offset, meta.src_rdma_rank);
Chenggang Zhao's avatar
Chenggang Zhao committed
1086
1087
                (lane_id == meta.src_rdma_rank) ? (total_offset += 1) : 0;

root's avatar
root committed
1088
1089
1090
1091
1092
1093
1094
                // Copy data
                UNROLLED_WARP_COPY(5, lane_id, hidden_int4, recv_x + recv_token_idx * hidden_int4,
                                   nvl_channel_x.buffer() + token_idx_in_buffer * hidden_int4,
                                   ld_nc_global, st_na_global);

                // Copy source meta
                if (lane_id == 0 and not kCachedMode)
Chenggang Zhao's avatar
Chenggang Zhao committed
1095
                    st_na_global(recv_src_meta + recv_token_idx, meta);
lijian6's avatar
lijian6 committed
1096

root's avatar
root committed
1097
1098
1099
1100
1101
1102
1103
1104
                // Copy scales
                UNROLLED_WARP_COPY(1, lane_id, num_scales,
                                   recv_x_scales + recv_token_idx * num_scales,
                                   nvl_channel_x_scales.buffer() + token_idx_in_buffer * num_scales,
                                   ld_nc_global, st_na_global);

                // Copy `topk_idx` and `topk_weights`
                if (lane_id < num_topk) {
lijian6's avatar
lijian6 committed
1105
1106
                    auto recv_idx   = recv_token_idx * num_topk + lane_id;
                    auto buffer_idx = token_idx_in_buffer * num_topk + lane_id;
root's avatar
root committed
1107
1108
1109
1110
1111
                    st_na_global(recv_topk_idx + recv_idx,
                                 static_cast<int64_t>(
                                     ld_nc_global(nvl_channel_topk_idx.buffer() + buffer_idx)));
                    st_na_global(recv_topk_weights + recv_idx,
                                 ld_nc_global(nvl_channel_topk_weights.buffer() + buffer_idx));
Chenggang Zhao's avatar
Chenggang Zhao committed
1112
1113
1114
                }
            }

root's avatar
root committed
1115
            // Move queue
lijian6's avatar
lijian6 committed
1116
            syncwarp();
root's avatar
root committed
1117
            if (lane_id == 0)
Chenggang Zhao's avatar
Chenggang Zhao committed
1118
                st_relaxed_sys_global(nvl_channel_head.buffer(), cached_channel_head_idx);
root's avatar
root committed
1119
        }
Chenggang Zhao's avatar
Chenggang Zhao committed
1120
    }
root's avatar
root committed
1121
    rocshmem::rocshmem_wg_ctx_destroy(&ctx);
Chenggang Zhao's avatar
Chenggang Zhao committed
1122
1123
}

lijian6's avatar
lijian6 committed
1124
1125
1126
1127
1128
1129
1130
1131
1132
1133
1134
1135
1136
void dispatch(void *recv_x, float *recv_x_scales, int64_t *recv_topk_idx, float *recv_topk_weights,
              void *recv_src_meta, const void *x, const float *x_scales, const int64_t *topk_idx,
              const float *topk_weights, int *send_rdma_head, int *send_nvl_head,
              int *recv_rdma_channel_prefix_matrix, int *recv_gbl_channel_prefix_matrix,
              const int *rdma_channel_prefix_matrix, const int *recv_rdma_rank_prefix_sum,
              const int *gbl_channel_prefix_matrix, const int *recv_gbl_rank_prefix_sum,
              const bool *is_token_in_rank, int num_tokens, int hidden_int4, int num_scales,
              int num_topk, int num_experts, int scale_token_stride, int scale_hidden_stride,
              void *rdma_buffer_ptr, int num_max_rdma_chunked_send_tokens,
              int num_max_rdma_chunked_recv_tokens, void **buffer_ptrs,
              int num_max_nvl_chunked_send_tokens, int num_max_nvl_chunked_recv_tokens, int rank,
              int num_ranks, bool is_cached_dispatch, hipStream_t stream, int num_channels,
              bool low_latency_mode) {
Chenggang Zhao's avatar
Chenggang Zhao committed
1137
    constexpr int kNumDispatchRDMASenderWarps = 7;
lijian6's avatar
lijian6 committed
1138
1139
1140
1141
1142
1143
1144
1145
1146
1147
1148
1149
1150
1151
1152
1153
1154
1155
1156
1157
1158
1159
1160
1161
1162

#define DISPATCH_LAUNCH_CASE(num_rdma_ranks)                                                       \
    {                                                                                              \
        auto dispatch_func =                                                                       \
            low_latency_mode                                                                       \
                ? (is_cached_dispatch                                                              \
                       ? dispatch<true, num_rdma_ranks, true, kNumDispatchRDMASenderWarps>         \
                       : dispatch<true, num_rdma_ranks, false, kNumDispatchRDMASenderWarps>)       \
                : (is_cached_dispatch                                                              \
                       ? dispatch<false, num_rdma_ranks, true, kNumDispatchRDMASenderWarps>        \
                       : dispatch<false, num_rdma_ranks, false, kNumDispatchRDMASenderWarps>);     \
        LAUNCH_KERNEL_NON_COOPERATIVE(                                                             \
            &cfg, dispatch_func, reinterpret_cast<int4 *>(recv_x), recv_x_scales, recv_topk_idx,   \
            recv_topk_weights, reinterpret_cast<SourceMeta *>(recv_src_meta),                      \
            reinterpret_cast<const int4 *>(x), x_scales, topk_idx, topk_weights, send_rdma_head,   \
            send_nvl_head, recv_rdma_channel_prefix_matrix, recv_gbl_channel_prefix_matrix,        \
            rdma_channel_prefix_matrix, recv_rdma_rank_prefix_sum, gbl_channel_prefix_matrix,      \
            recv_gbl_rank_prefix_sum, is_token_in_rank, num_tokens, hidden_int4, num_scales,       \
            num_topk, num_experts, scale_token_stride, scale_hidden_stride, rdma_buffer_ptr,       \
            num_max_rdma_chunked_send_tokens, num_max_rdma_chunked_recv_tokens, buffer_ptrs,       \
            num_max_nvl_chunked_send_tokens, num_max_nvl_chunked_recv_tokens, rank, num_ranks);    \
    }                                                                                              \
    break

    EP_HOST_ASSERT((topk_idx == nullptr) == (topk_weights == nullptr));
Chenggang Zhao's avatar
Chenggang Zhao committed
1163
1164
    EP_HOST_ASSERT((recv_topk_idx == nullptr) == (recv_topk_weights == nullptr));

root's avatar
root committed
1165
1166
    SETUP_LAUNCH_CONFIG(num_channels * 2,
                        (kNumDispatchRDMASenderWarps + 1 + NUM_MAX_NVL_PEERS) * kWarpSize, stream);
Chenggang Zhao's avatar
Chenggang Zhao committed
1167
1168
1169
1170
    SWITCH_RDMA_RANKS(DISPATCH_LAUNCH_CASE);
#undef DISPATCH_LAUNCH_CASE
}

lijian6's avatar
lijian6 committed
1171
template <bool kLowLatencyMode>
lijian6's avatar
lijian6 committed
1172
__global__ void __launch_bounds__(1024, 1)
lijian6's avatar
lijian6 committed
1173
1174
1175
1176
1177
cached_notify(const int rdma_clean_offset, const int rdma_num_int_clean, const int nvl_clean_offset,
              const int nvl_num_int_clean, int *combined_rdma_head, int num_combined_tokens,
              int num_channels, const int *rdma_channel_prefix_matrix,
              const int *rdma_rank_prefix_sum, int *combined_nvl_head, void *rdma_buffer_ptr,
              void **buffer_ptrs, int **barrier_signal_ptrs, int rank, int num_ranks,
1178
              bool is_cached_dispatch, const shmem_team_t rdma_team) {
lijian6's avatar
lijian6 committed
1179
1180
    auto sm_id       = static_cast<int>(blockIdx.x);
    auto thread_id   = static_cast<int>(threadIdx.x);
Chenggang Zhao's avatar
Chenggang Zhao committed
1181
    auto num_threads = static_cast<int>(blockDim.x);
lijian6's avatar
lijian6 committed
1182
1183
1184
    auto num_warps   = num_threads / kWarpSize;
    auto warp_id     = thread_id / kWarpSize;
    auto lane_id     = get_lane_id();
Chenggang Zhao's avatar
Chenggang Zhao committed
1185

lijian6's avatar
lijian6 committed
1186
    auto nvl_rank       = rank % NUM_MAX_NVL_PEERS;
Chenggang Zhao's avatar
Chenggang Zhao committed
1187
1188
1189
1190
1191
    auto num_rdma_ranks = num_ranks / NUM_MAX_NVL_PEERS;

    // Using two SMs, which clean the RDMA/NVL buffer respectively
    if (sm_id == 0) {
        // Barrier for RDMA
lishen's avatar
lishen committed
1192
        if (thread_id == kWarpSize)
root's avatar
root committed
1193
            shmem_sync_with_same_gpu_idx<kLowLatencyMode>(rdma_team);
Shangyan Zhou's avatar
Fix  
Shangyan Zhou committed
1194

lishen's avatar
lishen committed
1195
1196
        // Barrier for NVL
        barrier_block<NUM_MAX_NVL_PEERS>(barrier_signal_ptrs, nvl_rank);
Chenggang Zhao's avatar
Chenggang Zhao committed
1197

lishen's avatar
lishen committed
1198
        // Clean RDMA buffer
lijian6's avatar
lijian6 committed
1199
        auto rdma_buffer_ptr_int = reinterpret_cast<int *>(rdma_buffer_ptr);
Chenggang Zhao's avatar
Chenggang Zhao committed
1200
1201
        for (int i = thread_id; i < rdma_num_int_clean; i += num_threads)
            rdma_buffer_ptr_int[rdma_clean_offset + i] = 0;
lijian6's avatar
lijian6 committed
1202

lishen's avatar
lishen committed
1203
        // Clean NVL buffer
lijian6's avatar
lijian6 committed
1204
        auto nvl_buffer_ptr_int = reinterpret_cast<int *>(buffer_ptrs[nvl_rank]);
Chenggang Zhao's avatar
Chenggang Zhao committed
1205
1206
        for (int i = thread_id; i < nvl_num_int_clean; i += num_threads)
            nvl_buffer_ptr_int[nvl_clean_offset + i] = 0;
lishen's avatar
lishen committed
1207

1208
        __syncthreads();
Chenggang Zhao's avatar
Chenggang Zhao committed
1209

lishen's avatar
lishen committed
1210
1211
        // Barrier again
        if (thread_id == kWarpSize)
root's avatar
root committed
1212
            shmem_sync_with_same_gpu_idx<kLowLatencyMode>(rdma_team);
lishen's avatar
lishen committed
1213

Chenggang Zhao's avatar
Chenggang Zhao committed
1214
        // Barrier again
1215
        barrier_block<NUM_MAX_NVL_PEERS>(barrier_signal_ptrs, nvl_rank);
1216
    } else { 
Chenggang Zhao's avatar
Chenggang Zhao committed
1217
1218
1219
        if (is_cached_dispatch)
            return;

1220
1221
1222
1223
1224
1225
1226
1227
1228
1229
1230
1231
1232
1233
1234
1235
1236
1237
1238
1239
1240
1241
        constexpr int num_clean_sms = 1;
        const int logical_block = sm_id - num_clean_sms;
        const int total_blocks  = gridDim.x - num_clean_sms;

        if (logical_block < 0) return;
        if (combined_rdma_head != nullptr) {
            EP_DEVICE_ASSERT(num_rdma_ranks <= kWarpSize);
            for (int chan = logical_block; chan < num_channels; chan += total_blocks) {
                int token_start_idx, token_end_idx;
                get_channel_task_range(num_combined_tokens, num_channels, chan, token_start_idx, token_end_idx);

                for (int token_idx = token_end_idx - 1 - warp_id; token_idx >= token_start_idx; token_idx -= num_warps) {
                    int last_head = 1 << 25;
                    if (lane_id < num_rdma_ranks) {
                        auto ptr = combined_rdma_head + token_idx * num_rdma_ranks + lane_id;
                        int current_head = __ldg(ptr);
                        if (current_head < 0) {
                            *ptr = -last_head - 1;
                        } else {
                            last_head = current_head;
                        }
                    }
Chenggang Zhao's avatar
Chenggang Zhao committed
1242
1243
1244
                }
            }
        }
1245
1246
1247
1248
1249
1250
1251
1252
1253
1254
1255
1256
1257
1258
1259
1260
1261
1262
1263
1264
1265
1266
1267
1268
1269
        if (combined_nvl_head != nullptr) {
            EP_DEVICE_ASSERT(rdma_channel_prefix_matrix != nullptr);
            EP_DEVICE_ASSERT(rdma_rank_prefix_sum != nullptr);
            EP_STATIC_ASSERT(NUM_MAX_NVL_PEERS <= kWarpSize, "Too many NVL peers");

            for (int chan = logical_block; chan < num_channels; chan += total_blocks) {
                for (int dst_rdma_rank = 0; dst_rdma_rank < num_rdma_ranks; ++dst_rdma_rank) {
                    int token_start_idx = (chan == 0) ? 0 : rdma_channel_prefix_matrix[dst_rdma_rank * num_channels + chan - 1];
                    int token_end_idx = rdma_channel_prefix_matrix[dst_rdma_rank * num_channels + chan];
                    int shift = (dst_rdma_rank == 0) ? 0 : rdma_rank_prefix_sum[dst_rdma_rank - 1];
                    token_start_idx += shift;
                    token_end_idx   += shift;

                    for (int token_idx = token_end_idx - 1 - warp_id; token_idx >= token_start_idx; token_idx -= num_warps) {
                        int last_head = 1 << 25;
                        if (lane_id < NUM_MAX_NVL_PEERS) {
                            auto ptr = combined_nvl_head +
                                       token_idx * NUM_MAX_NVL_PEERS + lane_id;
                            int current_head = __ldg(ptr);
                            if (current_head < 0) {
                                *ptr = -last_head - 1;
                            } else {
                                last_head = current_head;
                            }
                        }
1270
                    }
Chenggang Zhao's avatar
Chenggang Zhao committed
1271
1272
1273
1274
1275
1276
1277
                }
            }
        }
    }
}

void cached_notify(int hidden_int4, int num_scales, int num_topk_idx, int num_topk_weights,
lijian6's avatar
lijian6 committed
1278
1279
1280
1281
1282
1283
                   int num_ranks, int num_channels, int num_combined_tokens,
                   int *combined_rdma_head, const int *rdma_channel_prefix_matrix,
                   const int *rdma_rank_prefix_sum, int *combined_nvl_head, void *rdma_buffer_ptr,
                   int num_max_rdma_chunked_recv_tokens, void **buffer_ptrs,
                   int num_max_nvl_chunked_recv_tokens, int **barrier_signal_ptrs, int rank,
                   hipStream_t stream, int64_t num_rdma_bytes, int64_t num_nvl_bytes,
Chenggang Zhao's avatar
Chenggang Zhao committed
1284
                   bool is_cached_dispatch, bool low_latency_mode) {
lijian's avatar
lijian committed
1285
    const int  num_threads    = ::min(1024, ::max(128, kWarpSize * num_channels));
Chenggang Zhao's avatar
Chenggang Zhao committed
1286
1287
1288
    const auto num_rdma_ranks = num_ranks / NUM_MAX_NVL_PEERS;

    // Get clean meta
lijian6's avatar
lijian6 committed
1289
1290
1291
1292
1293
1294
1295
1296
1297
1298
    auto rdma_clean_meta =
        get_rdma_clean_meta(hidden_int4, num_scales, num_topk_idx, num_topk_weights, num_rdma_ranks,
                            num_max_rdma_chunked_recv_tokens, num_channels);
    auto nvl_clean_meta =
        get_nvl_clean_meta(hidden_int4, num_scales, num_topk_idx, num_topk_weights, num_rdma_ranks,
                           NUM_MAX_NVL_PEERS, num_max_nvl_chunked_recv_tokens, num_channels);
    EP_HOST_ASSERT((rdma_clean_meta.first + rdma_clean_meta.second) * sizeof(int) <=
                       num_rdma_bytes);
    EP_HOST_ASSERT((nvl_clean_meta.first + nvl_clean_meta.second) * sizeof(int) <=
                       num_nvl_bytes);
Chenggang Zhao's avatar
Chenggang Zhao committed
1299
1300
    EP_HOST_ASSERT(num_rdma_bytes < std::numeric_limits<int>::max());
    EP_HOST_ASSERT(num_nvl_bytes < std::numeric_limits<int>::max());
root's avatar
root committed
1301
    EP_HOST_ASSERT(num_channels * 2 > 2);
Chenggang Zhao's avatar
Chenggang Zhao committed
1302
1303

    // Launch kernel
lijian6's avatar
lijian6 committed
1304
    auto cached_notify_func = low_latency_mode ? cached_notify<true> : cached_notify<false>;
root's avatar
root committed
1305
    SETUP_LAUNCH_CONFIG(num_channels * 2, num_threads, stream);
lijian6's avatar
lijian6 committed
1306
1307
1308
1309
1310
1311
    LAUNCH_KERNEL_NON_COOPERATIVE(
        &cfg, cached_notify_func, rdma_clean_meta.first, rdma_clean_meta.second,
        nvl_clean_meta.first, nvl_clean_meta.second, combined_rdma_head, num_combined_tokens,
        num_channels, rdma_channel_prefix_matrix, rdma_rank_prefix_sum, combined_nvl_head,
        rdma_buffer_ptr, buffer_ptrs, barrier_signal_ptrs, rank, num_ranks, is_cached_dispatch,
        cpu_rdma_team);
Chenggang Zhao's avatar
Chenggang Zhao committed
1312
1313
}

root's avatar
root committed
1314
template <int kNumRanks, typename dtype_t, int kMaxNumRanks, int kWidth, typename ReceiveFn, typename ReceiveTWFn>
lijian6's avatar
lijian6 committed
1315
1316
1317
__device__ int combine_token(bool is_token_in_rank, int head_idx,
                             int lane_id, int hidden_int4, int num_topk,
                             int4* combined_row, float* combined_topk_weights,
root's avatar
root committed
1318
                             int num_max_recv_tokens, const ReceiveFn& recv_fn, const ReceiveTWFn& recv_tw_fn) {
Chenggang Zhao's avatar
Chenggang Zhao committed
1319
1320
1321
1322
    constexpr auto kDtypePerInt4 = sizeof(int4) / sizeof(dtype_t);

    // Broadcast current heads
    // Lane `i` holds the head of rank `i` and `is_token_in_rank`
root's avatar
root committed
1323
    EP_STATIC_ASSERT(kMaxNumRanks <= kWidth, "Too many ranks");
Chenggang Zhao's avatar
Chenggang Zhao committed
1324
    int num_topk_ranks = 0, topk_ranks[kMaxNumRanks], slot_indices[kMaxNumRanks];
lijian6's avatar
lijian6 committed
1325
    #pragma unroll
root's avatar
root committed
1326
1327
    for (int i = 0; i < kNumRanks; ++ i) if (shfl_sync(is_token_in_rank, i, kWidth)) {
        slot_indices[num_topk_ranks] = shfl_sync(head_idx, i, kWidth) % num_max_recv_tokens;
lijian6's avatar
lijian6 committed
1328
1329
        topk_ranks[num_topk_ranks ++] = i;
    }
Chenggang Zhao's avatar
Chenggang Zhao committed
1330
1331
    EP_DEVICE_ASSERT(num_topk_ranks <= kMaxNumRanks);

root's avatar
root committed
1332
    // Reduce data    
lijian6's avatar
lijian6 committed
1333
    #pragma unroll
root's avatar
root committed
1334
1335
1336
1337
1338
1339
    for (int i = lane_id; i < hidden_int4; i += kWidth) {
        float values[kDtypePerInt4] = {0};

        // Temporary buffer
        int4 temp;

lijian6's avatar
lijian6 committed
1340
        #pragma unroll
lishen's avatar
lishen committed
1341
        for (int j = 0; j < num_topk_ranks; ++j) {
root's avatar
root committed
1342
1343
1344
            temp = recv_fn(topk_ranks[j], slot_indices[j], i);
            const dtype_t* d = reinterpret_cast<const dtype_t*>(&temp);

lijian6's avatar
lijian6 committed
1345
            #pragma unroll
lishen's avatar
lishen committed
1346
            for (int k = 0; k < kDtypePerInt4; ++k)
root's avatar
root committed
1347
                values[k] += static_cast<float>(d[k]);
Shangyan Zhou's avatar
Shangyan Zhou committed
1348
        }
Chenggang Zhao's avatar
Chenggang Zhao committed
1349

lijian6's avatar
lijian6 committed
1350
        int4 out_int4;
root's avatar
root committed
1351
1352
        dtype_t* out_dtypes = reinterpret_cast<dtype_t*>(&out_int4);

lijian6's avatar
lijian6 committed
1353
        #pragma unroll
root's avatar
root committed
1354
        for (int j = 0; j < kDtypePerInt4; ++j)
lijian6's avatar
lijian6 committed
1355
            out_dtypes[j] = static_cast<dtype_t>(values[j]);
root's avatar
root committed
1356

lijian6's avatar
lijian6 committed
1357
        st_na_global(combined_row + i, out_int4);
Chenggang Zhao's avatar
Chenggang Zhao committed
1358
1359
1360
1361
1362
    }

    // Reduce `topk_weights`
    if (lane_id < num_topk) {
        float value = 0;
lijian6's avatar
lijian6 committed
1363
1364
        #pragma unroll
        for (int i = 0; i < num_topk_ranks; ++ i)
Chenggang Zhao's avatar
Chenggang Zhao committed
1365
1366
1367
1368
1369
1370
1371
1372
            value += recv_tw_fn(topk_ranks[i], slot_indices[i], lane_id);
        st_na_global(combined_topk_weights + lane_id, value);
    }

    // Return the minimum top-k rank
    return topk_ranks[0];
}

root's avatar
root committed
1373
1374
1375
1376
1377
1378
1379
1380
1381
1382
1383
1384
1385
1386
1387
1388
1389
1390
template<bool kLowLatencyMode,
         int kNumRDMARanks, typename dtype_t,
         int kNumCombineForwarderWarps,
         int kNumTopkRDMARanks = get_num_topk_rdma_ranks(kNumRDMARanks),
         int kNumWarpsPerForwarder = (kNumCombineForwarderWarps / kNumRDMARanks > 0) ? kNumCombineForwarderWarps / kNumRDMARanks : 1,
         int kNumForwarders = kNumRDMARanks * kNumWarpsPerForwarder,
         int kNumRDMAReceivers = kNumRDMARanks <=8 ? kNumForwarders + NUM_MAX_NVL_PEERS / 2: kNumForwarders + NUM_MAX_NVL_PEERS,
         int kBlockThreads = (kNumRDMARanks > 8) ? ((NUM_MAX_NVL_PEERS + kNumForwarders) * kEmulatedWarpSize + kWarpSize) : ((NUM_MAX_NVL_PEERS/2 + 1 + kNumForwarders) * kWarpSize) >
__global__ void __launch_bounds__(kBlockThreads, 1)
combine(int4* combined_x, float* combined_topk_weights,
        const bool* is_combined_token_in_rank,
        const int4* x, const float* topk_weights, const int4 *bias_0, const int4 *bias_1,
        const int* combined_rdma_head, const int* combined_nvl_head,
        const SourceMeta* src_meta, const int* rdma_channel_prefix_matrix, const int* rdma_rank_prefix_sum, const int* gbl_channel_prefix_matrix,
        int num_tokens, int num_combined_tokens, int hidden, int num_topk,
        void* rdma_buffer_ptr, int num_max_rdma_chunked_send_tokens, int num_max_rdma_chunked_recv_tokens,
        void** buffer_ptrs, int num_max_nvl_chunked_send_tokens, int num_max_nvl_chunked_recv_tokens,
        int rank, int num_ranks) {
lijian6's avatar
lijian6 committed
1391
1392
1393
1394
    enum class WarpRole {
        kNVLSender,
        kNVLAndRDMAForwarder,
        kRDMAReceiver,
root's avatar
root committed
1395
        kCoordinator
lijian6's avatar
lijian6 committed
1396
    };
root's avatar
root committed
1397
1398
1399
1400
1401
1402
1403
1404
1405
1406

    constexpr auto kNVLPeersHyb = (kNumRDMARanks > 8) ? NUM_MAX_NVL_PEERS : NUM_MAX_NVL_PEERS / 2;
    constexpr auto  kWarpHyb = kNumRDMARanks > 8 ? kEmulatedWarpSize : kWarpSize;
    const auto sm_id = static_cast<int>(blockIdx.x);
    const auto num_threads = static_cast<int>(blockDim.x);
    const int num_warps = kNumRDMARanks > 8 ? (num_threads / kEmulatedWarpSize - 1) : (num_threads / kWarpSize);
    auto thread_id = static_cast<int>(threadIdx.x);
    const auto num_channels = static_cast<int>(gridDim.x) / 2, channel_id = sm_id / 2;
    const bool is_rdma_receiver_sm = sm_id % 2 == 1;

lijian6's avatar
lijian6 committed
1407
#if !defined(FORCE_DUSHMEM_API) && !defined(ROCM_DISABLE_CTX)
1408
1409
    __shared__ shmem_ctx_t ctx;
    shmem_wg_ctx_create(&ctx);
lijian6's avatar
lijian6 committed
1410

root's avatar
root committed
1411
#endif
lijian6's avatar
lijian6 committed
1412

root's avatar
root committed
1413
1414
    EP_DEVICE_ASSERT(num_topk <= kEmulatedWarpSize);
    EP_DEVICE_ASSERT(hidden % (sizeof(int4) / sizeof(dtype_t)) == 0);
Chenggang Zhao's avatar
Chenggang Zhao committed
1415
1416
1417
1418
    const auto hidden_int4 = hidden / (sizeof(int4) / sizeof(dtype_t));

    // NOTES: we decouple a channel into 2 SMs
    const auto rdma_rank = rank / NUM_MAX_NVL_PEERS, nvl_rank = rank % NUM_MAX_NVL_PEERS;
root's avatar
root committed
1419
1420
1421
1422
1423
1424
1425
1426
1427
1428
1429
    auto role_meta = [=]() -> std::pair<WarpRole, int> {
        auto warp_id = thread_id / kWarpHyb;
        if (not is_rdma_receiver_sm) {
            if (warp_id < kNVLPeersHyb) {
                auto shuffled_warp_id = warp_id;
                shuffled_warp_id = (shuffled_warp_id + channel_id) % kNVLPeersHyb;
                return {WarpRole::kNVLSender, shuffled_warp_id};
            } else if (warp_id < kNVLPeersHyb + kNumForwarders) {
                auto shuffled_warp_id = warp_id - kNVLPeersHyb;
                shuffled_warp_id = (shuffled_warp_id + channel_id) % kNumForwarders;
                return {WarpRole::kNVLAndRDMAForwarder, shuffled_warp_id};
Chenggang Zhao's avatar
Chenggang Zhao committed
1430
            } else {
root's avatar
root committed
1431
                return {WarpRole::kCoordinator, 0};
Chenggang Zhao's avatar
Chenggang Zhao committed
1432
1433
            }
        } else {
root's avatar
root committed
1434
            if (warp_id < kNVLPeersHyb + kNumForwarders) {
lijian6's avatar
lijian6 committed
1435
                return {WarpRole::kRDMAReceiver, warp_id};
Chenggang Zhao's avatar
Chenggang Zhao committed
1436
            } else {
root's avatar
root committed
1437
                return {WarpRole::kCoordinator, 0};
Chenggang Zhao's avatar
Chenggang Zhao committed
1438
1439
1440
1441
            }
        }
    }();
    auto warp_role = role_meta.first;
root's avatar
root committed
1442
    auto warp_id = role_meta.second;
Chenggang Zhao's avatar
Chenggang Zhao committed
1443
1444
1445

    auto num_max_nvl_chunked_recv_tokens_per_rdma = num_max_nvl_chunked_recv_tokens / kNumRDMARanks;

lijian6's avatar
lijian6 committed
1446
    // This approach is designed to sync multiple warps in a loop
lijian6's avatar
lijian6 committed
1447
    constexpr int num_sync_large_iteration = 64;
root's avatar
root committed
1448
1449
1450
1451
1452
1453
1454
1455
1456
    __shared__ volatile int rdma_receiver_counter[1];
    __shared__ volatile int rdma_forwarder_counter[1];
    __shared__ volatile uint8_t sync_large_warp_counters[2 * kNumRDMARanks * num_sync_large_iteration ];
    if (threadIdx.x==0){
        rdma_receiver_counter[0] = 0;
        rdma_forwarder_counter[0] = 0;
    }
    
    for (int i = thread_id; i < 2 * kNumRDMARanks * num_sync_large_iteration; i += num_threads) {
lijian6's avatar
lijian6 committed
1457
1458
1459
        sync_large_warp_counters[i] = 0;
    }
    __syncthreads();
Chenggang Zhao's avatar
Chenggang Zhao committed
1460
    if (warp_role == WarpRole::kNVLSender) {
root's avatar
root committed
1461
1462
1463
        // NVL producers
        const int dst_nvl_rank = kNumRDMARanks <= 8 ? (warp_id * 2 + (thread_id % kWarpSize) / kEmulatedWarpSize) : warp_id;  
        auto lane_id = get_lane_id() % kEmulatedWarpSize;
Chenggang Zhao's avatar
Chenggang Zhao committed
1464
1465
1466
1467

        // NVL layouts
        // NOTES: to avoid deadlocks, we use separate NVL buffers for different RDMA sources
        auto dst_buffer_ptr = buffer_ptrs[dst_nvl_rank], local_buffer_ptr = buffer_ptrs[nvl_rank];
lijian6's avatar
lijian6 committed
1468
1469
1470
1471
1472
        auto nvl_channel_x = AsymBuffer<int4>(dst_buffer_ptr, num_max_nvl_chunked_recv_tokens * hidden_int4, NUM_MAX_NVL_PEERS, channel_id, num_channels, nvl_rank).advance_also(local_buffer_ptr);
        auto nvl_channel_src_meta = AsymBuffer<SourceMeta>(dst_buffer_ptr, num_max_nvl_chunked_recv_tokens, NUM_MAX_NVL_PEERS, channel_id, num_channels, nvl_rank).advance_also(local_buffer_ptr);
        auto nvl_channel_topk_weights = AsymBuffer<float>(dst_buffer_ptr, num_max_nvl_chunked_recv_tokens * num_topk, NUM_MAX_NVL_PEERS, channel_id, num_channels, nvl_rank).advance_also(local_buffer_ptr);
        auto nvl_channel_head = AsymBuffer<int>(local_buffer_ptr, kNumRDMARanks, NUM_MAX_NVL_PEERS, channel_id, num_channels, dst_nvl_rank).advance_also(dst_buffer_ptr);
        auto nvl_channel_tail = AsymBuffer<int>(dst_buffer_ptr, kNumRDMARanks, NUM_MAX_NVL_PEERS, channel_id, num_channels, nvl_rank).advance_also(local_buffer_ptr);
1473

Chenggang Zhao's avatar
Chenggang Zhao committed
1474
1475
        // Get tasks for each RDMA lane
        int token_start_idx = 0, token_end_idx = 0;
root's avatar
root committed
1476
1477
        if (lane_id < kNumRDMARanks) {
            int prefix_idx = (lane_id * NUM_MAX_NVL_PEERS + dst_nvl_rank) * num_channels + channel_id;
Chenggang Zhao's avatar
Chenggang Zhao committed
1478
            token_start_idx = gbl_channel_prefix_matrix[prefix_idx];
root's avatar
root committed
1479
            token_end_idx = (prefix_idx == num_channels * num_ranks - 1) ? num_tokens : gbl_channel_prefix_matrix[prefix_idx + 1];
Chenggang Zhao's avatar
Chenggang Zhao committed
1480
        }
lijian6's avatar
lijian6 committed
1481
        syncwarp();
Chenggang Zhao's avatar
Chenggang Zhao committed
1482
1483
1484

        // NOTES: here the cached value of each lane is only responsible for a single RDMA buffer
        int cached_channel_head_idx = 0, cached_channel_tail_idx = 0;
root's avatar
root committed
1485
        EP_STATIC_ASSERT(kNumRDMARanks <= kEmulatedWarpSize, "Invalid number of RDMA peers");
Chenggang Zhao's avatar
Chenggang Zhao committed
1486
1487

        // Iterate over all tokens and send by chunks
root's avatar
root committed
1488
        while (true) {
Chenggang Zhao's avatar
Chenggang Zhao committed
1489
            // Exit if possible
root's avatar
root committed
1490
            if (__all_sync(kFullWarpMask, token_start_idx >= token_end_idx))
Chenggang Zhao's avatar
Chenggang Zhao committed
1491
1492
                break;

lijian6's avatar
lijian6 committed
1493
            // Decide next RDMA buffer to send
Chenggang Zhao's avatar
Chenggang Zhao committed
1494
            bool is_lane_ready = false;
root's avatar
root committed
1495
1496
            auto start_time = wall_clock64();
            while (true) {
Chenggang Zhao's avatar
Chenggang Zhao committed
1497
                int num_used_slots = cached_channel_tail_idx - cached_channel_head_idx;
root's avatar
root committed
1498
1499
1500
1501
                is_lane_ready = lane_id < kNumRDMARanks and token_start_idx < token_end_idx and num_max_nvl_chunked_recv_tokens_per_rdma - num_used_slots >= num_max_nvl_chunked_send_tokens;
                if(__any_sync(kFirstHalfMask, is_lane_ready))
                    break;
                if(__any_sync(kSecondHalfMask, is_lane_ready))
Chenggang Zhao's avatar
Chenggang Zhao committed
1502
1503
                    break;
                // Retry
root's avatar
root committed
1504
                if (lane_id < kNumRDMARanks and token_start_idx < token_end_idx)
lijian6's avatar
lijian6 committed
1505
                    cached_channel_head_idx = ld_volatile_global(nvl_channel_head.buffer() + lane_id);
Chenggang Zhao's avatar
Chenggang Zhao committed
1506
1507

                // Timeout check
root's avatar
root committed
1508
1509
1510
1511
1512
                long long int elapsed_time = wall_clock64() > start_time ? wall_clock64() - start_time : 0;
                if (elapsed_time > NUM_TIMEOUT_CYCLES and lane_id < kNumRDMARanks) {
                    printf("DeepEP combine NVL sender timeout, channel: %d, RDMA: %d, nvl: %d, dst NVL: %d, RDMA lane: %d, head: %d, tail: %d, start: %d, end: %d\n",
                           channel_id, rdma_rank, nvl_rank, dst_nvl_rank, lane_id, ld_volatile_global(nvl_channel_head.buffer() + lane_id), cached_channel_tail_idx,
                           token_start_idx, token_end_idx);
Chenggang Zhao's avatar
Chenggang Zhao committed
1513
1514
                    trap();
                }
root's avatar
root committed
1515
                __builtin_amdgcn_s_sleep(1);
Chenggang Zhao's avatar
Chenggang Zhao committed
1516
1517
1518
            }

            // Sync token start index and count
root's avatar
root committed
1519
1520
            for (int current_rdma_idx = 0; current_rdma_idx < kNumRDMARanks; ++ current_rdma_idx) {
                if (shfl_sync((token_start_idx >= token_end_idx) or (not is_lane_ready), current_rdma_idx, kEmulatedWarpSize))
Chenggang Zhao's avatar
Chenggang Zhao committed
1521
1522
1523
                    continue;

                // Sync token start index
root's avatar
root committed
1524
1525
                auto token_idx = static_cast<int64_t>(shfl_sync(token_start_idx, current_rdma_idx, kEmulatedWarpSize));
                int num_tokens_in_chunk = shfl_sync(min(num_max_nvl_chunked_send_tokens, token_end_idx - token_start_idx), current_rdma_idx, kEmulatedWarpSize);
Chenggang Zhao's avatar
Chenggang Zhao committed
1526
1527

                // Send by chunk
root's avatar
root committed
1528
                for (int chunk_idx = 0; chunk_idx < num_tokens_in_chunk; ++ chunk_idx, ++ token_idx) {
Chenggang Zhao's avatar
Chenggang Zhao committed
1529
1530
                    // Get an empty slot
                    int dst_slot_idx = 0;
root's avatar
root committed
1531
1532
                    if (lane_id == current_rdma_idx) {
                        dst_slot_idx = (cached_channel_tail_idx ++) % num_max_nvl_chunked_recv_tokens_per_rdma;
lijian6's avatar
lijian6 committed
1533
                        dst_slot_idx = current_rdma_idx * num_max_nvl_chunked_recv_tokens_per_rdma + dst_slot_idx;
Chenggang Zhao's avatar
Chenggang Zhao committed
1534
                    }
root's avatar
root committed
1535
                    dst_slot_idx = shfl_sync(dst_slot_idx, current_rdma_idx, kEmulatedWarpSize);
lijian6's avatar
lijian6 committed
1536
1537
1538

                    // Copy data
                    auto shifted_x_buffers = nvl_channel_x.buffer() + dst_slot_idx * hidden_int4;
root's avatar
root committed
1539
1540
                    auto shifted_x = x + token_idx * hidden_int4;
                    UNROLLED_WARP_COPY_EMULATED(5, lane_id, hidden_int4, shifted_x_buffers, shifted_x, ld_nc_global, st_na_global);
Chenggang Zhao's avatar
Chenggang Zhao committed
1541

lijian6's avatar
lijian6 committed
1542
                    // Copy source meta
root's avatar
root committed
1543
                    if (lane_id == num_topk)
lijian6's avatar
lijian6 committed
1544
                        st_na_global(nvl_channel_src_meta.buffer() + dst_slot_idx, ld_nc_global(src_meta + token_idx));
Chenggang Zhao's avatar
Chenggang Zhao committed
1545

lijian6's avatar
lijian6 committed
1546
                    // Copy `topk_weights`
root's avatar
root committed
1547
1548
                    if (lane_id < num_topk)
                        st_na_global(nvl_channel_topk_weights.buffer() + dst_slot_idx * num_topk + lane_id, ld_nc_global(topk_weights + token_idx * num_topk + lane_id));
Chenggang Zhao's avatar
Chenggang Zhao committed
1549
1550
1551
1552
1553
                }
                lane_id == current_rdma_idx ? (token_start_idx = static_cast<int>(token_idx)) : 0;
            }

            // Move queue tail
lijian6's avatar
lijian6 committed
1554
            syncwarp();
root's avatar
root committed
1555
1556
            if (lane_id < kNumRDMARanks and is_lane_ready)
                st_relaxed_sys_global(nvl_channel_tail.buffer() + lane_id, cached_channel_tail_idx);
Chenggang Zhao's avatar
Chenggang Zhao committed
1557
1558
        }
    } else {
root's avatar
root committed
1559
        auto lane_id = get_lane_id() % kWarpHyb;
Chenggang Zhao's avatar
Chenggang Zhao committed
1560
1561
        // Combiners and coordinators
        // RDMA symmetric layout
lijian6's avatar
lijian6 committed
1562
        auto hidden_bytes = hidden_int4 * sizeof(int4);
lijian6's avatar
lijian6 committed
1563
        auto num_bytes_per_rdma_token = get_num_bytes_per_rdma_token(hidden_int4, 0, 0, num_topk);
lijian6's avatar
lijian6 committed
1564
1565
1566
        auto rdma_channel_data = SymBuffer<int8_t>(rdma_buffer_ptr, num_max_rdma_chunked_recv_tokens * num_bytes_per_rdma_token, kNumRDMARanks, channel_id, num_channels);
        auto rdma_channel_head = SymBuffer<uint64_t, false>(rdma_buffer_ptr, 1, kNumRDMARanks, channel_id, num_channels);
        auto rdma_channel_tail = SymBuffer<uint64_t, false>(rdma_buffer_ptr, 1, kNumRDMARanks, channel_id, num_channels);
Chenggang Zhao's avatar
Chenggang Zhao committed
1567
1568

        // NVL layouts
lijian6's avatar
lijian6 committed
1569
1570
1571
1572
        void* local_nvl_buffer = buffer_ptrs[nvl_rank];
        void* nvl_buffers[NUM_MAX_NVL_PEERS];
        #pragma unroll
        for (int i = 0; i < NUM_MAX_NVL_PEERS; ++ i)
Chenggang Zhao's avatar
Chenggang Zhao committed
1573
            nvl_buffers[i] = buffer_ptrs[i];
lijian6's avatar
lijian6 committed
1574
1575
1576
1577
1578
        auto nvl_channel_x = AsymBuffer<int4>(local_nvl_buffer, num_max_nvl_chunked_recv_tokens * hidden_int4, NUM_MAX_NVL_PEERS, channel_id, num_channels).advance_also<NUM_MAX_NVL_PEERS>(nvl_buffers);
        auto nvl_channel_src_meta = AsymBuffer<SourceMeta>(local_nvl_buffer, num_max_nvl_chunked_recv_tokens, NUM_MAX_NVL_PEERS, channel_id, num_channels).advance_also<NUM_MAX_NVL_PEERS>(nvl_buffers);
        auto nvl_channel_topk_weights = AsymBuffer<float>(local_nvl_buffer, num_max_nvl_chunked_recv_tokens * num_topk, NUM_MAX_NVL_PEERS, channel_id, num_channels).advance_also<NUM_MAX_NVL_PEERS>(nvl_buffers);
        auto nvl_channel_head = AsymBuffer<int, NUM_MAX_NVL_PEERS>(nvl_buffers, kNumRDMARanks, NUM_MAX_NVL_PEERS, channel_id, num_channels, nvl_rank).advance_also(local_nvl_buffer);
        auto nvl_channel_tail = AsymBuffer<int>(local_nvl_buffer, kNumRDMARanks, NUM_MAX_NVL_PEERS, channel_id, num_channels).advance_also<NUM_MAX_NVL_PEERS>(nvl_buffers);
Chenggang Zhao's avatar
Chenggang Zhao committed
1579
1580

        // Combiner warp synchronization
lijian6's avatar
lijian6 committed
1581
        __shared__ volatile int forwarder_nvl_head[kNumForwarders][NUM_MAX_NVL_PEERS];
Chenggang Zhao's avatar
Chenggang Zhao committed
1582
        __shared__ volatile bool forwarder_retired[kNumForwarders];
lijian6's avatar
lijian6 committed
1583
        __shared__ volatile int rdma_receiver_rdma_head[kNumRDMAReceivers][kNumRDMARanks];
Chenggang Zhao's avatar
Chenggang Zhao committed
1584
        __shared__ volatile bool rdma_receiver_retired[kNumRDMAReceivers];
lijian6's avatar
lijian6 committed
1585

root's avatar
root committed
1586
1587
1588
1589
1590
1591
1592
1593
1594
1595
1596
1597
1598
1599
1600
        auto sync_forwarder_smem = [&]() {
            if (lane_id==0) {
                volatile int ret = __hip_atomic_fetch_add(&rdma_forwarder_counter[0], 1, __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_WORKGROUP);
            }
            syncwarp();
            while(rdma_forwarder_counter[0]<(kNumForwarders + 1)){}
        };
        auto sync_rdma_receiver_smem = [&]() {
            if (lane_id==0) {
                volatile int ret = __hip_atomic_fetch_add(&rdma_receiver_counter[0], 1, __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_WORKGROUP);
            }
            syncwarp();
            while(rdma_receiver_counter[0]<(kNumRDMAReceivers+1)){} 
        };

Chenggang Zhao's avatar
Chenggang Zhao committed
1601
1602
1603
        if (warp_role == WarpRole::kNVLAndRDMAForwarder) {
            // Receive from NVL ranks and forward to RDMA ranks
            // NOTES: this part is using "large warps" for each RDMA ranks
root's avatar
root committed
1604
1605
1606
            const auto dst_rdma_rank = warp_id / kNumWarpsPerForwarder;
            const auto sub_warp_id = warp_id % kNumWarpsPerForwarder;
            auto send_buffer = dst_rdma_rank == rdma_rank ? rdma_channel_data.recv_buffer(dst_rdma_rank) : rdma_channel_data.send_buffer(dst_rdma_rank);
lijian6's avatar
lijian6 committed
1607
            auto sync_large_warp = [=](const int iter, const int mode) {
Chenggang Zhao's avatar
Chenggang Zhao committed
1608
                if (kNumWarpsPerForwarder == 1) {
lijian6's avatar
lijian6 committed
1609
                    syncwarp();
Chenggang Zhao's avatar
Chenggang Zhao committed
1610
                } else {
root's avatar
root committed
1611
1612
1613
1614
1615
1616
1617
1618
1619
1620
1621
1622
1623
1624
1625
1626
1627
1628
1629
                    // LDS index to store for sync
                    int lds_dst_rdma_rank = dst_rdma_rank + (iter % num_sync_large_iteration) * kNumRDMARanks + mode * kNumRDMARanks * num_sync_large_iteration;
                    //reset index in the LDS to avoid race condition due to warp scheduling
                    int reset_idx =         dst_rdma_rank + ((iter + num_sync_large_iteration/2) % num_sync_large_iteration) * kNumRDMARanks + mode * kNumRDMARanks * num_sync_large_iteration;
                    // if (lane_id==0)
                    //     printf("rank %d dst_rdma_rank %d iter %d  warp_id %d  val %d\n", rank, dst_rdma_rank, iter, warp_id, sync_large_warp_counters[lds_dst_rdma_rank]);
                    auto start_time = clock64();
                    if (lane_id == 0){
                        volatile int ret = __hip_atomic_fetch_add(
                            &sync_large_warp_counters[lds_dst_rdma_rank], 1,
                            __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_WORKGROUP);
                    }
                    syncwarp();
                    //The while(...) loop polls the counter until all warps have arrived
                    if (lane_id == 0){
                        while (sync_large_warp_counters[lds_dst_rdma_rank] < (kNumWarpsPerForwarder)){
                            if (clock64() - start_time > NUM_TIMEOUT_CYCLES) {
                                printf("DeepEP combine sync timeout. current num_sync_large_iteration %d. double it.\n", num_sync_large_iteration );
                                trap();
lijian6's avatar
lijian6 committed
1630
1631
                            }
                        }
root's avatar
root committed
1632
1633
1634
1635
1636
1637
                    }
                    syncwarp();
                    if (lane_id == 0 && sync_large_warp_counters[reset_idx] == kNumWarpsPerForwarder){
                        sync_large_warp_counters[reset_idx] = 0;
                    }
                    syncwarp();
Chenggang Zhao's avatar
Chenggang Zhao committed
1638
1639
                }
            };
root's avatar
root committed
1640
            EP_STATIC_ASSERT(kNumWarpsPerForwarder == 1 or kNumRDMARanks + 2 <= 16, "Barriers are not enough");
1641

root's avatar
root committed
1642
1643
1644
            // In case of running less than 8 nodes
            constexpr bool kUseWave = (kNumRDMARanks <= 8);
            // Advance to the corresponding NVL buffer
lijian6's avatar
lijian6 committed
1645
            nvl_channel_x.advance(dst_rdma_rank * num_max_nvl_chunked_recv_tokens_per_rdma * hidden_int4);
lijian6's avatar
lijian6 committed
1646
            nvl_channel_src_meta.advance(dst_rdma_rank * num_max_nvl_chunked_recv_tokens_per_rdma);
lijian6's avatar
lijian6 committed
1647
            nvl_channel_topk_weights.advance(dst_rdma_rank * num_max_nvl_chunked_recv_tokens_per_rdma * num_topk);
Chenggang Zhao's avatar
Chenggang Zhao committed
1648
1649
1650
1651
            nvl_channel_head.advance(dst_rdma_rank);
            nvl_channel_tail.advance(dst_rdma_rank);

            // Clean shared memory and sync
lijian6's avatar
lijian6 committed
1652
            EP_STATIC_ASSERT(NUM_MAX_NVL_PEERS <= kWarpSize, "Invalid number of NVL peers");
root's avatar
root committed
1653
1654
1655
            lane_id < NUM_MAX_NVL_PEERS ? (forwarder_nvl_head[warp_id][lane_id] = 0) : 0;
            lane_id == 0 ? (forwarder_retired[warp_id] = false) : false;
            sync_forwarder_smem();
Chenggang Zhao's avatar
Chenggang Zhao committed
1656
1657
1658

            // Get count and cached head
            int cached_nvl_channel_tail_idx = 0;
lijian6's avatar
lijian6 committed
1659
1660
            int num_tokens_to_combine = rdma_channel_prefix_matrix[dst_rdma_rank * num_channels + channel_id];
            int num_tokens_prefix = channel_id == 0 ? 0 : rdma_channel_prefix_matrix[dst_rdma_rank * num_channels + channel_id - 1];
Chenggang Zhao's avatar
Chenggang Zhao committed
1661
1662
1663
1664
1665
            num_tokens_to_combine -= num_tokens_prefix;
            num_tokens_prefix += dst_rdma_rank == 0 ? 0 : rdma_rank_prefix_sum[dst_rdma_rank - 1];
            combined_nvl_head += num_tokens_prefix * NUM_MAX_NVL_PEERS;

            // Iterate over all tokens and combine by chunks
root's avatar
root committed
1666
            for (int token_start_idx = 0; token_start_idx < num_tokens_to_combine; token_start_idx += num_max_rdma_chunked_send_tokens) {
Chenggang Zhao's avatar
Chenggang Zhao committed
1667
                // Check destination queue emptiness, or wait a buffer to be released
root's avatar
root committed
1668
                auto token_end_idx = min(token_start_idx + num_max_rdma_chunked_send_tokens, num_tokens_to_combine);
Chenggang Zhao's avatar
Chenggang Zhao committed
1669
                auto num_chunked_tokens = token_end_idx - token_start_idx;
root's avatar
root committed
1670
1671
                auto start_time = wall_clock64();
                while (sub_warp_id == 0 and lane_id == 0) {
lijian6's avatar
lijian6 committed
1672
1673
1674
                    // Inequality: `num_max_rdma_chunked_recv_tokens - (tail - head) >= num_chunked_tokens`
                    // Here, `token_start_idx` is the actual tail
                    int num_used_slots = token_start_idx - ld_volatile_global(rdma_channel_head.buffer(dst_rdma_rank));
root's avatar
root committed
1675
                    if (num_max_rdma_chunked_recv_tokens - num_used_slots >= num_chunked_tokens)
Chenggang Zhao's avatar
Chenggang Zhao committed
1676
1677
1678
                        break;

                    // Timeout check
root's avatar
root committed
1679
1680
                    long long int elapsed_time = wall_clock64() > start_time ? wall_clock64() - start_time : 0;
                    if (elapsed_time > NUM_TIMEOUT_CYCLES) {
lijian6's avatar
lijian6 committed
1681
                        printf("DeepEP combine forwarder (RDMA check) timeout, channel: %d, RDMA: %d, nvl: %d, dst RDMA: %d, head: %ld, tail: %d, chunked: %d\n",
root's avatar
root committed
1682
                               channel_id, rdma_rank, nvl_rank, dst_rdma_rank, ld_volatile_global(rdma_channel_head.buffer(dst_rdma_rank)), token_start_idx, num_chunked_tokens);
Chenggang Zhao's avatar
Chenggang Zhao committed
1683
1684
1685
                        trap();
                    }
                }
lijian6's avatar
lijian6 committed
1686
                sync_large_warp(token_start_idx, 0);
Chenggang Zhao's avatar
Chenggang Zhao committed
1687
                // Combine and write to the RDMA buffer
root's avatar
root committed
1688
                for (int token_idx = token_start_idx + sub_warp_id; token_idx < token_end_idx; token_idx += kNumWarpsPerForwarder) {
Chenggang Zhao's avatar
Chenggang Zhao committed
1689
                    // Read expected head
lijian6's avatar
lijian6 committed
1690
                    EP_STATIC_ASSERT(kNumRDMARanks <= kWarpSize, "Invalid number of RDMA peers");
Chenggang Zhao's avatar
Chenggang Zhao committed
1691
                    int expected_head = -1;
root's avatar
root committed
1692
                    if (lane_id < NUM_MAX_NVL_PEERS)
lijian6's avatar
lijian6 committed
1693
                        expected_head = ld_nc_global(combined_nvl_head + token_idx * NUM_MAX_NVL_PEERS + lane_id);
Chenggang Zhao's avatar
Chenggang Zhao committed
1694
1695

                    // Wait lanes to be ready
lijian6's avatar
lijian6 committed
1696
                    start_time = wall_clock64();
root's avatar
root committed
1697
1698
                    while (cached_nvl_channel_tail_idx <= expected_head) {
                        cached_nvl_channel_tail_idx = ld_relaxed_sys_global(nvl_channel_tail.buffer(lane_id));
Chenggang Zhao's avatar
Chenggang Zhao committed
1699
                        // Timeout check
root's avatar
root committed
1700
1701
                        long long int elapsed_time = wall_clock64() > start_time ? wall_clock64() - start_time : 0;
                        if (elapsed_time > NUM_TIMEOUT_CYCLES and lane_id < NUM_MAX_NVL_PEERS) {
lijian6's avatar
lijian6 committed
1702
                            printf("DeepEP combine forwarder (NVL check) timeout, channel: %d, RDMA: %d, nvl: %d, src NVL: %d, dst RDMA: %d, tail: %d, waiting: %d, total: %d, sub: %d, large: %d, expected: %d\n",
root's avatar
root committed
1703
                                   channel_id, rdma_rank, nvl_rank, lane_id, dst_rdma_rank, cached_nvl_channel_tail_idx, token_idx, num_tokens_to_combine, sub_warp_id, kNumWarpsPerForwarder, expected_head);
Chenggang Zhao's avatar
Chenggang Zhao committed
1704
1705
                            trap();
                        }
root's avatar
root committed
1706
                        __builtin_amdgcn_s_sleep(1);
Chenggang Zhao's avatar
Chenggang Zhao committed
1707
1708
1709
                    }

                    // Combine current token
lijian6's avatar
lijian6 committed
1710
1711
                    auto rdma_slot_idx = token_idx % num_max_rdma_chunked_recv_tokens;
                    void* shifted = send_buffer + rdma_slot_idx * num_bytes_per_rdma_token;
root's avatar
root committed
1712
                    auto recv_fn = [&](int src_nvl_rank, int slot_idx, int hidden_int4_idx) -> int4 { return ld_nc_global(nvl_channel_x.buffer(src_nvl_rank) + slot_idx * hidden_int4 + hidden_int4_idx); };
lijian6's avatar
lijian6 committed
1713
                    auto recv_tw_fn = [&](int src_nvl_rank, int slot_idx, int topk_idx) -> float { return ld_nc_global(nvl_channel_topk_weights.buffer(src_nvl_rank) + slot_idx * num_topk + topk_idx); };
root's avatar
root committed
1714
1715
1716
1717
1718
1719
                    combine_token<NUM_MAX_NVL_PEERS, dtype_t, NUM_MAX_NVL_PEERS, kWarpHyb>(expected_head >= 0,
                                                                                     expected_head, lane_id,
                                                                                     hidden_int4, num_topk,
                                                                                     reinterpret_cast<int4*>(shifted),
                                                                                     reinterpret_cast<float*>(reinterpret_cast<int8_t*>(shifted) + hidden_bytes + sizeof(SourceMeta)),
                                                                                     num_max_nvl_chunked_recv_tokens_per_rdma, recv_fn, recv_tw_fn);
Chenggang Zhao's avatar
Chenggang Zhao committed
1720
1721

                    // Update head
root's avatar
root committed
1722
1723
                    if (lane_id < NUM_MAX_NVL_PEERS)
                        expected_head < 0 ? (forwarder_nvl_head[warp_id][lane_id] = -expected_head - 1) : (forwarder_nvl_head[warp_id][lane_id] = expected_head + 1);
Chenggang Zhao's avatar
Chenggang Zhao committed
1724
                }
lijian6's avatar
lijian6 committed
1725
                sync_large_warp(token_start_idx, 1);
Chenggang Zhao's avatar
Chenggang Zhao committed
1726
                // Issue RDMA send
root's avatar
root committed
1727
1728
1729
1730
                // TODO: Switch back to put_nbi_wave function
              
                if (sub_warp_id == kNumWarpsPerForwarder - 1 ) {
                    if (dst_rdma_rank != rdma_rank) {
Chenggang Zhao's avatar
Chenggang Zhao committed
1731
                        auto rdma_slot_idx = token_start_idx % num_max_rdma_chunked_recv_tokens;
root's avatar
root committed
1732
                    #ifdef FORCE_DUSHMEM_API
1733
                        shmemx_int8_put_nbi_warp(
lijian6's avatar
lijian6 committed
1734
                            rdma_channel_data.recv_buffer(rdma_rank) +
root's avatar
root committed
1735
1736
1737
1738
1739
1740
1741
1742
1743
1744
1745
1746
1747
1748
1749
1750
1751
1752
1753
1754
1755
1756
1757
1758
1759
1760
1761
1762
1763
                                    rdma_slot_idx * num_bytes_per_rdma_token,
                                rdma_channel_data.send_buffer(dst_rdma_rank) +
                                    rdma_slot_idx * num_bytes_per_rdma_token,
                                num_chunked_tokens * num_bytes_per_rdma_token,
                                translate_dst_rdma_rank<kLowLatencyMode>(dst_rdma_rank, nvl_rank));
                    #else
                        #if !defined(ROCM_DISABLE_CTX)
                        if constexpr (kUseWave){
                            shmem_ctx_schar_put_nbi_warp(ctx,
                                                 rdma_channel_data.recv_buffer(rdma_rank) + rdma_slot_idx * num_bytes_per_rdma_token,
                                                 rdma_channel_data.send_buffer(dst_rdma_rank) + rdma_slot_idx * num_bytes_per_rdma_token,
                                                 num_chunked_tokens * num_bytes_per_rdma_token,
                                                 translate_dst_rdma_rank<kLowLatencyMode>(dst_rdma_rank, nvl_rank));
                        }
                        #else
                        if constexpr (kUseWave){
                            shmemx_int8_put_nbi_warp(rdma_channel_data.recv_buffer(rdma_rank) + rdma_slot_idx * num_bytes_per_rdma_token,
                                                 rdma_channel_data.send_buffer(dst_rdma_rank) + rdma_slot_idx * num_bytes_per_rdma_token,
                                                 num_chunked_tokens * num_bytes_per_rdma_token,
                                                 translate_dst_rdma_rank<kLowLatencyMode>(dst_rdma_rank, nvl_rank));
                        } else {
                            if (lane_id == 0)
                                shmemx_int8_put_nbi(rdma_channel_data.recv_buffer(rdma_rank) + rdma_slot_idx * num_bytes_per_rdma_token,
                                                 rdma_channel_data.send_buffer(dst_rdma_rank) + rdma_slot_idx * num_bytes_per_rdma_token,
                                                 num_chunked_tokens * num_bytes_per_rdma_token,
                                                 translate_dst_rdma_rank<kLowLatencyMode>(dst_rdma_rank, nvl_rank));
                        }
                        #endif
                    #endif
lijian6's avatar
lijian6 committed
1764
#if !defined(FORCE_DUSHMEM_API) && !defined(ROCM_DISABLE_CTX)
root's avatar
root committed
1765
                        shmem_ctx_quiet(ctx);
1766
1767
1768
#else
                        shmem_fence();
#endif
Chenggang Zhao's avatar
Chenggang Zhao committed
1769
1770
1771
1772
1773
                    } else {
                        memory_fence();
                    }

                    // Write new RDMA tail
lijian6's avatar
lijian6 committed
1774
                    syncwarp();
root's avatar
root committed
1775
                    if (lane_id == 0)
lijian6's avatar
lijian6 committed
1776
#if !defined(FORCE_DUSHMEM_API) && !defined(ROCM_DISABLE_CTX)
1777
1778
1779
1780
                        shmem_ctx_ulong_atomic_add(ctx,
#else
                        shmem_signal_op_add(
#endif
root's avatar
root committed
1781
1782
                                           rdma_channel_tail.buffer(rdma_rank), num_chunked_tokens,
                                           translate_dst_rdma_rank<kLowLatencyMode>(dst_rdma_rank, nvl_rank));
Chenggang Zhao's avatar
Chenggang Zhao committed
1783
1784
1785
1786
                }
            }

            // Retired
lijian6's avatar
lijian6 committed
1787
            syncwarp();
root's avatar
root committed
1788
1789
1790
            if (lane_id == 0)
                forwarder_retired[warp_id] = true;
        } else if (warp_role == WarpRole::kRDMAReceiver) {
Chenggang Zhao's avatar
Chenggang Zhao committed
1791
1792
            // Receive from RDMA ranks and write to the output tensor
            // Clean shared memory and sync
lijian6's avatar
lijian6 committed
1793
            EP_DEVICE_ASSERT(kNumRDMARanks <= kWarpSize);
root's avatar
root committed
1794
1795
1796
            lane_id < kNumRDMARanks ? (rdma_receiver_rdma_head[warp_id][lane_id] = 0) : 0;
            lane_id == 0 ? (rdma_receiver_retired[warp_id] = false) : 0;
            sync_rdma_receiver_smem();
Chenggang Zhao's avatar
Chenggang Zhao committed
1797
1798
1799

            // The same tokens as the dispatch process
            int token_start_idx, token_end_idx;
lijian6's avatar
lijian6 committed
1800
            get_channel_task_range(num_combined_tokens, num_channels, channel_id, token_start_idx, token_end_idx);
Chenggang Zhao's avatar
Chenggang Zhao committed
1801

root's avatar
root committed
1802
            // Iterate over all tokens and combine
Chenggang Zhao's avatar
Chenggang Zhao committed
1803
            int cached_channel_tail_idx = 0;
root's avatar
root committed
1804
1805
1806
1807
1808
1809
1810
            for (int64_t token_idx = token_start_idx + warp_id; token_idx < token_end_idx; token_idx += kNumRDMAReceivers) {
                // Read expected head
                EP_STATIC_ASSERT(kNumRDMARanks <= kWarpSize, "Invalid number of RDMA peers");
                int expected_head = -1;
                if (lane_id < kNumRDMARanks) {
                    expected_head = ld_nc_global(combined_rdma_head + token_idx * kNumRDMARanks + lane_id);
                    (expected_head < 0) ? (rdma_receiver_rdma_head[warp_id][lane_id] = -expected_head - 1) : (rdma_receiver_rdma_head[warp_id][lane_id] = expected_head);
lishen's avatar
lishen committed
1811
                }
Chenggang Zhao's avatar
Chenggang Zhao committed
1812

root's avatar
root committed
1813
1814
1815
1816
1817
1818
1819
1820
1821
1822
1823
                // Wait lanes to be ready
                auto start_time = wall_clock64();
                while (cached_channel_tail_idx <= expected_head) {
                    cached_channel_tail_idx = static_cast<int>(ld_acquire_sys_global(rdma_channel_tail.buffer(lane_id)));

                    // Timeout check
                    long long int elapsed_time = wall_clock64() > start_time ? wall_clock64() - start_time : 0;
                    if (elapsed_time > NUM_TIMEOUT_CYCLES) {
                        printf("DeepEP combine RDMA receiver timeout, channel: %d, RDMA: %d, nvl: %d, src RDMA: %d, tail: %d, waiting: %ld, expect: %d\n",
                               channel_id, rdma_rank, nvl_rank, lane_id, cached_channel_tail_idx, token_idx, expected_head);
                        trap();
Chenggang Zhao's avatar
Chenggang Zhao committed
1824
                    }
root's avatar
root committed
1825
                    __builtin_amdgcn_s_sleep(1);
Chenggang Zhao's avatar
Chenggang Zhao committed
1826
                }
lijian6's avatar
lijian6 committed
1827
                syncwarp();
Chenggang Zhao's avatar
Chenggang Zhao committed
1828

root's avatar
root committed
1829
1830
1831
1832
1833
1834
1835
1836
1837
                // Combine current token
                auto recv_fn = [&](int src_rdma_rank, int slot_idx, int hidden_int4_idx) -> int4 { return ld_nc_global(reinterpret_cast<const int4*>(rdma_channel_data.recv_buffer(src_rdma_rank) + slot_idx * num_bytes_per_rdma_token) + hidden_int4_idx);};
                auto recv_tw_fn = [&](int src_rdma_rank, int slot_idx, int topk_idx) -> float { return ld_nc_global(reinterpret_cast<const float*>(rdma_channel_data.recv_buffer(src_rdma_rank) + slot_idx * num_bytes_per_rdma_token + hidden_bytes + sizeof(SourceMeta)) + topk_idx);};
                combine_token<kNumRDMARanks, dtype_t, kNumTopkRDMARanks, kWarpHyb>(expected_head >= 0,
                                                                                            expected_head, lane_id,
                                                                                            hidden_int4, num_topk,
                                                                                            combined_x + token_idx * hidden_int4,
                                                                                            combined_topk_weights + token_idx * num_topk,
                                                                                            num_max_rdma_chunked_recv_tokens, recv_fn, recv_tw_fn);
Chenggang Zhao's avatar
Chenggang Zhao committed
1838
1839
1840
            }

            // Retired
lijian6's avatar
lijian6 committed
1841
            syncwarp();
root's avatar
root committed
1842
1843
1844
1845
            if (lane_id == 0)
                rdma_receiver_retired[warp_id] = true;
        } else {
            auto lane_id = get_lane_id();
Chenggang Zhao's avatar
Chenggang Zhao committed
1846
1847
            // Coordinator
            // Sync shared memory status
root's avatar
root committed
1848
            is_rdma_receiver_sm ? sync_rdma_receiver_smem() : sync_forwarder_smem();
Chenggang Zhao's avatar
Chenggang Zhao committed
1849
1850
            const auto num_warps_per_rdma_rank = kNumForwarders / kNumRDMARanks;

root's avatar
root committed
1851
            int last_rdma_head = 0;
Chenggang Zhao's avatar
Chenggang Zhao committed
1852
            int last_nvl_head[kNumRDMARanks] = {0};
root's avatar
root committed
1853
1854
1855
1856
            int dst_rdma_rank = lane_id < kNumRDMARanks ? lane_id : 0;
            int dst_nvl_rank = lane_id < NUM_MAX_NVL_PEERS ? lane_id : 0;
            EP_STATIC_ASSERT(kNumCombineForwarderWarps <= kWarpSize, "Invalid number of forwarder warps");
            while (true) {
Chenggang Zhao's avatar
Chenggang Zhao committed
1857
                // Retired
root's avatar
root committed
1858
1859
1860
                if (is_rdma_receiver_sm and __all_sync(kFullWarpMask, lane_id >= kNumRDMAReceivers or rdma_receiver_retired[lane_id]))
                    break;
                if (not is_rdma_receiver_sm and __all_sync(kFullWarpMask, lane_id >= kNumForwarders or forwarder_retired[lane_id]))
Chenggang Zhao's avatar
Chenggang Zhao committed
1861
1862
1863
                    break;

                // Find minimum head for RDMA ranks
root's avatar
root committed
1864
                if (is_rdma_receiver_sm) {
Chenggang Zhao's avatar
Chenggang Zhao committed
1865
                    int min_head = std::numeric_limits<int>::max();
root's avatar
root committed
1866
1867
1868
                    #pragma unroll
                    for (int i = 0; i < kNumRDMAReceivers; ++ i) if (not rdma_receiver_retired[i])
                        min_head = min(min_head, rdma_receiver_rdma_head[i][dst_rdma_rank]);
lijian6's avatar
lijian6 committed
1869
                    if (min_head != std::numeric_limits<int>::max() and min_head >= last_rdma_head + num_max_rdma_chunked_send_tokens and lane_id < kNumRDMARanks) {
lijian6's avatar
lijian6 committed
1870
#if !defined(FORCE_DUSHMEM_API) && !defined(ROCM_DISABLE_CTX)
1871
1872
1873
1874
1875
                        shmem_ctx_ulong_atomic_add(ctx,
#else
                        shmem_signal_op_add(
#endif
                            rdma_channel_head.buffer(rdma_rank), min_head - last_rdma_head,
root's avatar
root committed
1876
                                           translate_dst_rdma_rank<kLowLatencyMode>(dst_rdma_rank, nvl_rank));
lijian6's avatar
lijian6 committed
1877

1878
1879
                        last_rdma_head = min_head;
                    }
root's avatar
root committed
1880
1881
1882
1883
1884
1885
1886
1887
1888
1889
1890
                } else {
                    // Find minimum head for NVL ranks
                    #pragma unroll
                    for (int i = 0; i < kNumRDMARanks; ++ i) {
                        int min_head = std::numeric_limits<int>::max();
                        #pragma unroll
                        for (int j = 0; j < num_warps_per_rdma_rank; ++ j) if (not forwarder_retired[i * num_warps_per_rdma_rank + j])
                            min_head = min(min_head, forwarder_nvl_head[i * num_warps_per_rdma_rank + j][dst_nvl_rank]);
                        if (min_head != std::numeric_limits<int>::max() and min_head > last_nvl_head[i] and lane_id < NUM_MAX_NVL_PEERS)
                            st_relaxed_sys_global(nvl_channel_head.buffer_by(dst_nvl_rank) + i, last_nvl_head[i] = min_head);
                    }
Chenggang Zhao's avatar
Chenggang Zhao committed
1891
1892
1893
                }

                // Nanosleep and let other warps work
lijian6's avatar
lijian6 committed
1894
                __builtin_amdgcn_s_sleep(NUM_WAIT_CYCLES_TIMES_64);
Chenggang Zhao's avatar
Chenggang Zhao committed
1895
1896
1897
            }
        }
    }
lijian6's avatar
lijian6 committed
1898
#if !defined(FORCE_DUSHMEM_API) && !defined(ROCM_DISABLE_CTX)
1899
1900
    shmem_wg_ctx_destroy(&ctx);
#endif
Chenggang Zhao's avatar
Chenggang Zhao committed
1901
1902
}

lijian6's avatar
lijian6 committed
1903

root's avatar
root committed
1904
1905
1906
1907
1908
1909
1910
1911
1912
1913
1914
1915
1916
1917
1918
1919
1920
1921
1922
1923
1924
1925
1926
1927
1928
1929
1930
1931
1932
1933
1934
1935
1936
void combine(hipDataType type,
             void* combined_x,
             float* combined_topk_weights,
             const bool* is_combined_token_in_rank,
             const void* x,
             const float* topk_weights,
             const void* bias_0,
             const void* bias_1,
             const int* combined_rdma_head,
             const int* combined_nvl_head,
             const void* src_meta,
             const int* rdma_channel_prefix_matrix,
             const int* rdma_rank_prefix_sum,
             const int* gbl_channel_prefix_matrix,
             int num_tokens,
             int num_combined_tokens,
             int hidden,
             int num_topk,
             void* rdma_buffer_ptr,
             int num_max_rdma_chunked_send_tokens,
             int num_max_rdma_chunked_recv_tokens,
             void** buffer_ptrs,
             int num_max_nvl_chunked_send_tokens,
             int num_max_nvl_chunked_recv_tokens,
             int rank,
             int num_ranks,
             hipStream_t stream,
             int num_channels,
             bool low_latency_mode)
{
    const int num_rdma_ranks = num_ranks / NUM_MAX_NVL_PEERS;

    EP_HOST_ASSERT(num_rdma_ranks > 0);
Chenggang Zhao's avatar
Chenggang Zhao committed
1937
    EP_HOST_ASSERT(num_max_nvl_chunked_recv_tokens % num_rdma_ranks == 0);
lijian6's avatar
lijian6 committed
1938
    EP_HOST_ASSERT(num_max_nvl_chunked_recv_tokens / num_rdma_ranks > std::max(num_max_rdma_chunked_send_tokens, num_max_nvl_chunked_send_tokens));
lijian6's avatar
lijian6 committed
1939
    EP_HOST_ASSERT(type == HIP_R_16BF);
Chenggang Zhao's avatar
Chenggang Zhao committed
1940

root's avatar
root committed
1941
1942
1943
1944
1945
1946
1947
1948
1949
1950
1951
1952
1953
1954
1955
1956
1957
1958
1959
1960
1961
1962
1963
1964
1965
1966
1967
1968
1969
1970
1971
1972
1973
1974
1975
1976
    // One case per compile-time NR specialization.
#define COMBINE_LAUNCH_CASE(NR) {                                                        \
    /* Per-case compile-time constants */                                                \
    constexpr int kNumCombineForwarderWarps = (NR < 9) ? 10 : 16;                        \
    constexpr int kWarpsPerForwarder = (kNumCombineForwarderWarps/NR) > 0                \
                                         ? (kNumCombineForwarderWarps/NR) : 1;           \
    constexpr int kNumForwarders = NR * kWarpsPerForwarder;                              \
    constexpr int kBlockThreads = (NR > 8)                                               \
        ? ((NUM_MAX_NVL_PEERS + kNumForwarders) * kEmulatedWarpSize + kWarpSize)         \
        : ((NUM_MAX_NVL_PEERS/2 + 1 + kNumForwarders) * kWarpSize);                      \
                                                                                         \
    SETUP_LAUNCH_CONFIG(num_channels * 2, kBlockThreads, stream);                        \
                                                                                         \
    using scalar_t =  hip_bfloat16;                                                      \
    auto fn = low_latency_mode                                                           \
        ? combine<true,  NR, scalar_t, kNumCombineForwarderWarps>                        \
        : combine<false, NR, scalar_t, kNumCombineForwarderWarps>;                       \
                                                                                         \
    /* Launch (backend-specific) */                                                      \
                                                                             \
    LAUNCH_KERNEL_NON_COOPERATIVE(&cfg, fn,                                                            \
            reinterpret_cast<int4*>(combined_x), combined_topk_weights, is_combined_token_in_rank, \
            reinterpret_cast<const int4*>(x), topk_weights,  \
            reinterpret_cast<const int4 *>(bias_0), reinterpret_cast<const int4 *>(bias_1), \
            combined_rdma_head, combined_nvl_head,                                       \
            reinterpret_cast<const SourceMeta*>(src_meta), rdma_channel_prefix_matrix,   \
            rdma_rank_prefix_sum, gbl_channel_prefix_matrix,                             \
            num_tokens, num_combined_tokens, hidden, num_topk,                           \
            rdma_buffer_ptr, num_max_rdma_chunked_send_tokens,                           \
            num_max_rdma_chunked_recv_tokens,                                            \
            buffer_ptrs, num_max_nvl_chunked_send_tokens,                                \
            num_max_nvl_chunked_recv_tokens,                                             \
            rank, num_ranks);                                                            \
} break

    // Dispatch on the runtime num_rdma_ranks, but each case is compile-time specialized.
Chenggang Zhao's avatar
Chenggang Zhao committed
1977
1978
1979
    SWITCH_RDMA_RANKS(COMBINE_LAUNCH_CASE);
#undef COMBINE_LAUNCH_CASE
}
lijian6's avatar
lijian6 committed
1980

Chenggang Zhao's avatar
Chenggang Zhao committed
1981
1982
1983
} // namespace internode

} // namespace deep_ep
lijian6's avatar
lijian6 committed
1984
1985

#endif