internode_ll.cu 55.5 KB
Newer Older
Chenggang Zhao's avatar
Chenggang Zhao committed
1
2
3
#include "configs.cuh"
#include "exception.cuh"
#include "launch.cuh"
4
5
6
7
#include "buffer.cuh"
#include "utils.cuh"
// #include <cooperative_groups.h>
#include <iostream>
lishen's avatar
lishen committed
8
9
10

#include "hip/hip_runtime.h"

11
#include "shmem_wrapper.cuh"
12
#include "internode_ll_logfmt.cuh"
Chenggang Zhao's avatar
Chenggang Zhao committed
13
14
15
16
17

namespace deep_ep {

namespace internode_ll {

lishen's avatar
lishen committed
18
19
20
21
22
23
24
template <typename dtype_a_t, typename dtype_b_t>
__device__ __forceinline__ dtype_b_t pack2(const dtype_a_t& x, const dtype_a_t& y) {
    EP_STATIC_ASSERT(sizeof(dtype_a_t) * 2 == sizeof(dtype_b_t), "Invalid dtypes");
    dtype_b_t packed;
    auto unpacked_ptr = reinterpret_cast<dtype_a_t*>(&packed);
    unpacked_ptr[0] = x, unpacked_ptr[1] = y;
    return packed;
25
26
27
28
29
}

__device__ void grid_barrier(int* global_counter, int num_blocks) {
    volatile int ret;
    __syncthreads();
lishen's avatar
lishen committed
30
    __threadfence();
31
    if (threadIdx.x == 0 ) {
lishen's avatar
lishen committed
32
        ret = __hip_atomic_fetch_add(&global_counter[0], 1, __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_AGENT);
33
34
35
    }
    __syncthreads();
    if (threadIdx.x == 0) {
lishen's avatar
lishen committed
36
        while (__hip_atomic_load(global_counter, __ATOMIC_RELAXED,__HIP_MEMORY_SCOPE_AGENT) != num_blocks);
37
38
39
    }
    __syncthreads();
}
lishen's avatar
lishen committed
40
41
42
43
44
45
46
47
48
49
50
template <typename dtype_t>
__host__ __device__ dtype_t ceil_div(dtype_t a, dtype_t b) {
    return (a + b - 1) / b;
}

template <typename dtype_a_t, typename dtype_b_t>
__device__ __forceinline__ void unpack2(const dtype_b_t& packed, dtype_a_t& x, dtype_a_t& y) {
    EP_STATIC_ASSERT(sizeof(dtype_a_t) * 2 == sizeof(dtype_b_t), "Invalid dtypes");
    auto unpacked_ptr = reinterpret_cast<const dtype_a_t*>(&packed);
    x = unpacked_ptr[0], y = unpacked_ptr[1];
}
51
52


Chenggang Zhao's avatar
Chenggang Zhao committed
53
template <int kNumThreads> __launch_bounds__(kNumThreads, 1)
54
__global__ void clean_low_latency_buffer(int64_t* clean_0, int num_clean_int_0,
lishen's avatar
lishen committed
55
                                         int64_t* clean_1, int num_clean_int_1) {
Chenggang Zhao's avatar
Chenggang Zhao committed
56
    // Barrier before cleaning (in case of unfinished chunked EP)
lishen's avatar
lishen committed
57
    if (threadIdx.x == 0)
58
        internode::shmem_device_barrier_all();
59
60

    // Clean
lishen's avatar
lishen committed
61
    auto thread_id = static_cast<int>(threadIdx.x);
62
63
64
65
66
67
68
    #pragma unroll
    for (int i = thread_id; i < num_clean_int_0; i += kNumThreads)
        clean_0[i] = 0;
    #pragma unroll
    for (int i = thread_id; i < num_clean_int_1; i += kNumThreads)
        clean_1[i] = 0;

lishen's avatar
lishen committed
69
    // Barrier after cleaning (make sure low-latency mode work
lishen's avatar
lishen committed
70
    if (threadIdx.x == 0)
71
        internode::shmem_device_barrier_all();
Chenggang Zhao's avatar
Chenggang Zhao committed
72
73
}

74
75
76
77
78
79
void clean_low_latency_buffer(int64_t* clean_0, int num_clean_int_0,
                              int64_t* clean_1, int num_clean_int_1,
                              hipStream_t stream) {
    constexpr int kNumThreads = 256;

    SETUP_LAUNCH_CONFIG(1, kNumThreads, stream);
lishen's avatar
lishen committed
80
81
    LAUNCH_KERNEL_NON_COOPERATIVE(&cfg, clean_low_latency_buffer<kNumThreads>,
                  clean_0, num_clean_int_0, clean_1, num_clean_int_1);
Chenggang Zhao's avatar
Chenggang Zhao committed
82
83
}

lishen's avatar
lishen committed
84
85
86
87
__device__ __forceinline__ void 
internode_ll_putmem_nbi(void* dst_ptr, void* src_ptr,
                        int num_ranks, int dst_rank, int expert_idx,
                        int msg_bytes) {
lishen's avatar
fix  
lishen committed
88
#if defined(FORCE_DUSHMEM_API)
lishen's avatar
lishen committed
89
90
91
92
93
94
95
96
97
98
99
100
101
        internode::shmemx_int8_put_nbi_warp(
            reinterpret_cast<signed char*>(dst_ptr), reinterpret_cast<signed char*>(src_ptr),
            msg_bytes, dst_rank);
#else
    #if defined(ROCM_DISABLE_MULTIQP)
        internode::shmemx_int8_put_nbi_warp(
            reinterpret_cast<signed char*>(dst_ptr), reinterpret_cast<signed char*>(src_ptr),
            msg_bytes, dst_rank);
    #else
        internode::shmemx_int8_put_nbi_warp_dp(
            reinterpret_cast<signed char*>(dst_ptr), reinterpret_cast<signed char*>(src_ptr),
            msg_bytes, (expert_idx + 1) * num_ranks + dst_rank, dst_rank);
    #endif
lishen's avatar
fix  
lishen committed
102
#endif // defined(FORCE_DUSHMEM_API)
lishen's avatar
lishen committed
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
}

__device__ __forceinline__ void 
internode_ll_long_atomic_add(long* dest, const long &value, 
                             int num_ranks, int dst_rank, int expert_idx) {
#if defined(FORCE_DUSHMEM_API)
        internode::shmem_long_atomic_add(dest, value, dst_rank);
#else
        #if defined(ROCM_DISABLE_MULTIQP)
        internode::shmem_long_atomic_add(dest, value, dst_rank);
        #else
        internode::shmem_long_atomic_add_dp(dest, value,
            (expert_idx + 1) * num_ranks + dst_rank, dst_rank);
        #endif
#endif // defined(FORCE_DUSHMEM_API)
}

120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
/**
 * @brief 将 K 个浮点数(BF16/FP32)量化并打包成 INT2(64位)存储
 * 
 * @tparam kQuantType 量化类型 (1: Int8, 2/3: FP8_E4M3/UE8M0, 4: FP8_E5M2)
 * @tparam kNumElemsPerRead 每次读取的元素数量 (通常为 2, 4, 8)
 * @tparam SrcT 源数据类型 (float 或 __hip_bfloat16)
 * @tparam DstT 目标数据类型 (int2 或 int4)
 * @param src_values 源数据数组 (长度 >= kNumElemsPerRead)
 * @param scale 缩放因子 (将 FP32 值映射到量化范围)
 * @param[out] dst_vec 输出的 64 位向量 (int2 或 int4)
 */
template <int kQuantType, int kNumElemsPerRead, typename SrcT, typename DstT>
__forceinline__ __device__ void pack_quantized_values(
    const SrcT* src_values, float scale, DstT& dst_vec) {

    if constexpr (kQuantType == 1) {
        // INT8 量化
        auto int8_ptr = reinterpret_cast<int8_t*>(&dst_vec);
        #pragma unroll
        for (int j = 0; j < kNumElemsPerRead; ++j) {
            // 如果源是 bfloat16,先提升为 float
            float fp32_value_scaled = static_cast<float>(src_values[j]) * scale;
            // 使用 nearbyintf 进行四舍五入
            int8_ptr[j] = static_cast<int8_t>(nearbyintf(fp32_value_scaled));
        }
    } else {
        // FP8 量化 (E4M3, UE8M0, E5M2)
        // 假设 dst_vec 能容纳 kNumElemsPerRead/2 个 fp8x2 元素
        auto fp8x2_ptr = reinterpret_cast<__hip_fp8x2_storage_t*>(&dst_vec);
        #pragma unroll
        for (int j = 0; j < kNumElemsPerRead; j += 2) {
            // 处理两个元素
            float2 fp32x2 = {static_cast<float>(src_values[j]) * scale, static_cast<float>(src_values[j + 1]) * scale};

            if constexpr (kQuantType == 4) {
                // FP8 E5M2
lishen's avatar
lishen committed
156
                fp8x2_ptr[j / 2] = __hip_cvt_float2_to_fp8x2(fp32x2, __HIP_SATFINITE, __HIP_E5M2);
157
158
            } else {
                // FP8 E4M3 或 UE8M0
lishen's avatar
lishen committed
159
                fp8x2_ptr[j / 2] = __hip_cvt_float2_to_fp8x2(fp32x2, __HIP_SATFINITE, __HIP_E4M3);
160
161
162
163
164
165
            }
        }
    }
}

template <int kHidden, int kQuantType=0, int kQuantGroupSize=0, int kMaxNumWarps=16>
lishen's avatar
lishen committed
166
__global__ __launch_bounds__(16 * kWarpSize, 1) void
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
    dispatch(void* packed_recv_x, void* packed_recv_x_scales,
             int* packed_recv_src_info, int64_t* packed_recv_layout_range,
             int* packed_recv_count,
             int* global_atomic_counter,
             void* rdma_recv_x, int64_t* rdma_recv_count, void* rdma_x,
             const void* x, const int64_t* topk_idx,
             int* atomic_counter_per_expert, int* atomic_finish_counter_per_expert,
             int64_t* next_clean, int num_next_clean_int,
             int num_tokens, int num_max_dispatch_tokens_per_rank,
             int num_topk, int num_experts, int rank, int num_ranks,
             int num_warp_groups, int num_warps_per_group,
             bool fp8_round_scale, int phases) {
    // 定义量化类型的枚举
    enum class QuantType {
        None        = 0,        // 不进行量化
        Int8        = 1,        // 采用 Int8 量化
lishen's avatar
lishen committed
183
        FP8_E4M3    = 2,        // 采用 FP8 量化 __HIP_E4M3
184
        FP8_UE8M0   = 3,        // 采用 FP8 量化 DeepseekV3.1的 UE8M0
lishen's avatar
lishen committed
185
        FP8_E5M2    = 4         // 采用 FP8 量化 __HIP_E5M2
186
187
    };

188
189
190
191
192
193
194
195
196
197
    const auto sm_id = static_cast<int>(blockIdx.x);
    const auto thread_id = static_cast<int>(threadIdx.x);
    const auto warp_id = thread_id / kWarpSize, lane_id = get_lane_id();
    const auto num_sms = static_cast<int>(gridDim.x);
    const auto num_warps = num_warp_groups * num_warps_per_group;
    const auto num_local_experts = num_experts / num_ranks;
    const auto warp_group_id = warp_id / num_warps_per_group;
    const auto sub_warp_id = warp_id % num_warps_per_group;
    const auto responsible_expert_idx = sm_id * num_warp_groups + warp_group_id;

lishen's avatar
lishen committed
198
    // May extract UE8M0 from the scales
199
200
    constexpr bool kUseQuant8Bit = kQuantType > 0;
    constexpr bool kUseUE8M0 = kQuantType == 3; // QuantType::FP8_UE8M0
lishen's avatar
lishen committed
201
202
203
204
    using scale_t = std::conditional_t<kUseUE8M0, uint8_t, float>;
    using packed_t = std::conditional_t<kUseUE8M0, uint32_t, float>;
    EP_STATIC_ASSERT(sizeof(packed_t) % sizeof(scale_t) == 0, "Invalid vector length");

205
    // FP8 staffs
206
    constexpr int kNumPerChannels = QUANTIZATION_GROUPSIZE;
lishen's avatar
lishen committed
207
    constexpr int kNumScales = kHidden / kNumPerChannels;
208
    const size_t hidden_bytes = kHidden * (kUseQuant8Bit ? sizeof(__hip_fp8_storage_t) : sizeof(hip_bfloat16));
209
210
    const size_t hidden_int4 = hidden_bytes / sizeof(int4);

lishen's avatar
lishen committed
211
    // Message package: hidden data, FP8 scales, index at source
212
    // NOTES: currently we have 3 reserved int fields for future use
213
    using vec_t = typename std::conditional<kUseQuant8Bit, int2, int4>::type;
lishen's avatar
lishen committed
214
215
    constexpr size_t num_bytes_per_msg = sizeof(int4) + 
        (kUseQuant8Bit ? (kHidden + (kQuantGroupSize == 0 ? 4 : kNumScales) * sizeof(float)) : (kHidden * sizeof(hip_bfloat16)));
lishen's avatar
lishen committed
216
217
    EP_STATIC_ASSERT(num_bytes_per_msg % sizeof(int4) == 0, "Invalid message size");
    constexpr size_t num_int4_per_msg = num_bytes_per_msg / sizeof(int4);
218

lishen's avatar
lishen committed
219
    // Expert counts
lishen's avatar
lishen committed
220
    __shared__ int shared_num_tokens_sent_per_expert[kMaxNumWarps];
lishen's avatar
lishen committed
221
222

    // Sending phase
223
224
225
226
227
228
    if ((phases & LOW_LATENCY_SEND_PHASE) == 0)
        goto LOW_LATENCY_DISPATCH_RECV;

    // There are 2 kinds of warps in this part:
    // 1. The first-kind warps for FP8 cast and sending top-k tokens
    // 2. The last warp for reading `topk_idx` and count for per-expert information
lishen's avatar
lishen committed
229
230
    if (warp_id < num_warps) {
        constexpr int kNumElemsPerRead = sizeof(int4) / sizeof(hip_bfloat16);
231
        constexpr int kNumThreadPerGroup = QUANTIZATION_GROUPSIZE / kNumElemsPerRead;
lishen's avatar
lishen committed
232
        // EP_DEVICE_ASSERT(kHidden % kNumElemsPerRead == 0);
233
        EP_STATIC_ASSERT(kNumElemsPerRead * kWarpSize % kNumPerChannels == 0, "Invalid vectorization");
lishen's avatar
lishen committed
234
        const auto num_threads = num_warps * kWarpSize;
lishen's avatar
lishen committed
235
        constexpr int hidden_bf16_int4 = kHidden / kNumElemsPerRead;
236
237

        for (int token_idx = sm_id; token_idx < num_tokens; token_idx += num_sms) {
lishen's avatar
lishen committed
238
239
            const auto x_int4 = reinterpret_cast<const int4*>(x) + token_idx * hidden_bf16_int4;
            const auto rdma_x_src_idx = reinterpret_cast<int*>(reinterpret_cast<uint8_t*>(rdma_x) + token_idx * num_bytes_per_msg);
240
241
242
            const auto rdma_x_vec = reinterpret_cast<vec_t*>(reinterpret_cast<uint8_t*>(rdma_x_src_idx) + sizeof(int4));
            const auto rdma_x_scales = reinterpret_cast<float*>(reinterpret_cast<uint8_t*>(rdma_x_vec) + hidden_bytes);

lishen's avatar
lishen committed
243
            // Overlap top-k index read and source token index write
244
245
246
            auto dst_expert_idx = warp_id < num_topk ? static_cast<int>(__ldg(topk_idx + token_idx * num_topk + warp_id)) : -1;
            thread_id == 0 ? (*rdma_x_src_idx = token_idx) : 0;

247
248
249
            // 用于记录per-channel量化的amax
            __shared__ float channel_amaxf[kNumScales];
            if constexpr(kUseQuant8Bit && kQuantGroupSize == 0) {
lishen's avatar
lishen committed
250
                if (thread_id < kNumScales) {
lishen's avatar
lishen committed
251
                    channel_amaxf[thread_id] = 0.0;
lishen's avatar
lishen committed
252
253
254
255
                }
                __syncthreads();
            }

256
257
258
259
260
261
            // FP8 cast
            #pragma unroll
            for (int i = thread_id; i < hidden_bf16_int4; i += num_threads) {
                // Read
                auto int4_value = __ldg(x_int4 + i);

262
                if constexpr(kUseQuant8Bit) {
263
264
265
                    // Calculate local amax
                    auto bf16_values = reinterpret_cast<hip_bfloat16*>(&int4_value);
                    float fp32_values[kNumElemsPerRead];
lishen's avatar
lishen committed
266
                    float amax = 0.0, scale, scale_inv;
267
                    #pragma unroll
lishen's avatar
lishen committed
268
                    for (int j = 0; j < kNumElemsPerRead; ++ j) {
269
270
271
272
273
                        fp32_values[j] = static_cast<float>(bf16_values[j]);
                        amax = fmaxf(amax, fabsf(fp32_values[j]));
                    }
                    // Reduce amax and scale
                    EP_STATIC_ASSERT(kNumElemsPerRead * kWarpSize / kNumPerChannels == 4, "Invalid vectorization");
274
275
                    amax = warp_reduce_max<kNumThreadPerGroup>(amax);
                    const int scale_offset = i * kNumElemsPerRead / QUANTIZATION_GROUPSIZE;
lishen's avatar
lishen committed
276

277
                    if constexpr(kQuantGroupSize == 0) {
lishen's avatar
lishen committed
278
                        // 记录每128个数的最大值
279
                        channel_amaxf[scale_offset] = fmaxf(amax, channel_amaxf[scale_offset]);
lishen's avatar
lishen committed
280
                    } else {
281
282
                        calculate_quant8bit_scales<kQuantType>(amax, scale, scale_inv, fp8_round_scale);
                        if (lane_id % kNumThreadPerGroup == 0)
lishen's avatar
lishen committed
283
284
285
286
                            rdma_x_scales[scale_offset] = scale_inv;

                        // Cast into send buffer
                        vec_t int2_value;
287
                        pack_quantized_values<kQuantType, kNumElemsPerRead>(fp32_values, scale, int2_value);
lishen's avatar
lishen committed
288
                        rdma_x_vec[i] = int2_value;
289
290
291
292
293
294
295
                    }
                } else {
                    // Reinterpret-cast is for C++14 compatibility
                    rdma_x_vec[i] = *reinterpret_cast<vec_t*>(&int4_value);
                }
            }
            __syncthreads();
lishen's avatar
lishen committed
296

297
            if constexpr(kUseQuant8Bit && kQuantGroupSize == 0) {
lishen's avatar
lishen committed
298
                float amax_per_token = 0.0;
lishen's avatar
lishen committed
299
300
301
302
303
                // 并行规约,计算每个token的amax
                for (int s = 0; s < kNumScales; s+=kWarpSize) {
                    int src_idx = s + lane_id;
                    float tmp_amaxf = 0;
                    if(src_idx < kNumScales) {
304
                        tmp_amaxf = channel_amaxf[src_idx];
lishen's avatar
lishen committed
305
306
                    }
                    tmp_amaxf = warp_reduce_max<kWarpSize>(tmp_amaxf);
307
                    channel_amaxf[0] = fmaxf(tmp_amaxf, channel_amaxf[0]);
lishen's avatar
lishen committed
308
309
                    __syncthreads();
                }
310
                amax_per_token = channel_amaxf[0];
lishen's avatar
lishen committed
311
312
313

                // 根据最大值计算scale
                float scale, scale_inv;
lishen's avatar
lishen committed
314
                calculate_quant8bit_scales<kQuantType>(amax_per_token, scale, scale_inv, fp8_round_scale);
lishen's avatar
lishen committed
315
316
317
318
319
320
321
322
323
324
325
                if (thread_id == 0) {
                    rdma_x_scales[0] = scale_inv;
                }

                for (int i = thread_id; i < hidden_bf16_int4; i += num_threads) {
                    // Read
                    auto int4_value = __ldg(x_int4 + i);
                    auto bf16_values = reinterpret_cast<hip_bfloat16*>(&int4_value);

                    // Cast into send buffer
                    vec_t int2_value;
326
                    pack_quantized_values<kQuantType, kNumElemsPerRead>(bf16_values, scale, int2_value);
lishen's avatar
lishen committed
327
328
329
330
331
                    rdma_x_vec[i] = int2_value;
                }
                __syncthreads();
            }

332
333
334
335
            // Issue IBGDA sends
            if (dst_expert_idx >= 0) {
                int slot_idx = lane_id == 0 ? atomicAdd(atomic_counter_per_expert + dst_expert_idx, 1) : 0;
                slot_idx = shfl_sync(slot_idx, 0);
lishen's avatar
lishen committed
336
337
                const auto dst_rank = dst_expert_idx / num_local_experts;
                const auto dst_expert_local_idx = dst_expert_idx % num_local_experts;
338
339
340
                const auto src_ptr = reinterpret_cast<uint64_t>(rdma_x_src_idx);
                const auto dst_ptr = reinterpret_cast<uint64_t>(rdma_recv_x) +
                                     dst_expert_local_idx * num_ranks * num_max_dispatch_tokens_per_rank * num_bytes_per_msg +
lishen's avatar
lishen committed
341
342
                                     rank * num_max_dispatch_tokens_per_rank * num_bytes_per_msg +
                                     slot_idx * num_bytes_per_msg;
lishen's avatar
lishen committed
343
344
345
346
347

                // 通过 shmem_get_p2p_ptr 获取 当前远程指针能否可达
                uint64_t p2p_ptr = internode::shmem_get_p2p_ptr((void*)dst_ptr, rank, dst_rank);
                if (p2p_ptr == 0) {  // RDMA
                    internode_ll_putmem_nbi((void*)dst_ptr, (void*)src_ptr,
lishen's avatar
lishen committed
348
349
                                            num_ranks, dst_rank, dst_expert_local_idx,
                                            num_bytes_per_msg);
lishen's avatar
lishen committed
350
                } else { //  本地 GPU 和 同一计算节点的 其他 GPU 地址
351
352
                    // NOTES: only 2 load iterations for 7K hidden with 8 unrolls
                    const auto* src_int4_ptr = reinterpret_cast<const int4*>(src_ptr);
lishen's avatar
lishen committed
353
                    const auto* dst_int4_ptr = reinterpret_cast<int4*>(p2p_ptr);
lishen's avatar
lishen committed
354
                    UNROLLED_WARP_COPY_LL(8, lane_id, num_int4_per_msg, dst_int4_ptr, src_int4_ptr, ld_nc_global, st_na_global);
355
                }
lishen's avatar
lishen committed
356

357
358
359
360
361
                // Increase counter after finishing
                syncwarp();
                lane_id == 0 ? atomic_add_release_global(atomic_finish_counter_per_expert + dst_expert_idx, 1) : 0;
            }
        }
lishen's avatar
lishen committed
362
363
    }
    if (warp_id == num_warps - 1) {
lishen's avatar
lishen committed
364
        // EP_DEVICE_ASSERT(num_sms > 1);
365
        if (sm_id == 0) {
lishen's avatar
lishen committed
366
            // The first SM is also responsible for checking QPs
367
368
369
370
371
372
373
374
375
376
377
378
            // The first SM is also responsible for cleaning the next buffer
            #pragma unroll
            for (int i = lane_id; i < num_next_clean_int; i += kWarpSize)
                next_clean[i] = 0;

            // Notify before executing `int_p`
            syncwarp();
            #pragma unroll
            for (int i = lane_id; i < num_experts; i += kWarpSize)
                atomic_add_release_global(atomic_finish_counter_per_expert + i, FINISHED_SUM_TAG);
        }
        // This SM should be responsible for some destination experts, read `topk_idx` for them
lishen's avatar
lishen committed
379
        int expert_count[kMaxNumWarps] = {0};
380
381
382
383
384
385
386
387
        const auto expert_begin_idx = sm_id * num_warp_groups;
        const auto expert_end_idx = min(expert_begin_idx + num_warp_groups, num_experts);

        // Per lane count
        #pragma unroll 8
        for (int i = lane_id; i < num_tokens * num_topk; i += kWarpSize) {
            auto idx = static_cast<int>(__ldg(topk_idx + i));
            if (idx >= expert_begin_idx and idx < expert_end_idx)
lishen's avatar
lishen committed
388
                expert_count[idx - expert_begin_idx] ++;
389
390
391
392
        }

        // Warp reduce
        #pragma unroll
lishen's avatar
lishen committed
393
        for (int i = expert_begin_idx; i < expert_end_idx; ++ i) {
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
            auto sum = warp_reduce_sum(expert_count[i - expert_begin_idx]);
            if (lane_id == 0) {
                shared_num_tokens_sent_per_expert[i - expert_begin_idx] = sum;
                atomic_add_release_global(atomic_finish_counter_per_expert + i, FINISHED_SUM_TAG - sum);
            }
        }
    }
    __syncthreads();

    // Issue count sends
    if (responsible_expert_idx < num_experts and sub_warp_id == 0 and lane_id == 0) {
        const auto dst_rank = responsible_expert_idx / num_local_experts;
        const auto dst_expert_local_idx = responsible_expert_idx % num_local_experts;
        const auto num_tokens_sent = shared_num_tokens_sent_per_expert[responsible_expert_idx - sm_id * num_warp_groups];

        // Wait local sends issued and send expert counts
        while (ld_acquire_global(atomic_finish_counter_per_expert + responsible_expert_idx) != FINISHED_SUM_TAG * 2);
lishen's avatar
lishen committed
411
412
413
414
415
416
417
418
419

        auto dst_ptr = rdma_recv_count + dst_expert_local_idx * num_ranks + rank;
        // 通过 shmem_get_p2p_ptr 获取 当前远程指针能否可达
        uint64_t p2p_ptr = internode::shmem_get_p2p_ptr((void*)dst_ptr, rank, dst_rank);
        if (p2p_ptr == 0) {  // RDMA
            internode_ll_long_atomic_add(dst_ptr, -num_tokens_sent - 1, 
                                         num_ranks, dst_rank, dst_expert_local_idx);
        } else { //  本地 GPU 和 同一计算节点的 其他 GPU 地址
            st_na_release(reinterpret_cast<int *>(p2p_ptr), -num_tokens_sent - 1);
420
421
422
423
424
425
426
427
428
429
430
431
        }

        // Clean workspace for next use
        atomic_counter_per_expert[responsible_expert_idx] = 0;
        atomic_finish_counter_per_expert[responsible_expert_idx] = 0;

        // Clean `packed_recv_count`
        if (dst_rank == 0)
            packed_recv_count[dst_expert_local_idx] = 0;
    }
    syncwarp();

lishen's avatar
lishen committed
432
433
    // Receiving phase
LOW_LATENCY_DISPATCH_RECV:
434
435
436
437
438
439
440
441
    if ((phases & LOW_LATENCY_RECV_PHASE) == 0)
        return;

    // For send-and-recv kernels, we need a grid sync for making `packed_recv_count` visible
    if (phases & LOW_LATENCY_SEND_PHASE){
        grid_barrier(global_atomic_counter, num_sms);
    }

lishen's avatar
lishen committed
442
443
444
445
    // 16 is the max possible number of warps in AMD GPUs
    constexpr int num_sync_large_iteration = kMaxNumWarps ;
    __shared__ volatile int sync_large_warp_counters[num_sync_large_iteration];

446
    #pragma unroll
lishen's avatar
lishen committed
447
448
449
450
451
    for (int i = thread_id; i < num_sync_large_iteration; i += blockDim.x) {
        sync_large_warp_counters[i] = 0;
    }
    __syncthreads();

452
453
454
455
    // Receiving and packing
    if (responsible_expert_idx < num_experts) {
        const auto src_rank = responsible_expert_idx / num_local_experts;
        const auto local_expert_idx = responsible_expert_idx % num_local_experts;
lishen's avatar
lishen committed
456
        const auto rdma_recv_x_uint8 = reinterpret_cast<uint8_t*>(rdma_recv_x) +
lishen's avatar
lishen committed
457
458
                                       local_expert_idx * num_ranks * num_max_dispatch_tokens_per_rank * num_bytes_per_msg +
                                       src_rank * num_max_dispatch_tokens_per_rank * num_bytes_per_msg;
lishen's avatar
lishen committed
459
        const auto recv_x_int4 = reinterpret_cast<int4*>(packed_recv_x) +
lishen's avatar
lishen committed
460
                                 local_expert_idx * num_ranks * num_max_dispatch_tokens_per_rank * hidden_int4;
461
462
        const auto recv_src_info = packed_recv_src_info + local_expert_idx * num_ranks * num_max_dispatch_tokens_per_rank;
        const auto recv_range = packed_recv_layout_range + local_expert_idx * num_ranks;
lishen's avatar
lishen committed
463
        const auto num_aligned_scales = ALIGN<int>(kNumScales, sizeof(float) / sizeof(scale_t));
lishen's avatar
lishen committed
464
        const auto recv_x_scales = static_cast<scale_t*>(packed_recv_x_scales) +
lishen's avatar
lishen committed
465
                                   local_expert_idx * num_ranks * num_max_dispatch_tokens_per_rank *
lishen's avatar
lishen committed
466
                                       (kQuantGroupSize == 0 ? 1 : num_aligned_scales);
467
468

        // Shared between sub-warps in warp groups
lishen's avatar
lishen committed
469
        __shared__ int shared_num_recv_tokens[kMaxNumWarps], shared_recv_token_begin_idx[kMaxNumWarps];
470
471
472

        // Wait tokens to arrive
        // NOTES: using sub-warp 1 to overlap with sub-warp 0
lishen's avatar
lishen committed
473
        int num_recv_tokens, recv_token_begin_idx;
lishen's avatar
lishen committed
474
        // EP_DEVICE_ASSERT(num_warps_per_group > 1);
475
476

        if (sub_warp_id == 1 and lane_id == 0) {
lishen's avatar
lishen committed
477
            while ((num_recv_tokens = ld_acquire_global(reinterpret_cast<int*>(rdma_recv_count + local_expert_idx * num_ranks + src_rank))) == 0);
478
            num_recv_tokens = -num_recv_tokens - 1;
lishen's avatar
lishen committed
479
480
            recv_token_begin_idx = atomicAdd(packed_recv_count + local_expert_idx, num_recv_tokens);
            shared_num_recv_tokens[warp_group_id] = num_recv_tokens;
481
            shared_recv_token_begin_idx[warp_group_id] = recv_token_begin_idx;
lishen's avatar
lishen committed
482
            recv_range[src_rank] = pack2<int, int64_t>(num_recv_tokens, recv_token_begin_idx);
483
484
485
486
        }

        // no needs to reset because there is no iteration
        if (lane_id == 0){
lishen's avatar
lishen committed
487
            volatile int ret = __hip_atomic_fetch_add(&sync_large_warp_counters[warp_group_id], 1, __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_WORKGROUP);
488
489
490
        }
        syncwarp();

lishen's avatar
lishen committed
491
        while (sync_large_warp_counters[warp_group_id] < num_warps_per_group);
492
493
494
495
        num_recv_tokens = shared_num_recv_tokens[warp_group_id];
        recv_token_begin_idx = shared_recv_token_begin_idx[warp_group_id];

        // Copy tokens
lishen's avatar
lishen committed
496
        EP_STATIC_ASSERT(kNumScales <= 64, "Invalid hidden size");
497
498
499
500
501
502
503
504
505
506
507
        for (int i = sub_warp_id; i < num_recv_tokens; i += num_warps_per_group) {
            // Copy source info
            const auto src_src_idx = reinterpret_cast<int*>(rdma_recv_x_uint8 + i * num_bytes_per_msg);
            if (lane_id == 0)
                recv_src_info[recv_token_begin_idx + i] = ld_nc_global(src_src_idx);
            syncwarp();

            // Copy data
            // NOTES: only 2 load iterations for 7K hidden with 7 unrolls
            const auto src_data = reinterpret_cast<int4*>(reinterpret_cast<uint8_t*>(src_src_idx) + sizeof(int4));
            const auto dst_data = recv_x_int4 + (recv_token_begin_idx + i) * hidden_int4;
lishen's avatar
lishen committed
508
            UNROLLED_WARP_COPY_LL(7, lane_id, hidden_int4, dst_data, src_data, ld_nc_global, st_na_global);
509
510

            // Copy scales
511
            if constexpr(kUseQuant8Bit) {
512
                const auto src_scales = reinterpret_cast<float*>(reinterpret_cast<uint8_t*>(src_data) + hidden_bytes);
lishen's avatar
lishen committed
513
514
515
516
517
                const auto num_elems_per_pack = static_cast<int>(sizeof(packed_t) / sizeof(scale_t));
                const auto token_idx = recv_token_begin_idx + i;
                const auto token_stride = num_elems_per_pack;
                const auto pack_stride = num_ranks * num_max_dispatch_tokens_per_rank * num_elems_per_pack;

lishen's avatar
lishen committed
518
                if constexpr(kQuantGroupSize == 0) {
lishen's avatar
lishen committed
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
                    if (lane_id == 0) {
                        recv_x_scales[token_idx] = ld_nc_global(src_scales);
                    }
                } else {
                    if (lane_id < kNumScales) {
                        const auto pack_idx = lane_id / num_elems_per_pack;
                        const auto elem_idx = lane_id % num_elems_per_pack;
                        auto scale = extract_required_scale_format<kUseUE8M0>(ld_nc_global(src_scales + lane_id));
                        recv_x_scales[token_idx * token_stride + pack_idx * pack_stride + elem_idx] = scale;
                    }
                    if (lane_id + kWarpSize < kNumScales) {
                        const auto pack_idx = (lane_id + kWarpSize) / num_elems_per_pack;
                        const auto elem_idx = (lane_id + kWarpSize) % num_elems_per_pack;
                        auto scale = extract_required_scale_format<kUseUE8M0>(ld_nc_global(src_scales + lane_id + kWarpSize));
                        recv_x_scales[token_idx * token_stride + pack_idx * pack_stride + elem_idx] = scale;
                    }
lishen's avatar
lishen committed
535
                }
536
537
538
            }
        }
    }
Chenggang Zhao's avatar
Chenggang Zhao committed
539
540
}

lishen's avatar
lishen committed
541
void dispatch(void* packed_recv_x, void* packed_recv_x_scales,
lishen's avatar
lishen committed
542
              int* packed_recv_src_info, int64_t* packed_recv_layout_range,
543
              int* packed_recv_count,
544
              int* global_atomic_counter,
lishen's avatar
lishen committed
545
546
547
548
              void* rdma_recv_x, int64_t* rdma_recv_count, void* rdma_x,
              const void* x, const int64_t* topk_idx,
              int64_t* next_clean, int num_next_clean_int,
              int num_tokens, int hidden, int num_max_dispatch_tokens_per_rank,
lishen's avatar
lishen committed
549
              int num_topk, int num_experts, int rank, int num_ranks,
550
              int quant_type, int quant_group_size, bool fp8_round_scale,
lishen's avatar
lishen committed
551
              void* workspace, int num_device_sms,
lishen's avatar
lishen committed
552
              hipStream_t stream, int phases) {
553
    constexpr int kMaxNumWarps = 16;
554
    constexpr int kNumMaxTopK = 11;
lishen's avatar
lishen committed
555
    const int num_warp_groups = ceil_div(num_experts, num_device_sms);
556
    const int num_warps_per_group = kMaxNumWarps / num_warp_groups;
557
558
559
560
    EP_HOST_ASSERT(num_warp_groups > 0 and num_warps_per_group > 0);
    EP_HOST_ASSERT(kNumMaxTopK + 1 <= num_warp_groups * num_warps_per_group);

    const auto num_warps = num_warp_groups * num_warps_per_group;
lishen's avatar
lishen committed
561
    const auto num_sms = ceil_div(num_experts, num_warp_groups);
562
563
564
    EP_HOST_ASSERT(num_topk <= kNumMaxTopK);

    // Workspace checks
lishen's avatar
lishen committed
565
    auto atomic_counter_per_expert = reinterpret_cast<int*>(workspace);
566
567
568
    auto atomic_finish_counter_per_expert = atomic_counter_per_expert + num_experts;
    EP_HOST_ASSERT(num_experts * sizeof(int) * 2 <= NUM_WORKSPACE_BYTES);

569
570
571
572
573
574
    // 限制groupsize的大小
    EP_HOST_ASSERT(quant_group_size == 0 || quant_group_size == 128);

    /*量化类型枚举
    0 -> None          不量化,保持原始精度
    1 -> Int8          使用 INT8 对称量化
lishen's avatar
lishen committed
575
    2 -> FP8_E4M3      使用 FP8 E4M3 格式 (__HIP_E4M3)
576
    3 -> FP8_UE8M0     使用 DeepSeekV3.1 提出的 UE8M0 格式 (仅支持round_scale=True)
lishen's avatar
lishen committed
577
    4 -> FP8_E5M2      使用 FP8 E5M2 格式 (__HIP_E5M2)
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
    */

#define DISPATCH_LAUNCH_CASE(hidden)                                                \
  {                                                                                 \
    auto dispatch_func = dispatch<hidden, 0, 0, kMaxNumWarps>;                      \
    if (quant_group_size == 0) {                                                    \
        switch (quant_type) {                                                       \
            case 1: dispatch_func = dispatch<hidden, 1, 0, kMaxNumWarps>; break;    \
            case 2: dispatch_func = dispatch<hidden, 2, 0, kMaxNumWarps>; break;    \
            case 3: dispatch_func = dispatch<hidden, 3, 0, kMaxNumWarps>; break;    \
            case 4: dispatch_func = dispatch<hidden, 4, 0, kMaxNumWarps>; break;    \
        }                                                                           \
    } else {                                                                        \
        switch (quant_type) {                                                       \
            case 1: dispatch_func = dispatch<hidden, 1, 128, kMaxNumWarps>; break;  \
            case 2: dispatch_func = dispatch<hidden, 2, 128, kMaxNumWarps>; break;  \
            case 3: dispatch_func = dispatch<hidden, 3, 128, kMaxNumWarps>; break;  \
            case 4: dispatch_func = dispatch<hidden, 4, 128, kMaxNumWarps>; break;  \
        }                                                                           \
    }                                                                               \
    LAUNCH_KERNEL_NON_COOPERATIVE(&cfg, dispatch_func,                              \
        packed_recv_x, packed_recv_x_scales,                                        \
        packed_recv_src_info, packed_recv_layout_range, packed_recv_count,          \
        global_atomic_counter,                                                      \
        rdma_recv_x, rdma_recv_count, rdma_x, x, topk_idx,                          \
        atomic_counter_per_expert, atomic_finish_counter_per_expert,                \
        next_clean, num_next_clean_int,                                             \
        num_tokens, num_max_dispatch_tokens_per_rank,                               \
        num_topk, num_experts, rank, num_ranks,                                     \
        num_warp_groups, num_warps_per_group, fp8_round_scale, phases);             \
  }                                                                                 \
  break
610
611
612
613

    SETUP_LAUNCH_CONFIG(num_sms, num_warps * kWarpSize, stream);
    SWITCH_HIDDEN(DISPATCH_LAUNCH_CASE);
#undef DISPATCH_LAUNCH_CASE
614
615
}

616
template <bool kUseLogFMT, int kHidden, int kNumMaxTopk, int kMaxNumWarps=16>
lishen's avatar
lishen committed
617
__global__ __launch_bounds__(16 * kWarpSize, 1) void
lishen's avatar
lishen committed
618
619
620
621
622
combine(void* combined_x,
        void* rdma_recv_x, int64_t* rdma_recv_flag, void* rdma_send_x,
        const void* x, const int64_t* topk_idx, const float* topk_weights,
        const int* src_info, const int64_t* layout_range,
        int* global_atomic_counter,
lishen's avatar
lishen committed
623
        int64_t* combine_wait_recv_cost_stats,
lishen's avatar
lishen committed
624
625
626
627
628
        int64_t* next_clean, int num_next_clean_int,
        int* atomic_clean_flag,
        int num_combined_tokens, int hidden, int num_topk,
        int num_max_dispatch_tokens_per_rank,
        int num_experts, int rank, int num_ranks,
lishen's avatar
lishen committed
629
        int num_warp_groups, int num_warps_per_group,
lishen's avatar
lishen committed
630
631
632
633
634
635
636
        int phases, bool zero_copy) {
    const auto sm_id = static_cast<int>(blockIdx.x);
    const auto num_sms = static_cast<int>(gridDim.x);
    const auto thread_id = static_cast<int>(threadIdx.x);
    const auto num_threads = static_cast<int>(blockDim.x);
    const auto warp_id = thread_id / kWarpSize, lane_id = get_lane_id();
    const auto num_local_experts = num_experts / num_ranks;
lishen's avatar
lishen committed
637
638
    const auto warp_group_id = warp_id / num_warps_per_group;
    const auto sub_warp_id = warp_id % num_warps_per_group;
lishen's avatar
lishen committed
639
    const auto num_warps = num_threads / kWarpSize;
lishen's avatar
lishen committed
640
    const auto responsible_expert_idx = sm_id * num_warp_groups + warp_group_id;
lishen's avatar
lishen committed
641
642
643
644
645
646

    // Data type staffs
    constexpr int kNumElemsPerInt4 = sizeof(int4) / sizeof(hip_bfloat16);
    const size_t hidden_bf16_int4 = kHidden / kNumElemsPerInt4;

    // Message package
647
    EP_STATIC_ASSERT(kHidden % QUANTIZATION_GROUPSIZE == 0, "Invalid hidden");
648
649
650
651
652
653
654
655
656
657
658
659

    /////////////// LogFMT使用 ///////////////
    constexpr int bSupportLogFMT = kUseLogFMT && hidden_bf16_int4 % (kWarpSize * 2) == 0;
    constexpr int kNumSendUnrolls = bSupportLogFMT ? 2 : 1;
    constexpr int kNumRecvUnrolls = bSupportLogFMT ? 2 : 1;
    constexpr int kNumMsgInt4ElemPerWarp = kWarpSize * kNumSendUnrolls; // 每个warp发送的int4元素数据量,即每个warp发送 kNumMsgInt4ElemPerWarp*sizeof(int4)/sizeof(bfloat16)
    EP_STATIC_ASSERT(hidden_bf16_int4 % (kNumSendUnrolls * kWarpSize) == 0, "Invalid hidden");
    EP_STATIC_ASSERT(kNumSendUnrolls >= kNumRecvUnrolls, "Invalid unroll factors");

    constexpr int kNumDivisions = kHidden / QUANTIZATION_GROUPSIZE;
    constexpr int kNumMetaBytes = kNumDivisions * sizeof(__hip_bfloat162);  // 用于记录数据的最大最小值
    constexpr int kNumSendLogFMTBytes = kNumMsgInt4ElemPerWarp * sizeof(int4);
lishen's avatar
lishen committed
660
    constexpr int kNumStages = 3;  // 使用kNumStages>1,则需要的LDS大于64KB
661
662
663
664
665
    constexpr int kLogFMTShmemSize = kMaxNumWarps * (kNumStages * kNumSendLogFMTBytes + kNumMetaBytes);
    __shared__ uint8_t smem_buffer[kLogFMTShmemSize];
    /////////////////////////////////////////////

    constexpr size_t num_bytes_per_slot = kHidden * sizeof(hip_bfloat16) + kNumMetaBytes;
lishen's avatar
lishen committed
666
    EP_STATIC_ASSERT(num_bytes_per_slot % sizeof(int4) == 0, "Invalid vectorization");
lishen's avatar
lishen committed
667

668
    // 初始化用于细粒度warp间同步的计数器数组
lishen's avatar
lishen committed
669
670
671
672
673
674
675
676
    __shared__ volatile int sync_large_warp_counters[kMaxNumWarps];
    if (threadIdx.x==0){
        #pragma unroll
        for (int i = 0; i < kMaxNumWarps; ++i) {
            sync_large_warp_counters[i] = 0;
        }
    }
    __syncthreads();
677

lishen's avatar
lishen committed
678
679
680
    // Sending phase
    if ((phases & LOW_LATENCY_SEND_PHASE) == 0)
        goto LOW_LATENCY_COMBINE_RECV;
Chenggang Zhao's avatar
Chenggang Zhao committed
681

lishen's avatar
lishen committed
682
683
684
685
686
    // Clean up next buffer
    if (sm_id == 0 and warp_group_id == 0 and sub_warp_id == 0) {
        #pragma unroll
        for (int i = lane_id; i < num_next_clean_int; i += kWarpSize)
            next_clean[i] = 0;
687

lishen's avatar
lishen committed
688
689
690
691
692
        // Notify before executing `int_p`
        syncwarp();
        if (lane_id == 0)
            atomic_add_release_global(atomic_clean_flag, num_experts);
    }
693

lishen's avatar
lishen committed
694
695
696
697
698
699
700
    // Issue IBGDA sends
    if (responsible_expert_idx < num_experts) {
        const auto dst_rank = responsible_expert_idx / num_local_experts;
        const auto local_expert_idx = responsible_expert_idx % num_local_experts;
        const auto global_expert_idx = rank * num_local_experts + local_expert_idx;
        const auto layout = __ldg(layout_range + local_expert_idx * num_ranks + dst_rank);
        const auto local_x = reinterpret_cast<const int4*>(x) +
lishen's avatar
lishen committed
701
                             local_expert_idx * num_ranks * num_max_dispatch_tokens_per_rank * hidden_bf16_int4;
lishen's avatar
lishen committed
702
703
        const auto local_src_info = src_info + local_expert_idx * num_ranks * num_max_dispatch_tokens_per_rank;
        const auto rdma_send_x_vec = reinterpret_cast<uint8_t*>(rdma_send_x) +
lishen's avatar
lishen committed
704
                                     local_expert_idx * num_ranks * num_max_dispatch_tokens_per_rank * num_bytes_per_slot;
705
706
707
708
709
710
        // 用于logfmt的LDS
        auto smem_ptr = smem_buffer + warp_id * (kNumStages * kNumSendLogFMTBytes + kNumMetaBytes);
        // 存储logfmt的起始地址,并根据stage_idx进行索引块
        auto logfmt_buffers = PatternVisitor([=](const int& i) { return reinterpret_cast<int4*>(smem_ptr + i * kNumSendLogFMTBytes); });
        // 存储logfmt的最大最小值
        auto meta_buffers = bSupportLogFMT ? reinterpret_cast<__hip_bfloat162*>(smem_ptr + kNumStages * kNumSendLogFMTBytes) : nullptr;
lishen's avatar
lishen committed
711
712
713
714
715
716
717
718
719
720
721
        // 用于多buffer时临时存储
        auto get_num_logfmt_bytes = [&](const int& offset_int4) {
            return min(kNumSendLogFMTBytes, static_cast<int>((hidden_bf16_int4 - offset_int4) * sizeof(int4)));
        };
        // 简化从global到LDS的存储写法
        auto logfmt_load_global2lds = [&](const int& stage_idx, const int4* gmem_ptr, const int& num_bytes) {
            UNROLLED_WARP_COPY_LL(1, lane_id, num_bytes / sizeof(int4),
                reinterpret_cast<int4 *>(logfmt_buffers[stage_idx]),
                reinterpret_cast<const int4 *>(gmem_ptr),
                ld_direct_global, st_na_global);
        };
lishen's avatar
lishen committed
722
723
724
725
726
727

        // Unpack layout
        int offset, num_tokens_to_send;
        unpack2(layout, num_tokens_to_send, offset);

        // Issue IBGDA send
lishen's avatar
lishen committed
728
        for (int token_idx = offset + sub_warp_id; token_idx < offset + num_tokens_to_send; token_idx += num_warps_per_group) {
lishen's avatar
lishen committed
729
730
            const auto x_int4 = local_x + token_idx * hidden_bf16_int4;
            const auto rdma_send_type_row = reinterpret_cast<int*>(rdma_send_x_vec + token_idx * num_bytes_per_slot);
lishen's avatar
lishen committed
731
            const auto rdma_send_x_vec_row = reinterpret_cast<uint8_t*>(rdma_send_type_row);
lishen's avatar
lishen committed
732
733

            // Copy directly to local rank, or copy to buffer and issue RDMA
734
            const auto src_idx = __ldg(local_src_info + token_idx);
lishen's avatar
lishen committed
735
            const auto buf_ptr = reinterpret_cast<int64_t>(rdma_send_x_vec_row);
lishen's avatar
lishen committed
736
            const auto dst_ptr = reinterpret_cast<uint64_t>(rdma_recv_x) + (global_expert_idx * num_max_dispatch_tokens_per_rank + src_idx) * num_bytes_per_slot;
lishen's avatar
lishen committed
737

738
739
740
741
742
743
            // 采用logfmt或者直接拷贝
            uint64_t dst_p2p_ptr = internode::shmem_get_p2p_ptr((void*)dst_ptr, rank, dst_rank);
            int num_send_bytes = hidden * sizeof(hip_bfloat16);

            if (not zero_copy or dst_p2p_ptr != 0) {
                const auto cpy_src_int4_ptr = zero_copy ? reinterpret_cast<int4*>(buf_ptr) : x_int4;
lishen's avatar
lishen committed
744
                const auto cpy_dst_int4_ptr = dst_p2p_ptr == 0 ? reinterpret_cast<int4*>(buf_ptr) : reinterpret_cast<int4*>(dst_p2p_ptr);
745

lishen's avatar
lishen committed
746
747
                constexpr int kNumIters = hidden_bf16_int4 / kNumMsgInt4ElemPerWarp;
                EP_STATIC_ASSERT(kNumIters >= 1, "hidden length too small");
748

lishen's avatar
lishen committed
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
                if constexpr (bSupportLogFMT) {
                    // ===== LogFMT 路径:使用 LDS + encode + 多级流水 =====
                    int logfmt_offset_bytes = kNumMetaBytes;
                    // meta_buffers 存储的thread间隔
                    constexpr int kNumInt4PerDivision = 128 / kNumElemsPerInt4;
                    // 记录S1~S3的编码字节数
                    int encoded_bytes[kNumStages];

                    // Prefetch: iter0执行S1
                    logfmt_load_global2lds(0, cpy_src_int4_ptr, get_num_logfmt_bytes(0));
                    syncwarp();

                    // Prefetch: iter0执行S2, iter1执行S1
                    if (kNumStages > 2 && kNumIters > 1) {
                        int warp_offset = /*1 * */kNumMsgInt4ElemPerWarp;
                        logfmt_load_global2lds(1, cpy_src_int4_ptr + warp_offset, get_num_logfmt_bytes(warp_offset));

                        int thread_offset = /*0 + */lane_id * kNumSendUnrolls;
                        int num_bytes = logfmt_encode<kNumSendUnrolls>(
                            logfmt_buffers[0],
                            (thread_offset % kNumInt4PerDivision == 0) ? meta_buffers + thread_offset / kNumInt4PerDivision : nullptr,
                            lane_id
                        );
                        encoded_bytes[0] = num_bytes;
773
774
775
                    }
                    syncwarp();

lishen's avatar
lishen committed
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
                    // 采用3级流水
                    for (int iter_idx = 0; iter_idx < kNumIters; ++iter_idx) {
                        // 流水线S1: 加载第 (kNumStages-1) 轮之后的数据
                        const int stage_last_iter = iter_idx + kNumStages - 1;  // 当前iter所在stage中的最后一个,初始为S3的读取数据
                        if (stage_last_iter < kNumIters) {
                            int stage_idx = stage_last_iter % kNumStages;
                            int warp_offset = stage_last_iter * kNumMsgInt4ElemPerWarp;
                            logfmt_load_global2lds(stage_idx, cpy_src_int4_ptr + warp_offset, get_num_logfmt_bytes(warp_offset));
                        }

                        // 流水线S2: 处理下一轮的数据量化
                        const int stage_next_iter = iter_idx + 1;
                        if (stage_next_iter < kNumIters) {
                            int stage_idx = stage_next_iter % kNumStages;
                            int warp_offset = stage_next_iter * kNumMsgInt4ElemPerWarp;
                            int thread_offset = warp_offset + lane_id * kNumSendUnrolls;
                            int num_bytes = logfmt_encode<kNumSendUnrolls>(
                                logfmt_buffers[stage_idx],
                                (thread_offset % kNumInt4PerDivision == 0) ? meta_buffers + thread_offset / kNumInt4PerDivision : nullptr,
                                lane_id
                            );
                            encoded_bytes[stage_idx] = num_bytes;
                        }

                        // 流水线S3:当前轮进行数据拷贝到通信显存
                        if (iter_idx < kNumIters) {
                            int stage_idx = iter_idx % kNumStages;
                            using vec_type = uint64_t;
                            int nvecs = encoded_bytes[stage_idx] / sizeof(vec_type);
                            if (nvecs > 0) {
                                UNROLLED_WARP_COPY_LL(1, lane_id, nvecs,
                                    reinterpret_cast<vec_type*>(reinterpret_cast<uint8_t*>(cpy_dst_int4_ptr) + logfmt_offset_bytes),
                                    reinterpret_cast<vec_type*>(logfmt_buffers[stage_idx]),
                                    ld_direct_global, st_na_global);
                            }
                            logfmt_offset_bytes += encoded_bytes[stage_idx];
                        }

                        syncwarp();
                    }

817
818
                    num_send_bytes = logfmt_offset_bytes;

lishen's avatar
lishen committed
819
820
821
822
823
824
                    // Store metadata
                    using meta_vec_type = uint32_t;
                    UNROLLED_WARP_COPY_LL(1, lane_id, kNumMetaBytes / sizeof(meta_vec_type),
                        reinterpret_cast<meta_vec_type*>(cpy_dst_int4_ptr),
                        reinterpret_cast<meta_vec_type*>(meta_buffers),
                        ld_direct_global, st_na_global);
825

lishen's avatar
lishen committed
826
827
828
829
830
831
832
833
834
                } else {
                    // ===== 非 LogFMT 路径:直接 global -> global,不经过 LDS =====
                    for (int iter_idx = 0; iter_idx < kNumIters; ++iter_idx) {
                        int warp_offset = iter_idx * kNumMsgInt4ElemPerWarp;
                        UNROLLED_WARP_COPY_LL(kNumSendUnrolls, lane_id, kNumMsgInt4ElemPerWarp,
                            cpy_dst_int4_ptr + warp_offset,
                            cpy_src_int4_ptr + warp_offset,
                            ld_direct_global, st_na_global);
                        syncwarp();
835
                    }
lishen's avatar
lishen committed
836
837
                    // 非 LogFMT 时,发送字节数为原始大小
                    num_send_bytes = hidden_bf16_int4 * sizeof(int4); // 或根据实际计算
838
                }
lishen's avatar
lishen committed
839

840
841
                syncwarp();
            }
842

843
            if (dst_p2p_ptr == 0) {
lishen's avatar
lishen committed
844
                internode_ll_putmem_nbi((void*)dst_ptr, (void*)buf_ptr,
845
846
                                        num_ranks, dst_rank, local_expert_idx,
                                        num_send_bytes);
lishen's avatar
lishen committed
847
            }
lishen's avatar
lishen committed
848
        }
Chenggang Zhao's avatar
Chenggang Zhao committed
849

lishen's avatar
lishen committed
850
        // Put finishing flag
lishen's avatar
lishen committed
851
        // EP_DEVICE_ASSERT(num_warps_per_group > 1);
lishen's avatar
lishen committed
852
        if (lane_id == 0){
lishen's avatar
lishen committed
853
            volatile int ret = __hip_atomic_fetch_add(&sync_large_warp_counters[warp_group_id], 1,__ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_WORKGROUP);
lishen's avatar
lishen committed
854
855
        }
        syncwarp();
lishen's avatar
lishen committed
856
        while (sync_large_warp_counters[warp_group_id] < num_warps_per_group);
lishen's avatar
lishen committed
857

lishen's avatar
lishen committed
858
859
        if (sub_warp_id == 1 and lane_id == 0) {
            while (ld_acquire_global(atomic_clean_flag) == 0);
lishen's avatar
lishen committed
860
861
862
863
864
865
866
867

            auto dst_ptr = rdma_recv_flag + global_expert_idx;
            // 通过 shmem_get_p2p_ptr 获取 当前远程指针能否可达
            uint64_t p2p_ptr = internode::shmem_get_p2p_ptr((void*)dst_ptr, rank, dst_rank);
            if (p2p_ptr == 0) {  // RDMA
                internode_ll_long_atomic_add(dst_ptr, 1, num_ranks, dst_rank, local_expert_idx);
            } else { //  本地 GPU 和 同一计算节点的 其他 GPU 地址
                st_na_release(reinterpret_cast<int *>(p2p_ptr), 1);
lishen's avatar
lishen committed
868
            }
lishen's avatar
lishen committed
869

lishen's avatar
lishen committed
870
871
872
            atomic_add_release_global(atomic_clean_flag, -1);
        }
        syncwarp();
873
874
    }

lishen's avatar
lishen committed
875
    // Receiving phase
lishen's avatar
lishen committed
876
LOW_LATENCY_COMBINE_RECV:
lishen's avatar
lishen committed
877
878
    if ((phases & LOW_LATENCY_RECV_PHASE) == 0)
        return;
879

lishen's avatar
lishen committed
880
881
    // Wait all ranks to arrive and notify PCIe usage
    if (responsible_expert_idx < num_experts) {
lishen's avatar
lishen committed
882
        // EP_DEVICE_ASSERT(num_warps_per_group > 1);
lishen's avatar
lishen committed
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
        if (sub_warp_id == 0 and lane_id == 0) {
            const auto src_rank = responsible_expert_idx / num_local_experts;
            auto start_time = wall_clock64();
            uint64_t wait_recv_cost = 0;
            while (ld_acquire_global(reinterpret_cast<int*>(rdma_recv_flag + responsible_expert_idx)) == 0  // recv not ready
                   && (wait_recv_cost = wall_clock64() - start_time) <= NUM_TIMEOUT_CYCLES   // not timeout
            );

            // Mask rank if timeout
            if (wait_recv_cost > NUM_TIMEOUT_CYCLES) {
                printf("Warning: DeepEP timeout for combine receive, rank %d, local_expert_idx %d, src_rank %d\n",
                       rank, responsible_expert_idx % num_local_experts, src_rank);
            }

            if (combine_wait_recv_cost_stats != nullptr) {
                atomicAdd(reinterpret_cast<unsigned long long*>(combine_wait_recv_cost_stats + src_rank), wait_recv_cost);
            }
lishen's avatar
lishen committed
900
        }
901
    }
lishen's avatar
lishen committed
902
903
904
    grid_barrier(global_atomic_counter, num_sms);

    // Reduce tokens with FP8 cast
lishen's avatar
lishen committed
905
    // EP_DEVICE_ASSERT(num_topk <= kWarpSize and hidden_bf16_int4 <= num_threads);
lishen's avatar
lishen committed
906
    EP_STATIC_ASSERT(kHidden % (kWarpSize * kNumElemsPerInt4) == 0, "Invalid vectorization");
907

908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
    // 计算需要多少个warp
    constexpr int num_decode_warps = hidden_bf16_int4 / (kNumRecvUnrolls * kWarpSize);

    // 每128个数据记录一个max/min值,即该数为总的max/min值数量
    constexpr int kNumDivisionBytes = kNumDivisions * sizeof(float);
    // 每个warp内总的BF16值的数量
    constexpr int kNumBF16PerWarpBytes = kWarpSize * kNumRecvUnrolls * sizeof(int4);
    constexpr int kNumLogFMTPerWarpBytes = kNumBF16PerWarpBytes * 10 / 16;

    // 用于记录 max/min 值的 log 值
    auto log_amax_buffers =
        PatternVisitor([=](const int& i) { return reinterpret_cast<float*>(smem_buffer + i * kNumDivisionBytes); });
    auto log_amin_buffers = PatternVisitor([=](const int& i) {
      return reinterpret_cast<float*>(smem_buffer + kNumStages * kNumDivisionBytes + i * kNumDivisionBytes);
    });
    auto cast_info_buffers = PatternVisitor([=](const int& i) {
      return reinterpret_cast<int*>(smem_buffer + kNumStages * kNumDivisionBytes * 2 + i * kNumDivisionBytes);
    });

    // 初始化 topk_idx 和 topk_weights
    int topk_idx_by_lane = -1;
    float topk_weights_by_lane = -1;
    int stage_idx = 0;
    for (int token_idx = sm_id; token_idx < num_combined_tokens; token_idx += num_sms) {
        if (lane_id < num_topk) {
            topk_idx_by_lane = static_cast<int>(__ldg(topk_idx + token_idx * num_topk + lane_id));
            topk_weights_by_lane = __ldg(topk_weights + token_idx * num_topk + lane_id);
        }

lishen's avatar
lishen committed
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
        for (int w_i = warp_id; w_i < num_decode_warps; w_i += num_warps) {
            float combined_values[kNumElemsPerInt4 * kNumRecvUnrolls] = {0.0f};
            #pragma unroll
            for (int i = 0; i < num_topk; ++ i) {
                int topk_idx_reg = shfl_sync(topk_idx_by_lane, i);
                if (topk_idx_reg < 0)
                    continue;
                const auto& topk_weight_reg = shfl_sync(topk_weights_by_lane, i);

                // Read from sources
                auto rdma_buffer_type = reinterpret_cast<const uint8_t*>(reinterpret_cast<uint8_t*>(rdma_recv_x) +
                    (topk_idx_reg * num_max_dispatch_tokens_per_rank + token_idx) * num_bytes_per_slot);

                if constexpr(bSupportLogFMT) {
                    // 接收到的数据位置
                    const uint8_t* data_buffer = rdma_buffer_type + kNumMetaBytes;

                    // 读取max/min数据
                    if(w_i == 0) {
                        // 因为每个warp能处理数据量为 kWarpSize*sizeof(int4)/sizeof(bfloat16) * kNumSendUnrolls
                        // 即不考虑kNumSendUnrolls,一共 kWarpSize*sizeof(int4)/sizeof(bfloat16)/128 组, 代入参数 = kWarpSize / 16 个warp,nv上为2,dcu上为4
                        logfmt_check_amaxmin<kNumDivisions / (kWarpSize / 16), kNumSendUnrolls, kNumRecvUnrolls>(
                            /*meta_buffer*/rdma_buffer_type,
                            reinterpret_cast<int4*>(log_amax_buffers[stage_idx]),
                            reinterpret_cast<int4*>(log_amin_buffers[stage_idx]),
                            cast_info_buffers[stage_idx],
                            lane_id);
                    }
965

lishen's avatar
lishen committed
966
                    __syncthreads();
967

lishen's avatar
lishen committed
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
1001
                    // 获取cast_info_buffers
                    const auto& info = cast_info_buffers[stage_idx][w_i];
                    bool enable_cast = info & 1;
                    int num_casted_prefix = info >> 1; // 可用的

                    // 计算偏移(与TMA版本逻辑一致)
                    int warp_offset = kNumLogFMTPerWarpBytes * num_casted_prefix +
                                      kNumBF16PerWarpBytes * (w_i - num_casted_prefix);
                    int lane_offset = (enable_cast ? kNumLogFMTPerWarpBytes : kNumBF16PerWarpBytes) / kWarpSize * lane_id;

                    // 使用临时缓冲区进行归约
                    const uint8_t* thread_data_ptr = data_buffer + warp_offset + lane_offset;

                    /**
                    一共有kNumDivisions个max/min数据对,读取时每warp默认处理256bit的max/min,所以logfmt_check_amaxmin的kNumLanes设置为 kNumDivisions/2
                    保存数据时每个log_amax_buffers为float2数据类型,保存总的warpkNumDivisions / 2
                    实际保存数据时,每个warp保存的实际数据个数为 kWarpSize*kNumRecvUnrolls*sizeof(int4)/sizeof(hip_bfloat16)
                    实际每个warp读取的max/min的 warp_idx=kWarpSize*kNumRecvUnrolls*sizeof(int4)/sizeof(hip_bfloat16) / 128 = kNumRecvUnrolls * 2
                    具体的lane_id处理的数据量为 warp_idx / kWarpSize
                    */
                    int log_amaxmin_per_warp = kNumRecvUnrolls * kWarpSize * sizeof(int4) / sizeof(hip_bfloat16) / QUANTIZATION_GROUPSIZE;
                    int division_idx = w_i * log_amaxmin_per_warp + lane_id * log_amaxmin_per_warp / kWarpSize;

                    // 反量化
                    decode_and_accumulate<kNumRecvUnrolls>(
                        reinterpret_cast<const uint32_t*>(thread_data_ptr),  // 直接使用全局内存地址
                        combined_values,
                        log_amax_buffers[stage_idx][division_idx],
                        log_amin_buffers[stage_idx][division_idx],
                        enable_cast,
                        topk_weight_reg);
                } else {
                    // 接收到的数据位置
                    const uint8_t* data_buffer = rdma_buffer_type;
1002

lishen's avatar
lishen committed
1003
1004
1005
1006
1007
                    // 计算偏移
                    int warp_offset = kNumBF16PerWarpBytes * w_i;
                    int lane_offset = kNumBF16PerWarpBytes / kWarpSize * lane_id;
                    // 使用临时缓冲区进行归约
                    const uint8_t* thread_data_ptr = data_buffer + warp_offset + lane_offset;
1008
1009

                    #pragma unroll
lishen's avatar
lishen committed
1010
1011
1012
1013
1014
1015
1016
1017
1018
                    for (int j = 0; j < kNumRecvUnrolls; ++j) {
                        auto tmp_rdma_value = ld_nc_global(reinterpret_cast<const int4*>(thread_data_ptr) + j);
                        const auto x_bf16 = reinterpret_cast<const hip_bfloat16*>(&tmp_rdma_value);

                        #pragma unroll
                        for (int k = 0; k < kNumElemsPerInt4; ++k) {
                            int combined_idx = j * kNumElemsPerInt4 + k;
                            combined_values[combined_idx] += static_cast<float>(x_bf16[k]) * topk_weight_reg;
                        }
1019
1020
                    }
                }
lishen's avatar
lishen committed
1021
            }
1022

lishen's avatar
lishen committed
1023
1024
1025
1026
1027
1028
1029
            // Write results,kNumRecvUnrolls==2时则写256bit的数
            int4 combined_int4[kNumRecvUnrolls];
            auto combined_bf16 = reinterpret_cast<hip_bfloat16 *>(&combined_int4[0]);
            #pragma unroll
            for (int j = 0; j < kNumElemsPerInt4 * kNumRecvUnrolls; ++ j) {
                combined_bf16[j] = static_cast<hip_bfloat16>(combined_values[j]);
            }
1030

lishen's avatar
lishen committed
1031
1032
1033
1034
            for(int j = 0; j < kNumRecvUnrolls; ++ j) {
                (reinterpret_cast<int4*>(combined_x) + token_idx * hidden_bf16_int4 +
                w_i * kWarpSize * kNumRecvUnrolls)[lane_id * kNumRecvUnrolls + j] = combined_int4[j];
            }
lishen's avatar
lishen committed
1035
1036
        }
    }
1037
1038
}

lishen's avatar
lishen committed
1039
1040
1041
1042
1043
void combine(void* combined_x,
             void* rdma_recv_x, int64_t* rdma_recv_flag, void* rdma_send_x,
             const void* x, const int64_t* topk_idx, const float* topk_weights,
             const int* src_info, const int64_t* layout_range,
             int* global_atomic_counter,
lishen's avatar
lishen committed
1044
             int64_t* combine_wait_recv_cost_stats,
lishen's avatar
lishen committed
1045
1046
1047
             int64_t* next_clean, int num_next_clean_int,
             int num_combined_tokens, int hidden, int num_max_dispatch_tokens_per_rank,
             int num_topk, int num_experts, int rank, int num_ranks,
1048
             bool use_logfmt,
lishen's avatar
lishen committed
1049
             void* workspace, int num_device_sms, hipStream_t stream,
lishen's avatar
lishen committed
1050
             int phases, bool zero_copy) {
lishen's avatar
lishen committed
1051
    constexpr int kMaxNumWarps = 8;
lishen's avatar
lishen committed
1052
1053
    constexpr int kNumMaxTopk = 11;
    const int num_warp_groups = ceil_div(num_experts, num_device_sms);
1054
    const int num_warps_per_group = kMaxNumWarps / num_warp_groups; // num_warps_per_group>1, "Requires more than one warp per group"
lishen's avatar
lishen committed
1055
1056
    const int num_recv_per_sm = ceil_div(num_combined_tokens, num_device_sms);
    EP_HOST_ASSERT(num_warp_groups > 0 and num_warps_per_group > 0 and num_recv_per_sm >= 0);
lishen's avatar
lishen committed
1057

lishen's avatar
lishen committed
1058
1059
1060
    const auto num_warps = num_warp_groups * num_warps_per_group;
    const auto num_sms =
        max(ceil_div(num_experts, num_warp_groups), num_recv_per_sm == 0 ? 1 : ceil_div(num_combined_tokens, num_recv_per_sm));
lishen's avatar
lishen committed
1061
1062
1063
1064
1065
1066

    // Check workspace
    auto atomic_clean_flag = reinterpret_cast<int*>(workspace);
    EP_HOST_ASSERT(sizeof(int) <= NUM_WORKSPACE_BYTES);
    EP_HOST_ASSERT(num_topk <= kNumMaxTopk);

1067
1068
#define COMBINE_LAUNCH_CASE(hidden)                                            \
  {                                                                            \
1069
1070
1071
    auto combine_func = use_logfmt ?                                           \
         combine<true, hidden, kNumMaxTopk, kMaxNumWarps> :                    \
         combine<false, hidden, kNumMaxTopk, kMaxNumWarps>;                    \
1072
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
    LAUNCH_KERNEL_NON_COOPERATIVE(&cfg, combine_func,                          \
        combined_x, rdma_recv_x, rdma_recv_flag, rdma_send_x,                  \
        x, topk_idx, topk_weights, src_info, layout_range,                     \
        global_atomic_counter, combine_wait_recv_cost_stats,                   \
        next_clean, num_next_clean_int,                                        \
        atomic_clean_flag, num_combined_tokens, hidden,                        \
        num_topk, num_max_dispatch_tokens_per_rank,                            \
        num_experts, rank, num_ranks,                                          \
        num_warp_groups, num_warps_per_group, phases, zero_copy);              \
  }                                                                            \
  break
lishen's avatar
lishen committed
1083
1084
1085
1086

    SETUP_LAUNCH_CONFIG(num_sms, num_warps * kWarpSize, stream);
    SWITCH_HIDDEN(COMBINE_LAUNCH_CASE);
#undef COMBINE_LAUNCH_CASE
1087
1088
}

Chenggang Zhao's avatar
Chenggang Zhao committed
1089
1090
1091
} // namespace internode_ll

} // namespace deep_ep