intranode.cu 47.8 KB
Newer Older
Chenggang Zhao's avatar
Chenggang Zhao committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
#include "configs.cuh"
#include "buffer.cuh"
#include "exception.cuh"
#include "launch.cuh"
#include "utils.cuh"

namespace deep_ep {

namespace intranode {

template<int kNumRanks>
__global__ void
notify_dispatch(const int* num_tokens_per_rank, int* moe_recv_counter_mapped,
                const int* num_tokens_per_expert, int* moe_recv_expert_counter_mapped, int num_experts,
                int num_tokens, int num_channels, const bool* is_token_in_rank, int* channel_prefix_matrix,
                int* rank_prefix_matrix_copy, int num_memset_int, int expert_alignment,
17
                void** buffer_ptrs, int** barrier_signal_ptrs, int rank) {
Chenggang Zhao's avatar
Chenggang Zhao committed
18
19
20
21
22
23
    auto sm_id = static_cast<int>(blockIdx.x);
    auto thread_id = static_cast<int>(threadIdx.x), num_threads = static_cast<int>(blockDim.x);
    auto lane_id = thread_id % 32, warp_id = thread_id / 32, num_warps = num_threads / 32;

    if (sm_id == 0) {
        // Barrier first
24
        barrier_block<kNumRanks, true>(barrier_signal_ptrs, rank);
Chenggang Zhao's avatar
Chenggang Zhao committed
25
26
27

        int *per_rank_buffer, *per_expert_buffer;
        if (thread_id < kNumRanks) {
28
            per_rank_buffer = static_cast<int*>(buffer_ptrs[thread_id]);
Chenggang Zhao's avatar
Chenggang Zhao committed
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
            per_expert_buffer = per_rank_buffer + kNumRanks * kNumRanks;
        }

        // After this loop:
        //  - `per_rank_buffer[rank][i, j]` means the number of tokens from rank i to rank j
        //  - `per_expert_buffer[rank][i, j]` means the number of tokens from rank i to local expert j
        int num_experts_per_rank = num_experts / kNumRanks;
        if (thread_id < kNumRanks) {
            #pragma unroll
            for (int i = 0; i < kNumRanks; ++ i)
                per_rank_buffer[rank * kNumRanks + i] = num_tokens_per_rank[i];
            #pragma unroll
            for (int i = 0; i < num_experts_per_rank; ++ i)
                per_expert_buffer[rank * num_experts_per_rank + i] = num_tokens_per_expert[thread_id * num_experts_per_rank + i];
        }

        // Wait for all ranks to be finished
46
        barrier_block<kNumRanks>(barrier_signal_ptrs, rank);
Chenggang Zhao's avatar
Chenggang Zhao committed
47
48
49

        // Sum per-rank counts and return to CPU
        // Also pre-compute the prefix sum for data sending
50
        auto local_per_rank_buffer = static_cast<int*>(buffer_ptrs[rank]);
Chenggang Zhao's avatar
Chenggang Zhao committed
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
        if (thread_id < kNumRanks) {
            #pragma unroll
            for (int i = 1; i < kNumRanks; ++ i)
                local_per_rank_buffer[i * kNumRanks + thread_id] += local_per_rank_buffer[(i - 1) * kNumRanks + thread_id];
            if (thread_id == rank)
                *moe_recv_counter_mapped = local_per_rank_buffer[(kNumRanks - 1) * kNumRanks + rank];
        }

        // Sum per-experts counts and return to CPU
        auto local_per_expert_buffer = local_per_rank_buffer + kNumRanks * kNumRanks;
        if (thread_id < num_experts_per_rank) {
            int sum = 0;
            #pragma unroll
            for (int i = 0; i < kNumRanks; ++ i)
                sum += local_per_expert_buffer[i * num_experts_per_rank + thread_id];
            sum = (sum + expert_alignment - 1) / expert_alignment * expert_alignment;
            moe_recv_expert_counter_mapped[thread_id] = sum;
        }
        __syncthreads();

        // Copy rank size prefix matrix to another tensor
        #pragma unroll
        for (int i = thread_id; i < kNumRanks * kNumRanks; i += num_threads)
            rank_prefix_matrix_copy[i] = local_per_rank_buffer[i];

        // Extra memset for later communication queue
        #pragma unroll
        for (int i = thread_id; i < num_memset_int; i += num_threads)
            local_per_expert_buffer[i] = 0;

        // Barrier
82
        barrier_block<kNumRanks>(barrier_signal_ptrs, rank);
Chenggang Zhao's avatar
Chenggang Zhao committed
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
    } else {
        int dst_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;
            get_channel_task_range(num_tokens, num_channels, channel_id, token_start_idx, token_end_idx);

            // Iterate over tokens
            int count = 0;
            for (int64_t i = token_start_idx + lane_id; i < token_end_idx; i += 32)
                count += is_token_in_rank[i * kNumRanks + dst_rank];
            count = warp_reduce_sum(count);
            if (lane_id == 0)
                channel_prefix_matrix[dst_rank * num_channels + channel_id] = count;
        }
        __syncthreads();

        // Pre-compute prefix sum for all channels
        if (thread_id == 0) {
            #pragma unroll
            for (int i = 1; i < num_channels; ++ i)
                channel_prefix_matrix[dst_rank * num_channels + i] += channel_prefix_matrix[dst_rank * num_channels + i - 1];
        }
    }
}

void notify_dispatch(const int* num_tokens_per_rank, int* moe_recv_counter_mapped, int num_ranks,
                     const int* num_tokens_per_expert, int* moe_recv_expert_counter_mapped, int num_experts,
                     int num_tokens, const bool* is_token_in_rank, int* channel_prefix_matrix,
                     int* rank_prefix_matrix_copy, int num_memset_int, int expert_alignment,
112
                     void** buffer_ptrs, int** barrier_signal_ptrs, int rank,
Chenggang Zhao's avatar
Chenggang Zhao committed
113
114
115
116
117
118
119
                     cudaStream_t stream, int num_channels) {
#define NOTIFY_DISPATCH_LAUNCH_CASE(ranks) \
    LAUNCH_KERNEL(&cfg, notify_dispatch<ranks>, \
        num_tokens_per_rank, moe_recv_counter_mapped, \
        num_tokens_per_expert, moe_recv_expert_counter_mapped, num_experts, \
        num_tokens, num_channels, is_token_in_rank, channel_prefix_matrix, \
        rank_prefix_matrix_copy, num_memset_int, expert_alignment, \
120
        buffer_ptrs, barrier_signal_ptrs, rank); \
Chenggang Zhao's avatar
Chenggang Zhao committed
121
122
123
124
125
126
127
128
129
130
131
132
133
134
    break

    constexpr int kNumThreads = 128;
    EP_HOST_ASSERT(num_experts % num_ranks == 0);
    EP_HOST_ASSERT(num_experts / num_ranks <= kNumThreads and num_ranks <= kNumThreads);

    SETUP_LAUNCH_CONFIG(1 + num_ranks, kNumThreads, stream);
    SWITCH_RANKS(NOTIFY_DISPATCH_LAUNCH_CASE);
#undef NOTIFY_DISPATCH_LAUNCH_CASE
}

template<int kNumRanks>
__global__ void
cached_notify_dispatch(const int* rank_prefix_matrix, int num_memset_int,
135
                       void** buffer_ptrs, int** barrier_signal_ptrs, int rank) {
Chenggang Zhao's avatar
Chenggang Zhao committed
136
    // A simplified version for cached handles
137
    barrier_block<kNumRanks, true>(barrier_signal_ptrs, rank);
Chenggang Zhao's avatar
Chenggang Zhao committed
138
139
140

    // Copy and clean
    auto thread_id = static_cast<int>(threadIdx.x), num_threads = static_cast<int>(blockDim.x);
141
    auto ptr = static_cast<int*>(buffer_ptrs[rank]);
Chenggang Zhao's avatar
Chenggang Zhao committed
142
143
144
145
146
147
148
149
    #pragma unroll
    for (int i = thread_id; i < kNumRanks * kNumRanks; i += num_threads)
        ptr[i] = rank_prefix_matrix[i];
    #pragma unroll
    for (int i = thread_id; i < num_memset_int; i += num_threads)
        ptr[kNumRanks * kNumRanks + i] = 0;

    // Barrier after cleaning
150
    barrier_block<kNumRanks>(barrier_signal_ptrs, rank);
Chenggang Zhao's avatar
Chenggang Zhao committed
151
152
153
}

void cached_notify_dispatch(const int* rank_prefix_matrix, int num_memset_int,
154
155
                            void** buffer_ptrs, int** barrier_signal_ptrs,
                            int rank, int num_ranks, cudaStream_t stream) {
Chenggang Zhao's avatar
Chenggang Zhao committed
156
157
#define CACHED_NOTIFY_DISPATCH_LAUNCH_CASE(ranks) \
    LAUNCH_KERNEL(&cfg, cached_notify_dispatch<ranks>, \
158
        rank_prefix_matrix, num_memset_int, buffer_ptrs, barrier_signal_ptrs, rank); \
Chenggang Zhao's avatar
Chenggang Zhao committed
159
160
161
162
163
164
165
    break

    SETUP_LAUNCH_CONFIG(1, 128, stream);
    SWITCH_RANKS(CACHED_NOTIFY_DISPATCH_LAUNCH_CASE);
#undef CACHED_NOTIFY_DISPATCH_LAUNCH_CASE
}

166
template <int kNumRanks, int kNumThreads, int kNumTMABytesPerWarp>
Chenggang Zhao's avatar
Chenggang Zhao committed
167
__global__ void __launch_bounds__(kNumThreads, 1)
Chenggang Zhao's avatar
Chenggang Zhao committed
168
169
170
dispatch(int4* recv_x, float* recv_x_scales, int* recv_src_idx, int64_t* recv_topk_idx, float* recv_topk_weights, int* recv_channel_offset,
         int* send_head, const int4* x, const float* x_scales, const int64_t* topk_idx, const float* topk_weights,
         const bool* is_token_in_rank, const int* channel_prefix_matrix,
171
         int num_tokens, int num_worst_tokens, int hidden_int4, int num_topk, int num_experts, int num_scales,
Shifang Xu's avatar
Shifang Xu committed
172
         int scale_token_stride, int scale_hidden_stride,
173
         void** buffer_ptrs, int rank,
Chenggang Zhao's avatar
Chenggang Zhao committed
174
175
         int num_max_send_tokens, int num_recv_buffer_tokens) {
    const auto num_sms = static_cast<int>(gridDim.x), sm_id = static_cast<int>(blockIdx.x);
176
    const auto thread_id = static_cast<int>(threadIdx.x), lane_id = get_lane_id();
Chenggang Zhao's avatar
Chenggang Zhao committed
177
178
179
    const bool is_sender = sm_id % 2 == 0;
    EP_DEVICE_ASSERT(num_sms % 2 == 0);

Chenggang Zhao's avatar
Chenggang Zhao committed
180
181
    // Several warps are response for a single rank
    const auto num_threads_per_rank = kNumThreads / kNumRanks;
Chenggang Zhao's avatar
Chenggang Zhao committed
182
    const auto num_channels = num_sms / 2;
Chenggang Zhao's avatar
Chenggang Zhao committed
183
184
    const auto responsible_rank = (static_cast<int>(thread_id)) / num_threads_per_rank;
    // Even-numbered blocks for sending, odd-numbered blocks for receiving.
Chenggang Zhao's avatar
Chenggang Zhao committed
185
186
187
188
189
190
191
192
193
194
    const auto responsible_channel = sm_id / 2;

    int num_experts_per_rank = num_experts / kNumRanks;
    EP_DEVICE_ASSERT(num_experts_per_rank > 0 or num_topk == 0);
    EP_DEVICE_ASSERT(num_topk <= 32);
    EP_DEVICE_ASSERT((topk_idx == nullptr)  == (topk_weights == nullptr));
    EP_DEVICE_ASSERT((recv_topk_idx == nullptr) == (recv_topk_weights == nullptr));

    // Calculate pointers by the specific layout
    // `rank_prefix_matrix`: kNumRanks * kNumRanks * sizeof(int)
195
    auto ptr = reinterpret_cast<void*>(static_cast<int8_t*>(buffer_ptrs[is_sender ? responsible_rank : rank]) + kNumRanks * kNumRanks * sizeof(int));
Chenggang Zhao's avatar
Chenggang Zhao committed
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
    int target_rank = is_sender ? rank : responsible_rank;
    auto num_channels_total = num_channels * kNumRanks;
    auto channel_rank_offset = responsible_channel * kNumRanks + target_rank;

    // Channel buffer metadata
    // Senders are responsible for tails, and receivers are responsible for heads
    // Stored on the receiver side
    // The retired signals are actually boolean flags, but to align with 16 bytes, we make it `int64_t`
    // `start_offset`: kNumChannels * kNumRanks * sizeof(int)
    // `end_offset`: kNumChannels * kNumRanks * sizeof(int)
    // `head_idx`: kNumChannels * kNumRanks * sizeof(int)
    // `tail_idx`: kNumChannels * kNumRanks * sizeof(int)
    auto channel_start_offset = Buffer<int>(ptr, num_channels_total, channel_rank_offset);
    auto channel_end_offset = Buffer<int>(ptr, num_channels_total, channel_rank_offset);
    auto channel_head_idx = Buffer<int>(ptr, num_channels_total, channel_rank_offset);
    auto channel_tail_idx = Buffer<int>(ptr, num_channels_total, channel_rank_offset);

    // Channel data buffers, stored on the receiver side
    // `x_buffers`: kNumChannels * kNumRanks * num_recv_buffer_tokens * hidden_int4 * sizeof(int4)
    // `src_idx_buffers`: kNumChannels * kNumRanks * num_recv_buffer_tokens * sizeof(int)
    // `topk_idx_buffers`: kNumChannels * kNumRanks * num_recv_buffer_tokens * num_topk * sizeof(int64_t)
    // `topk_weights_buffers`: kNumChannels * kNumRanks * num_recv_buffer_tokens * num_topk * sizeof(float)
    // `x_scales_buffers`: kNumChannels * kNumRanks * num_recv_buffer_tokens * num_scales * sizeof(float)
    auto channel_x_buffers = Buffer<int4>(ptr, num_channels_total * num_recv_buffer_tokens * hidden_int4, channel_rank_offset * num_recv_buffer_tokens * hidden_int4);
    auto channel_src_idx_buffers = Buffer<int>(ptr, num_channels_total * num_recv_buffer_tokens, channel_rank_offset * num_recv_buffer_tokens);
    auto channel_topk_idx_buffers = Buffer<int64_t>(ptr, num_channels_total * num_recv_buffer_tokens * num_topk, channel_rank_offset * num_recv_buffer_tokens * num_topk);
    auto channel_topk_weights_buffers = Buffer<float>(ptr, num_channels_total * num_recv_buffer_tokens * num_topk, channel_rank_offset * num_recv_buffer_tokens * num_topk);
    auto channel_x_scales_buffers = Buffer<float>(ptr, num_channels_total * num_recv_buffer_tokens * num_scales, channel_rank_offset * num_recv_buffer_tokens * num_scales);

225
    // TMA stuffs
226
#ifndef DISABLE_SM90_FEATURES
227
228
229
230
231
232
233
234
235
236
237
238
239
    extern __shared__ __align__(1024) uint8_t smem_buffer[];
    auto half_hidden_int4 = hidden_int4 / 2;
    auto half_hidden_bytes = half_hidden_int4 * static_cast<int>(sizeof(int4));
    auto tma_buffer = smem_buffer + (thread_id / 32) * kNumTMABytesPerWarp;
    auto tma_mbarrier = reinterpret_cast<uint64_t*>(tma_buffer + half_hidden_bytes);
    uint32_t tma_phase = 0;
    if (lane_id == 0) {
        mbarrier_init(tma_mbarrier, 1);
        fence_view_async_shared();
        fence_barrier_init();
        EP_DEVICE_ASSERT(hidden_int4 % 2 == 0 and half_hidden_bytes + sizeof(uint64_t) <= kNumTMABytesPerWarp);
    }
    __syncwarp();
240
#endif
241

Chenggang Zhao's avatar
Chenggang Zhao committed
242
243
    if (is_sender) {
        // Workers for sending
Chenggang Zhao's avatar
Chenggang Zhao committed
244
245
        constexpr int num_send_warps = kNumThreads / 32;
        constexpr int num_send_warps_per_rank = num_send_warps / kNumRanks;
Chenggang Zhao's avatar
Chenggang Zhao committed
246
        const auto send_thread_id = thread_id;
Chenggang Zhao's avatar
Chenggang Zhao committed
247
        const auto send_warp_id_in_rank = send_thread_id % num_threads_per_rank / 32;
Chenggang Zhao's avatar
Chenggang Zhao committed
248
        EP_DEVICE_ASSERT(kNumRanks <= 32);
Chenggang Zhao's avatar
Chenggang Zhao committed
249
        EP_DEVICE_ASSERT(num_send_warps % kNumRanks == 0);
Chenggang Zhao's avatar
Chenggang Zhao committed
250
251
252

        // Send offset by `-value - 1`, e.g. 0 -> -1, 1 -> -2
        // NOTES: this is for distinguishing zero tokens
253
        if (lane_id == 0 and send_warp_id_in_rank == 0) {
Chenggang Zhao's avatar
Chenggang Zhao committed
254
            int value = responsible_channel > 0 ? channel_prefix_matrix[responsible_rank * num_channels + responsible_channel - 1] : 0;
Chenggang Zhao's avatar
Chenggang Zhao committed
255
            st_relaxed_sys_global(channel_start_offset.buffer(), -value - 1);
Chenggang Zhao's avatar
Chenggang Zhao committed
256
            value = channel_prefix_matrix[responsible_rank * num_channels + responsible_channel];
Chenggang Zhao's avatar
Chenggang Zhao committed
257
258
259
260
261
262
263
264
265
266
            st_relaxed_sys_global(channel_end_offset.buffer(), -value - 1);
        }
        __syncwarp();

        // Get tasks
        int token_start_idx, token_end_idx;
        get_channel_task_range(num_tokens, num_channels, responsible_channel, token_start_idx, token_end_idx);

        // Iterate over all tokens and send by chunks
        int cached_channel_tail_idx = 0;
Chenggang Zhao's avatar
Chenggang Zhao committed
267
        for (int64_t token_idx = token_start_idx; token_idx < token_end_idx; ) {
Chenggang Zhao's avatar
Chenggang Zhao committed
268
            // Check destination queue emptiness, or wait a buffer to be released (rare cases)
Chenggang Zhao's avatar
Chenggang Zhao committed
269
            // NOTES: the head index received by different warps may not be the same
Chenggang Zhao's avatar
Chenggang Zhao committed
270
            auto start_time = clock64();
271
            while (lane_id == 0) {
Chenggang Zhao's avatar
Chenggang Zhao committed
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
                // NOTES: we only consider the worst case, because counting the real numbers are time-consuming
                int num_used_slots = cached_channel_tail_idx - ld_volatile_global(channel_head_idx.buffer());
                if (num_recv_buffer_tokens - num_used_slots >= num_max_send_tokens)
                    break;

                // Rare cases to loop again
                if (clock64() - start_time > NUM_TIMEOUT_CYCLES) {
                    printf("DeepEP timeout for dispatch senders, rank %d, responsible_channel = %d\n", rank, responsible_channel);
                    trap();
                }
            }
            __syncwarp();

            int chunk_token_idx = 0;
            while (chunk_token_idx < num_max_send_tokens and token_idx < token_end_idx) {
287
                // NOTES: for the same token, the warp assigned to save `send_head` may be different from the warp assigned to send the following data
288
                if (lane_id == 0 and token_idx % num_send_warps_per_rank == send_warp_id_in_rank)
Chenggang Zhao's avatar
Chenggang Zhao committed
289
290
                    send_head[token_idx * kNumRanks + responsible_rank] = is_token_in_rank[token_idx * kNumRanks + responsible_rank] ? cached_channel_tail_idx : -1;

Chenggang Zhao's avatar
Chenggang Zhao committed
291
                // Skip if not selected
Chenggang Zhao's avatar
Chenggang Zhao committed
292
                if (not is_token_in_rank[token_idx * kNumRanks + responsible_rank]) {
Chenggang Zhao's avatar
Chenggang Zhao committed
293
294
295
296
297
298
                    token_idx ++;
                    continue;
                }

                // Get an empty slot
                int dst_slot_idx = (cached_channel_tail_idx ++) % num_recv_buffer_tokens;
Chenggang Zhao's avatar
Chenggang Zhao committed
299
300
301
302
                if (cached_channel_tail_idx % num_send_warps_per_rank == send_warp_id_in_rank) {
                    // Copy data
                    auto shifted_channel_x_buffers = channel_x_buffers.buffer() + dst_slot_idx * hidden_int4;
                    auto shifted_x = x + token_idx * hidden_int4;
303
                    UNROLLED_WARP_COPY(5, lane_id, hidden_int4, shifted_channel_x_buffers, shifted_x, __ldg, st_na_global);
Chenggang Zhao's avatar
Chenggang Zhao committed
304
305

                    // Copy source index
306
                    if (lane_id == 0)
Chenggang Zhao's avatar
Chenggang Zhao committed
307
308
309
                        channel_src_idx_buffers[dst_slot_idx] = static_cast<int>(token_idx);

                    // Copy `topk_idx` and `topk_weights` with transformed index
310
                    if (lane_id < num_topk) {
Chenggang Zhao's avatar
Chenggang Zhao committed
311
312
                        // Top-k index
                        int recv_expert_begin = responsible_rank * num_experts_per_rank, recv_expert_end = (responsible_rank + 1) * num_experts_per_rank;
313
                        auto idx_value = __ldg(topk_idx + token_idx * num_topk + lane_id);
Chenggang Zhao's avatar
Chenggang Zhao committed
314
                        idx_value = (idx_value >= recv_expert_begin and idx_value < recv_expert_end) ? idx_value - recv_expert_begin : -1;
315
                        channel_topk_idx_buffers[dst_slot_idx * num_topk + lane_id] = idx_value;
Chenggang Zhao's avatar
Chenggang Zhao committed
316
317

                        // Top-k weights
318
                        auto weight_value = __ldg(topk_weights + token_idx * num_topk + lane_id);
Chenggang Zhao's avatar
Chenggang Zhao committed
319
                        weight_value = (idx_value >= 0) ? weight_value : 0.0f;
320
                        channel_topk_weights_buffers[dst_slot_idx * num_topk + lane_id] = weight_value;
Chenggang Zhao's avatar
Chenggang Zhao committed
321
322
323
324
                    }

                    // Copy `x_scales`
                    #pragma unroll
Shifang Xu's avatar
Shifang Xu committed
325
326
327
328
                    for (int i = lane_id; i < num_scales; i += 32) {
                        auto offset = token_idx * scale_token_stride + i * scale_hidden_stride;
                        channel_x_scales_buffers[dst_slot_idx * num_scales + i] = __ldg(x_scales + offset);
                    }
Chenggang Zhao's avatar
Chenggang Zhao committed
329
330
331
332
333
334
335
                }

                // Move token index
                chunk_token_idx ++, token_idx ++;
            }

            // Move tail index
Chenggang Zhao's avatar
Chenggang Zhao committed
336
337
            // NOTES: here all warps should share the same new tail
            asm volatile("bar.sync %0, %1;" :: "r"(responsible_rank), "r"(num_threads_per_rank));
338
            if (send_warp_id_in_rank == 0 and lane_id == 0)
Chenggang Zhao's avatar
Chenggang Zhao committed
339
340
341
342
                st_release_sys_global(channel_tail_idx.buffer(), cached_channel_tail_idx);
        }
    } else {
        // Workers for receiving and copying into buffer
Chenggang Zhao's avatar
Chenggang Zhao committed
343
344
        constexpr int num_recv_warps = kNumThreads / 32;
        constexpr int num_recv_warps_per_rank = num_recv_warps / kNumRanks;
Chenggang Zhao's avatar
Chenggang Zhao committed
345
        const auto recv_thread_id = thread_id;
Chenggang Zhao's avatar
Chenggang Zhao committed
346
347
348
349
        const auto recv_thread_id_in_rank = recv_thread_id % num_threads_per_rank;
        const auto recv_warp_id_in_rank = recv_thread_id_in_rank / 32;
        EP_DEVICE_ASSERT(kNumRanks <= 32);
        EP_DEVICE_ASSERT(recv_thread_id >= 0 and num_recv_warps % kNumRanks == 0);
Chenggang Zhao's avatar
Chenggang Zhao committed
350
351

        // Calculate offset first
352
        auto rank_prefix_matrix = static_cast<int*>(buffer_ptrs[rank]);
Chenggang Zhao's avatar
Chenggang Zhao committed
353
        int rank_offset = responsible_rank > 0 ? rank_prefix_matrix[(responsible_rank - 1) * kNumRanks + rank] : 0;
Chenggang Zhao's avatar
Chenggang Zhao committed
354
355
356

        // Receive channel offset
        int total_offset, num_tokens_to_recv;
357
358
359
        while (lane_id == 0 and (total_offset = ld_volatile_global(channel_start_offset.buffer())) == 0);
        while (lane_id == 0 and (num_tokens_to_recv = ld_volatile_global(channel_end_offset.buffer())) == 0);
        if (lane_id == 0) {
Chenggang Zhao's avatar
Chenggang Zhao committed
360
            total_offset = -total_offset - 1, num_tokens_to_recv = -num_tokens_to_recv - 1;
Chenggang Zhao's avatar
Chenggang Zhao committed
361
362
            if (recv_warp_id_in_rank == 0)
                recv_channel_offset[responsible_rank * num_channels + responsible_channel] = total_offset;
Chenggang Zhao's avatar
Chenggang Zhao committed
363
364
365
366
367
368
            num_tokens_to_recv -= total_offset;
        }
        total_offset = __shfl_sync(0xffffffff, total_offset, 0);
        total_offset += rank_offset;
        num_tokens_to_recv = __shfl_sync(0xffffffff, num_tokens_to_recv, 0);

Chenggang Zhao's avatar
Chenggang Zhao committed
369
370
371
        // Shared tail indices for different warps
        __shared__ volatile int shared_channel_tail_idx[kNumRanks];

Chenggang Zhao's avatar
Chenggang Zhao committed
372
373
374
        auto start_time = clock64();
        int cached_channel_head_idx = 0, cached_channel_tail_idx = 0;
        while (num_tokens_to_recv > 0) {
375
            // NOTES: unlike the sender, the receiver must ensure that the tail indices hold by different warps are the same
Chenggang Zhao's avatar
Chenggang Zhao committed
376
            while (recv_thread_id_in_rank == 0) {
fzyzcjy's avatar
fzyzcjy committed
377
                cached_channel_tail_idx = ld_acquire_sys_global(channel_tail_idx.buffer());
Chenggang Zhao's avatar
Chenggang Zhao committed
378
379

                // Ready to copy
Chenggang Zhao's avatar
Chenggang Zhao committed
380
381
                if (cached_channel_head_idx != cached_channel_tail_idx) {
                    shared_channel_tail_idx[responsible_rank] = cached_channel_tail_idx;
Chenggang Zhao's avatar
Chenggang Zhao committed
382
                    break;
Chenggang Zhao's avatar
Chenggang Zhao committed
383
                }
Chenggang Zhao's avatar
Chenggang Zhao committed
384
385
386
387
388
389
390
391

                // Timeout check
                if (clock64() - start_time > NUM_TIMEOUT_CYCLES) {
                    printf("DeepEP timeout for dispatch receivers, rank %d, responsible_channel = %d, tokens remained: %d\n", rank, responsible_channel, num_tokens_to_recv);
                    trap();
                }
            }

Chenggang Zhao's avatar
Chenggang Zhao committed
392
393
394
            // Synchronize queue tail
            asm volatile("bar.sync %0, %1;" :: "r"(responsible_rank), "r"(num_threads_per_rank));
            cached_channel_tail_idx = shared_channel_tail_idx[responsible_rank];
Chenggang Zhao's avatar
Chenggang Zhao committed
395
396
397

            // Copy data
            int num_recv_tokens = cached_channel_tail_idx - cached_channel_head_idx;
Chenggang Zhao's avatar
Chenggang Zhao committed
398
            for (int chunk_idx = recv_warp_id_in_rank; chunk_idx < num_recv_tokens; chunk_idx += num_recv_warps_per_rank) {
Chenggang Zhao's avatar
Chenggang Zhao committed
399
400
401
                int token_idx_in_buffer = (cached_channel_head_idx + chunk_idx) % num_recv_buffer_tokens;
                auto shifted_buffer_x_int4 = channel_x_buffers.buffer() + token_idx_in_buffer * hidden_int4;
                auto shifted_recv_x_int4 = recv_x + static_cast<int64_t>(total_offset + chunk_idx) * hidden_int4;
402
#ifndef DISABLE_SM90_FEATURES
403
404
405
406
407
408
409
410
411
                #pragma unroll
                for (int i = 0; i < 2; ++ i) if (lane_id == 0) {
                    tma_store_wait();
                    tma_load_1d(tma_buffer, shifted_buffer_x_int4 + i * half_hidden_int4, tma_mbarrier, half_hidden_bytes);
                    mbarrier_arrive_and_expect_tx(tma_mbarrier, half_hidden_bytes);
                    mbarrier_wait(tma_mbarrier, tma_phase);
                    tma_store_1d(tma_buffer, shifted_recv_x_int4 + i * half_hidden_int4, half_hidden_bytes, false);
                }
                __syncwarp();
412
413
414
415
#else
                UNROLLED_WARP_COPY(5, lane_id, hidden_int4, shifted_recv_x_int4, shifted_buffer_x_int4,
                                   ld_nc_global, st_na_global);
#endif
Chenggang Zhao's avatar
Chenggang Zhao committed
416
417
418
419
            }

            // Copy `src_idx`
            #pragma unroll 4
Chenggang Zhao's avatar
Chenggang Zhao committed
420
            for (int chunk_idx = cached_channel_head_idx + recv_thread_id_in_rank; chunk_idx < cached_channel_tail_idx; chunk_idx += 32 * num_recv_warps_per_rank)
Chenggang Zhao's avatar
Chenggang Zhao committed
421
422
423
424
                recv_src_idx[total_offset + chunk_idx - cached_channel_head_idx] = ld_nc_global(channel_src_idx_buffers.buffer() + chunk_idx % num_recv_buffer_tokens);

            // Copy `topk_idx` and `topk_weights`
            #pragma unroll 4
Chenggang Zhao's avatar
Chenggang Zhao committed
425
            for (int idx = recv_thread_id_in_rank; idx < num_recv_tokens * num_topk; idx += 32 * num_recv_warps_per_rank) {
Chenggang Zhao's avatar
Chenggang Zhao committed
426
427
428
429
430
431
432
433
434
435
                int chunk_idx = idx / num_topk, token_topk_idx = idx % num_topk;
                int token_idx_in_buffer = (cached_channel_head_idx + chunk_idx) % num_recv_buffer_tokens;
                auto recv_idx = static_cast<int64_t>(total_offset + chunk_idx) * num_topk + token_topk_idx;
                auto buffer_idx = token_idx_in_buffer * num_topk + token_topk_idx;
                recv_topk_idx[recv_idx] = ld_nc_global(channel_topk_idx_buffers.buffer() + buffer_idx);
                recv_topk_weights[recv_idx] = ld_nc_global(channel_topk_weights_buffers.buffer() + buffer_idx);
            }

            // Copy `x_scales`
            #pragma unroll 4
Chenggang Zhao's avatar
Chenggang Zhao committed
436
            for (int i = recv_thread_id_in_rank; i < num_recv_tokens * num_scales; i += 32 * num_recv_warps_per_rank) {
Chenggang Zhao's avatar
Chenggang Zhao committed
437
438
439
440
441
442
443
444
445
                int chunk_idx = i / num_scales, scales_idx = i % num_scales;
                int token_idx_in_buffer = (cached_channel_head_idx + chunk_idx) % num_recv_buffer_tokens;
                recv_x_scales[static_cast<int64_t>(total_offset + chunk_idx) * num_scales + scales_idx] =
                        ld_nc_global(channel_x_scales_buffers.buffer() + token_idx_in_buffer * num_scales + scales_idx);
            }

            // Move queue
            cached_channel_head_idx += num_recv_tokens;
            total_offset += num_recv_tokens;
Chenggang Zhao's avatar
Chenggang Zhao committed
446
            asm volatile("bar.sync %0, %1;" :: "r"(responsible_rank), "r"(num_threads_per_rank));
447
            if (recv_warp_id_in_rank == num_recv_warps_per_rank - 1 and lane_id == 0)
Chenggang Zhao's avatar
Chenggang Zhao committed
448
449
450
451
452
                st_relaxed_sys_global(channel_head_idx.buffer(), cached_channel_head_idx);

            // Exit
            num_tokens_to_recv -= num_recv_tokens;
        }
453
454

        // Make TMA store visible to the next kernel
455
#ifndef DISABLE_SM90_FEATURES
456
457
        if (lane_id == 0)
            tma_store_wait();
458
#endif
Chenggang Zhao's avatar
Chenggang Zhao committed
459
    }
460
461
462
463
464
465
466
467
468
469
470
471
472


    // Clean unused `recv_topk_idx` as -1
    if (num_worst_tokens > 0) {
        auto rank_prefix_matrix = static_cast<int*>(buffer_ptrs[rank]);
        const auto num_recv_tokens = rank_prefix_matrix[(kNumRanks - 1) * kNumRanks + rank];
        const auto clean_start = num_recv_tokens * num_topk + sm_id * kNumThreads;
        const auto clean_end = num_worst_tokens * num_topk;
        const auto clean_stride = num_sms * kNumThreads;
        #pragma unroll
        for (int i = clean_start + thread_id; i < clean_end; i += clean_stride)
            recv_topk_idx[i] = -1;
    }
Chenggang Zhao's avatar
Chenggang Zhao committed
473
474
475
476
477
}

void dispatch(void* recv_x, float* recv_x_scales, int* recv_src_idx, int64_t* recv_topk_idx, float* recv_topk_weights, int* recv_channel_offset,
              int* send_head, const void* x, const float* x_scales, const int64_t* topk_idx, const float* topk_weights,
              const bool* is_token_in_rank, const int* channel_prefix_matrix,
478
              int num_tokens, int num_worst_tokens, int hidden_int4, int num_topk, int num_experts, int num_scales,
Shifang Xu's avatar
Shifang Xu committed
479
              int scale_token_stride, int scale_hidden_stride,
Chenggang Zhao's avatar
Chenggang Zhao committed
480
481
              void** buffer_ptrs, int rank, int num_ranks,
              cudaStream_t stream, int num_sms, int num_max_send_tokens, int num_recv_buffer_tokens) {
482
483
    constexpr int kNumThreads = 768;
    constexpr int kNumTMABytesPerWarp = 8192;
484
#ifndef DISABLE_SM90_FEATURES
485
    constexpr int smem_size = kNumTMABytesPerWarp * (kNumThreads / 32);
486
#endif
487

Shifang Xu's avatar
Shifang Xu committed
488
489
490
    // Make sure never OOB
    EP_HOST_ASSERT(static_cast<int64_t>(num_scales) * scale_hidden_stride < std::numeric_limits<int>::max());

491
492
#define DISPATCH_LAUNCH_CASE(ranks) { \
    auto kernel = dispatch<ranks, kNumThreads, kNumTMABytesPerWarp>; \
493
    SET_SHARED_MEMORY_FOR_TMA(kernel); \
494
495
496
497
    LAUNCH_KERNEL(&cfg, kernel, \
        reinterpret_cast<int4*>(recv_x), recv_x_scales, recv_src_idx, recv_topk_idx, recv_topk_weights, recv_channel_offset, \
        send_head, reinterpret_cast<const int4*>(x), x_scales, topk_idx, topk_weights, \
        is_token_in_rank, channel_prefix_matrix, \
498
        num_tokens, num_worst_tokens, hidden_int4, num_topk, num_experts, num_scales, \
Shifang Xu's avatar
Shifang Xu committed
499
        scale_token_stride, scale_hidden_stride, \
500
501
502
        buffer_ptrs, rank, \
        num_max_send_tokens, num_recv_buffer_tokens); \
    } break
Chenggang Zhao's avatar
Chenggang Zhao committed
503
504
505

    // Even-numbered blocks for sending, odd-numbered blocks for receiving.
    EP_HOST_ASSERT(num_sms % 2 == 0);
Chenggang Zhao's avatar
Chenggang Zhao committed
506
    SETUP_LAUNCH_CONFIG(num_sms, kNumThreads, stream);
Chenggang Zhao's avatar
Chenggang Zhao committed
507
508
509
510
511
512
513
    SWITCH_RANKS(DISPATCH_LAUNCH_CASE);
#undef DISPATCH_LAUNCH_CASE
}

template<int kNumRanks>
__global__ void
cached_notify_combine(void** buffer_ptrs, int* send_head, int num_channels, int num_recv_tokens, int num_memset_int,
514
                      int** barrier_signal_ptrs, int rank) {
Chenggang Zhao's avatar
Chenggang Zhao committed
515
516
517
    const auto sm_id = static_cast<int>(blockIdx.x);
    if (sm_id == 0) {
        // Barrier before cleaning
518
        barrier_block<kNumRanks, true>(barrier_signal_ptrs, rank);
Chenggang Zhao's avatar
Chenggang Zhao committed
519
520
521

        // Clean
        auto thread_id = static_cast<int>(threadIdx.x), num_threads = static_cast<int>(blockDim.x);
522
        auto ptr = static_cast<int*>(buffer_ptrs[rank]);
Chenggang Zhao's avatar
Chenggang Zhao committed
523
524
525
526
527
        #pragma unroll
        for (int i = thread_id; i < num_memset_int; i += num_threads)
            ptr[i] = 0;

        // Barrier after cleaning
528
        barrier_block<kNumRanks>(barrier_signal_ptrs, rank);
Chenggang Zhao's avatar
Chenggang Zhao committed
529
530
531
532
533
    } else {
        const auto channel_id = sm_id - 1;
        const auto thread_id = static_cast<int>(threadIdx.x);
        const auto rank_id = thread_id / 32;
        const auto lane_id = thread_id % 32;
534
535
        if (rank_id >= kNumRanks)
            return;
Chenggang Zhao's avatar
Chenggang Zhao committed
536
537
538
539
540
541
542
543
544
545
546

        int token_start_idx, token_end_idx;
        get_channel_task_range(num_recv_tokens, num_channels, channel_id, token_start_idx, token_end_idx);

        // NOTES: `1 << 25` is a heuristic large number
        int last_head = 1 << 25;
        #pragma unroll
        for (int token_idx_tail = token_end_idx - 1; token_idx_tail >= token_start_idx; token_idx_tail -= 32) {
            int token_idx = token_idx_tail - lane_id, expected_head = 0;
            auto current_head = (token_idx >= token_start_idx) ? __ldg(send_head + token_idx * kNumRanks + rank_id) : -1;
            for (int i = 0; i < min(32, token_idx_tail - token_start_idx + 1); ++ i) {
547
                const int head = __shfl_sync(0xffffffff, current_head, i);
Chenggang Zhao's avatar
Chenggang Zhao committed
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
                if (head < 0) {
                    if (lane_id == i)
                        expected_head = -last_head - 1;
                } else {
                    last_head = head;
                }
            }
            if (current_head < 0 and token_idx >= token_start_idx)
                send_head[token_idx * kNumRanks + rank_id] = expected_head;
        }
    }
}

void cached_notify_combine(void** buffer_ptrs, int* send_head, int num_channels,
                           int num_recv_tokens, int num_memset_int,
563
                           int** barrier_signal_ptrs, int rank, int num_ranks,
Chenggang Zhao's avatar
Chenggang Zhao committed
564
565
566
                           cudaStream_t stream) {
#define CACHED_NOTIFY_COMBINE(ranks) \
    LAUNCH_KERNEL(&cfg, cached_notify_combine<ranks>, \
567
        buffer_ptrs, send_head, num_channels, num_recv_tokens, num_memset_int, barrier_signal_ptrs, rank); \
Chenggang Zhao's avatar
Chenggang Zhao committed
568
569
570
571
572
573
574
575
576
577
578
    break

    const int num_threads = std::max(128, 32 * num_ranks);
    EP_HOST_ASSERT(num_ranks <= num_threads);
    EP_HOST_ASSERT(num_threads <= 1024);
    EP_HOST_ASSERT(1 + num_channels <= num_channels * 2);
    SETUP_LAUNCH_CONFIG(1 + num_channels, num_threads, stream);
    SWITCH_RANKS(CACHED_NOTIFY_COMBINE);
#undef CACHED_NOTIFY_COMBINE
}

579
template<typename dtype_t, int kNumRanks, int kNumThreads, int kNumTMABytesPerWarp>
Chenggang Zhao's avatar
Chenggang Zhao committed
580
581
582
__global__ void __launch_bounds__(kNumThreads, 1)
combine(dtype_t* recv_x, float* recv_topk_weights,
        const dtype_t* x, const float* topk_weights,
Shangyan Zhou's avatar
Shangyan Zhou committed
583
        const dtype_t* bias_0, const dtype_t* bias_1,
Chenggang Zhao's avatar
Chenggang Zhao committed
584
585
        const int* src_idx, const int* rank_prefix_matrix, const int* channel_prefix_matrix,
        int* send_head, int num_tokens, int num_recv_tokens, int hidden, int num_topk,
586
        void** buffer_ptrs, int rank,
Chenggang Zhao's avatar
Chenggang Zhao committed
587
588
589
        int num_max_send_tokens, int num_recv_buffer_tokens) {
    const auto num_sms = static_cast<int>(gridDim.x);
    const auto thread_id = static_cast<int>(threadIdx.x);
590
    const auto sm_id = static_cast<int>(blockIdx.x), lane_id = get_lane_id();
Chenggang Zhao's avatar
Chenggang Zhao committed
591
592
593
594
595
596
597
598
    const auto num_channels = num_sms / 2;
    const bool is_sender = sm_id % 2 == 0;
    const int responsible_channel = sm_id / 2;
    EP_DEVICE_ASSERT(num_topk <= 32);

    constexpr int kDtypePerInt4 = sizeof(int4) / sizeof(dtype_t);
    int hidden_int4 = hidden * sizeof(dtype_t) / sizeof(int4);
    auto x_int4 = reinterpret_cast<const int4*>(x);
Shangyan Zhou's avatar
Shangyan Zhou committed
599
600
    auto bias_0_int4 = reinterpret_cast<const int4*>(bias_0);
    auto bias_1_int4 = reinterpret_cast<const int4*>(bias_1);
Chenggang Zhao's avatar
Chenggang Zhao committed
601
602
    auto recv_int4 = reinterpret_cast<int4*>(recv_x);

603
    // TMA stuffs
604
#ifndef DISABLE_SM90_FEATURES
605
606
    extern __shared__ __align__(1024) uint8_t smem_buffer[];
    auto tma_buffer = smem_buffer + (thread_id / 32) * kNumTMABytesPerWarp;
607
#endif
608

Chenggang Zhao's avatar
Chenggang Zhao committed
609
610
611
    if (is_sender) {
        // Workers for sending
        // Several warps are responsible for a single rank
612
613
        constexpr int num_send_warps_per_rank = (kNumThreads / 32) / kNumRanks;
        constexpr int num_send_warps = num_send_warps_per_rank * kNumRanks;
Chenggang Zhao's avatar
Chenggang Zhao committed
614
615
        const auto num_threads_per_rank = num_send_warps_per_rank * 32;
        const auto send_thread_id = thread_id;
616
        const auto send_warp_id = send_thread_id / 32;
617
618
        const auto send_rank_id = (responsible_channel + send_warp_id) % kNumRanks;
        const auto send_warp_id_in_rank = send_warp_id / kNumRanks;
619
        EP_STATIC_ASSERT(num_send_warps * 32 == kNumThreads, "Invalid warp count");
Chenggang Zhao's avatar
Chenggang Zhao committed
620
621

        // Calculate pointers by the specific layout
622
        auto ptr = reinterpret_cast<void*>(static_cast<int8_t*>(buffer_ptrs[send_rank_id]));
Chenggang Zhao's avatar
Chenggang Zhao committed
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
        auto num_channels_total = num_channels * kNumRanks;
        auto channel_rank_offset = responsible_channel * kNumRanks + rank;

        // Channel meta data
        // `head_idx`: kNumChannels * kNumRanks * sizeof(int)
        // `tail_idx`: kNumChannels * kNumRanks * sizeof(int)
        // `x_buffers`: kNumChannels * kNumRanks * num_recv_buffer_tokens * hidden_int4 * sizeof(int4)
        // `src_idx_buffers`: kNumChannels * kNumRanks * num_recv_buffer_tokens * sizeof(int)
        // `topk_weights_buffers`: kNumChannels * kNumRanks * num_recv_buffer_tokens * num_topk * sizeof(float)
        auto channel_head_idx = Buffer<int>(ptr, num_channels_total, channel_rank_offset);
        auto channel_tail_idx = Buffer<int>(ptr, num_channels_total, channel_rank_offset);
        auto channel_x_buffers = Buffer<int4>(ptr, num_channels_total * num_recv_buffer_tokens * hidden_int4, channel_rank_offset * num_recv_buffer_tokens * hidden_int4);
        auto channel_src_idx_buffers = Buffer<int>(ptr, num_channels_total * num_recv_buffer_tokens, channel_rank_offset * num_recv_buffer_tokens);
        auto channel_topk_weights_buffers = Buffer<float>(ptr, num_channels_total * num_recv_buffer_tokens * num_topk, channel_rank_offset * num_recv_buffer_tokens * num_topk);

        // Get tasks
        // NOTES: `channel_offset` is already shifted
        int rank_offset = send_rank_id > 0 ? rank_prefix_matrix[(send_rank_id - 1) * kNumRanks + rank] : 0;
        int num_rank_tokens = rank_prefix_matrix[send_rank_id * kNumRanks + rank] - rank_offset;
        int channel_offset = channel_prefix_matrix[send_rank_id * num_channels + responsible_channel];
        int num_channel_tokens = (responsible_channel == num_channels - 1 ? num_rank_tokens : channel_prefix_matrix[send_rank_id * num_channels + responsible_channel + 1]) - channel_offset;
        int token_start_idx = rank_offset + channel_offset, token_end_idx = rank_offset + channel_offset + num_channel_tokens;

        // Iterate over all tokens and send by chunks
        int current_channel_tail_idx = 0;
        for (int64_t token_idx = token_start_idx; token_idx < token_end_idx; ) {
            // Check destination queue emptiness, or wait a buffer to be released (rare cases)
            auto start_time = clock64();
            int num_round_tokens = min(num_max_send_tokens, token_end_idx - static_cast<int>(token_idx));
652
            while (lane_id == 0) {
Chenggang Zhao's avatar
Chenggang Zhao committed
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
                // NOTES: we only consider the worst case, because counting the real numbers are time-consuming
                int num_used_slots = current_channel_tail_idx - ld_volatile_global(channel_head_idx.buffer());
                if (num_recv_buffer_tokens - num_used_slots >= num_round_tokens)
                    break;

                // Rare cases to loop again
                if (clock64() - start_time > NUM_TIMEOUT_CYCLES) {
                    printf("DeepEP timeout for combine senders, rank %d, responsible_channel = %d\n", rank, responsible_channel);
                    trap();
                }
            }
            __syncwarp();

            // Send by chunk
            #pragma unroll
            for (int i = send_warp_id_in_rank; i < num_round_tokens; i += num_send_warps_per_rank) {
                // Get an empty slot
                int dst_slot_idx = (current_channel_tail_idx + i) % num_recv_buffer_tokens;

                // Copy data
                auto shifted_x_buffers = channel_x_buffers.buffer() + dst_slot_idx * hidden_int4;
                auto shifted_x = x_int4 + (token_idx + i) * hidden_int4;
675
                UNROLLED_WARP_COPY(4, lane_id, hidden_int4, shifted_x_buffers, shifted_x, ld_nc_global, st_na_global);
Chenggang Zhao's avatar
Chenggang Zhao committed
676
677

                // Send source index
678
                if (lane_id == 0)
Chenggang Zhao's avatar
Chenggang Zhao committed
679
680
681
                    channel_src_idx_buffers[dst_slot_idx] = __ldg(src_idx + token_idx + i);

                // Send `topk_weights`
682
683
                if (num_topk > 0 and lane_id < num_topk)
                    channel_topk_weights_buffers[dst_slot_idx * num_topk + lane_id] = __ldg(topk_weights + (token_idx + i) * num_topk + lane_id);
Chenggang Zhao's avatar
Chenggang Zhao committed
684
685
686
687
688
689
            }
            token_idx += num_round_tokens;
            current_channel_tail_idx += num_round_tokens;

            // Move tail index
            asm volatile("bar.sync %0, %1;" :: "r"(send_rank_id), "r"(num_threads_per_rank));
690
            if (lane_id == 0 and send_warp_id_in_rank == 0)
Chenggang Zhao's avatar
Chenggang Zhao committed
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
                st_release_sys_global(channel_tail_idx.buffer(), current_channel_tail_idx);
        }
    } else {
        // Workers for receiving
        // One warp for moving the queue head, others for reduction
        constexpr int num_recv_warps = kNumThreads / 32;
        const auto recv_warp_id = thread_id / 32;
        EP_DEVICE_ASSERT(kNumRanks <= 32 and kNumThreads > 32);
        EP_DEVICE_ASSERT(thread_id >= 0 and kNumThreads % 32 == 0);

        // Shared head, tail and retired flags for receiver warps
        __shared__ volatile int warp_channel_head_idx[num_recv_warps][kNumRanks];
        __shared__ volatile int channel_tail_idx[kNumRanks];
        __shared__ volatile bool warp_retired[num_recv_warps];
        if (thread_id < num_recv_warps)
            warp_retired[thread_id] = false;
707
708
        if (lane_id < kNumRanks)
            warp_channel_head_idx[recv_warp_id][lane_id] = 0;
Chenggang Zhao's avatar
Chenggang Zhao committed
709
710
711
712
713
        if (thread_id < kNumRanks)
            channel_tail_idx[thread_id] = 0;
        asm volatile("bar.sync 0, %0;" :: "r"(kNumThreads));

        if (thread_id < 32) {
714
            int* channel_head_idx_ptr = static_cast<int*>(buffer_ptrs[rank]) + responsible_channel * kNumRanks + lane_id;
Chenggang Zhao's avatar
Chenggang Zhao committed
715
716
717
718
            int* channel_tail_idx_ptr = channel_head_idx_ptr + num_channels * kNumRanks;

            // Queue head updater
            int last_head = 0;
719
            while (lane_id < kNumRanks) {
Chenggang Zhao's avatar
Chenggang Zhao committed
720
721
722
723
724
725
726
727
728
                // Check retired
                bool retired = true;
                #pragma unroll
                for (int i = 1; i < num_recv_warps; ++ i)
                    retired = retired and warp_retired[i];
                if (retired)
                    break;

                // Update queue tail
729
                channel_tail_idx[lane_id] = ld_acquire_sys_global(channel_tail_idx_ptr);
Chenggang Zhao's avatar
Chenggang Zhao committed
730
731
732
733
734

                // Update minimum head
                int min_head = std::numeric_limits<int>::max();
                #pragma unroll
                for (int i = 1; i < num_recv_warps; ++ i) if (not warp_retired[i])
735
                    min_head = min(min_head, warp_channel_head_idx[i][lane_id]);
Chenggang Zhao's avatar
Chenggang Zhao committed
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
                if (min_head != std::numeric_limits<int>::max() and min_head > last_head)
                    st_relaxed_sys_global(channel_head_idx_ptr, last_head = min_head);
            }
        } else {
            // Receivers
            // Channel metadata
            // All lanes will use data buffer, but only rank lane will use `head/tail/src_idx`
            Buffer<int4> channel_x_buffers[kNumRanks];
            Buffer<float> channel_topk_weights_buffers[kNumRanks];

            // Calculate pointers by the specific layout
            #pragma unroll
            for (int i = 0; i < kNumRanks; ++ i) {
                auto channel_rank_offset = responsible_channel * kNumRanks + i;
                auto num_channels_total = num_channels * kNumRanks;
                // `head_idx` & `tail_idx`: kNumChannels * kNumRanks * sizeof(int)
752
                auto ptr = reinterpret_cast<void*>(static_cast<int8_t*>(buffer_ptrs[rank]) + 2 * num_channels * kNumRanks * sizeof(int));
Chenggang Zhao's avatar
Chenggang Zhao committed
753
754
755
756
757

                // `x_buffers`: kNumChannels * kNumRanks * num_recv_buffer_tokens * hidden_int4 * sizeof(int4)
                channel_x_buffers[i] = Buffer<int4>(ptr, num_channels_total * num_recv_buffer_tokens * hidden_int4, channel_rank_offset * num_recv_buffer_tokens * hidden_int4);

                // `src_idx_buffers`: kNumChannels * kNumRanks * num_recv_buffer_tokens * sizeof(int)
758
                ptr = reinterpret_cast<void*>(static_cast<int8_t*>(ptr) + num_channels_total * num_recv_buffer_tokens * sizeof(int));
Chenggang Zhao's avatar
Chenggang Zhao committed
759
760
761
762
763
764
765
766
767
768
769
770
771

                // `topk_weights_buffers`: kNumChannels * kNumRanks * num_recv_buffer_tokens * num_topk * sizeof(float)
                channel_topk_weights_buffers[i] = Buffer<float>(ptr, num_channels_total * num_recv_buffer_tokens * num_topk, channel_rank_offset * num_recv_buffer_tokens * num_topk);
            }

            // The same tokens as the dispatch process
            int token_start_idx, token_end_idx;
            get_channel_task_range(num_recv_tokens, num_channels, responsible_channel, token_start_idx, token_end_idx);

            // Iterate over all tokens and combine
            for (int64_t token_idx = token_start_idx + recv_warp_id - 1; token_idx < token_end_idx; token_idx += num_recv_warps - 1) {
                // Read expected head
                int expected_head = -1;
772
773
                if (lane_id < kNumRanks)
                    expected_head = ld_nc_global(send_head + token_idx * kNumRanks + lane_id);
774

Chenggang Zhao's avatar
Chenggang Zhao committed
775
                auto start_time = clock64();
776
                while (__any_sync(0xffffffff, channel_tail_idx[lane_id] <= expected_head and expected_head >= 0)) {
Chenggang Zhao's avatar
Chenggang Zhao committed
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
                    // Timeout check
                    if (clock64() - start_time > NUM_TIMEOUT_CYCLES) {
                        printf("DeepEP timeout for combine receivers, rank %d, responsible_channel = %d, expect = %d\n", rank, responsible_channel, expected_head);
                        trap();
                    }
                }
                __syncwarp();

                // Broadcast current heads
                int num_topk_ranks = 0, topk_ranks[kNumRanks], slot_indices[kNumRanks];
                #pragma unroll
                for (int i = 0; i < kNumRanks; ++ i) {
                    auto expected_head_i = __shfl_sync(0xffffffff, expected_head, i);
                    if (expected_head_i >= 0) {
                        slot_indices[num_topk_ranks] = expected_head_i % num_recv_buffer_tokens;
                        topk_ranks[num_topk_ranks ++] = i;
                    }
                }

796
                // Wait shared memory release
797
#ifndef DISABLE_SM90_FEATURES
798
799
800
                if (lane_id == 0)
                    tma_store_wait();
                __syncwarp();
801
#endif
802
803
804
805

                // Reduce data with pipeline
                constexpr int kNumStages = 8;
                EP_STATIC_ASSERT(kNumStages * 32 * sizeof(int4) <= kNumTMABytesPerWarp, "Invalid count");
Chenggang Zhao's avatar
Chenggang Zhao committed
806
                #pragma unroll
807
                for (int i = lane_id; i < hidden_int4; i += 32) {
Shangyan Zhou's avatar
Shangyan Zhou committed
808
809
810
811
812
                    // Read bias
                    // TODO: make it as a template
                    int4 bias_0_value_int4 = bias_0_int4 != nullptr ? __ldg(bias_0_int4 + token_idx * hidden_int4 + i) : make_int4(0, 0, 0, 0);
                    int4 bias_1_value_int4 = bias_1_int4 != nullptr ? __ldg(bias_1_int4 + token_idx * hidden_int4 + i) : make_int4(0, 0, 0, 0);

Chenggang Zhao's avatar
Chenggang Zhao committed
813
814
815
816
817
818
                    // Read buffers
                    int4 recv_value_int4[kNumRanks];
                    #pragma unroll
                    for (int j = 0; j < num_topk_ranks; ++ j)
                        recv_value_int4[j] = ld_nc_global(channel_x_buffers[topk_ranks[j]].buffer() + slot_indices[j] * hidden_int4 + i);

Shangyan Zhou's avatar
Shangyan Zhou committed
819
820
821
822
823
824
825
826
                    // Reduce bias
                    float values[kDtypePerInt4];
                    auto bias_0_values = reinterpret_cast<const dtype_t*>(&bias_0_value_int4);
                    auto bias_1_values = reinterpret_cast<const dtype_t*>(&bias_1_value_int4);
                    #pragma unroll
                    for (int j = 0; j < kDtypePerInt4; ++ j)
                        values[j] = static_cast<float>(bias_0_values[j]) + static_cast<float>(bias_1_values[j]);

Chenggang Zhao's avatar
Chenggang Zhao committed
827
828
829
830
831
832
833
834
835
                    // Reduce all-to-all results
                    #pragma unroll
                    for (int j = 0; j < num_topk_ranks; ++ j) {
                        auto recv_value_dtypes = reinterpret_cast<const dtype_t*>(&recv_value_int4[j]);
                        #pragma unroll
                        for (int k = 0; k < kDtypePerInt4; ++ k)
                            values[k] += static_cast<float>(recv_value_dtypes[k]);
                    }

836
                    // Cast back to `dtype_t`
Chenggang Zhao's avatar
Chenggang Zhao committed
837
838
839
840
841
                    int4 out_int4;
                    auto out_dtypes = reinterpret_cast<dtype_t*>(&out_int4);
                    #pragma unroll
                    for (int j = 0; j < kDtypePerInt4; ++ j)
                        out_dtypes[j] = static_cast<dtype_t>(values[j]);
842

843
#ifndef DISABLE_SM90_FEATURES
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
                    // Wait TMA arrival
                    if (lane_id == 0)
                        tma_store_wait<kNumStages - 1>();
                    __syncwarp();

                    // Write into TMA buffer
                    auto tma_stage_idx = (i / 32) % kNumStages;
                    reinterpret_cast<int4*>(tma_buffer)[tma_stage_idx * 32 + lane_id] = out_int4;

                    // Issue TMA
                    tma_store_fence();
                    __syncwarp();
                    if (lane_id == 0) {
                        auto tma_bytes = min(32, hidden_int4 - i) * static_cast<int>(sizeof(int4));
                        tma_store_1d(reinterpret_cast<int4*>(tma_buffer) + tma_stage_idx * 32,
                                     recv_int4 + token_idx * hidden_int4 + i, tma_bytes, false);
                    }
                    __syncwarp();
862
863
864
#else
                    recv_int4[token_idx * hidden_int4 + i] = out_int4;
#endif
Chenggang Zhao's avatar
Chenggang Zhao committed
865
866
867
                }

                // Reduce `topk_weights`
868
                if (lane_id < num_topk) {
Chenggang Zhao's avatar
Chenggang Zhao committed
869
870
871
                    float value = 0;
                    #pragma unroll
                    for (int i = 0; i < num_topk_ranks; ++ i)
872
873
                        value += ld_nc_global(channel_topk_weights_buffers[topk_ranks[i]].buffer() + slot_indices[i] * num_topk + lane_id);
                    recv_topk_weights[token_idx * num_topk + lane_id] = value;
Chenggang Zhao's avatar
Chenggang Zhao committed
874
                }
875
876

                // Update head
877
878
                if (lane_id < kNumRanks)
                    warp_channel_head_idx[recv_warp_id][lane_id] = (expected_head < 0) ? -expected_head - 1 : expected_head + 1;
Chenggang Zhao's avatar
Chenggang Zhao committed
879
880
881
882
            }

            // Retired
            __syncwarp();
883
            if (lane_id == 0)
Chenggang Zhao's avatar
Chenggang Zhao committed
884
                warp_retired[recv_warp_id] = true;
885
886

            // Make TMA store visible to the next kernel
887
#ifndef DISABLE_SM90_FEATURES
888
889
            if (lane_id == 0)
                tma_store_wait();
890
#endif
Chenggang Zhao's avatar
Chenggang Zhao committed
891
892
893
894
895
896
897
        }
    }
}

void combine(cudaDataType_t type,
             void* recv_x, float* recv_topk_weights,
             const void* x, const float* topk_weights,
Shangyan Zhou's avatar
Shangyan Zhou committed
898
             const void* bias_0, const void* bias_1,
Chenggang Zhao's avatar
Chenggang Zhao committed
899
900
901
902
903
904
             const int* src_idx, const int* rank_prefix_matrix, const int* channel_prefix_matrix,
             int* send_head, int num_tokens, int num_recv_tokens, int hidden, int num_topk,
             void** buffer_ptrs, int rank, int num_ranks,
             cudaStream_t stream, int num_sms,
             int num_max_send_tokens, int num_recv_buffer_tokens) {
    constexpr int kNumThreads = 768;
905
    constexpr int kNumTMABytesPerWarp = 4096;
906
#ifndef DISABLE_SM90_FEATURES
907
    constexpr int smem_size = kNumTMABytesPerWarp * (kNumThreads / 32);
908
#endif
909
910
911

#define COMBINE_LAUNCH_CASE(dtype, ranks) { \
    auto kernel = combine<dtype, ranks, kNumThreads, kNumTMABytesPerWarp>; \
912
    SET_SHARED_MEMORY_FOR_TMA(kernel); \
913
    LAUNCH_KERNEL(&cfg, kernel, \
Chenggang Zhao's avatar
Chenggang Zhao committed
914
915
        reinterpret_cast<dtype*>(recv_x), recv_topk_weights, \
        reinterpret_cast<const dtype*>(x), topk_weights,   \
Shangyan Zhou's avatar
Shangyan Zhou committed
916
        reinterpret_cast<const dtype*>(bias_0), reinterpret_cast<const dtype*>(bias_1), \
Chenggang Zhao's avatar
Chenggang Zhao committed
917
918
919
        src_idx, rank_prefix_matrix, channel_prefix_matrix, \
        send_head, num_tokens, num_recv_tokens, hidden, num_topk, \
        buffer_ptrs, rank, \
920
        num_max_send_tokens, num_recv_buffer_tokens); } \
Chenggang Zhao's avatar
Chenggang Zhao committed
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
    break
#define COMBINE_DTYPE_LAUNCH_CASE(dtype) SWITCH_RANKS_WITH_DTYPE(dtype, COMBINE_LAUNCH_CASE); break

    // Even-numbered blocks for sending, odd-numbered blocks for receiving
    EP_HOST_ASSERT(num_sms % 2 == 0);
    EP_HOST_ASSERT(kNumThreads >= num_ranks * 32);
    SETUP_LAUNCH_CONFIG(num_sms, kNumThreads, stream);
    SWITCH_TYPES(COMBINE_DTYPE_LAUNCH_CASE);
#undef COMBINE_DTYPE_LAUNCH_CASE
#undef COMBINE_LAUNCH_CASE
}

} // namespace intranode

} // namespace deep_ep