Skip to content

Commit f5748fd

Browse files
xtyao66xtyao66
committed
examples/ep: support configurable NVL group size
Signed-off-by: xt66 <60164575+xtyao66@users.noreply.github.com>
1 parent b775042 commit f5748fd

11 files changed

Lines changed: 494 additions & 193 deletions

File tree

‎examples/device/ep/README.md‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -38,7 +38,7 @@ buffer.disconnect_ranks(ranks)
3838

3939
## Key APIs
4040

41-
- `Buffer(rank_id, ...)`: Initialize the NIXL communication buffer
41+
- `Buffer(rank_id, ..., nvl_group_size=8)`: Initialize the NIXL communication buffer; `nvl_group_size` must satisfy `0 < nvl_group_size <= 8` and divide `8`
4242
- `update_memory_buffers(num_ranks, num_experts_per_rank, num_rdma_bytes, num_nvl_bytes=0)`: Prepare buffers for up to `num_ranks` ranks and `num_experts_per_rank` experts
4343
- `connect_ranks(remote_ranks, activate=True)`: Establish NIXL connections to new peers (can be called multiple times); in low-latency mode, use `activate=False` to keep new peers masked until explicitly unmasked.
4444
- `disconnect_ranks(remote_ranks)`: Clean up connections to departing peers

‎examples/device/ep/csrc/config.hpp‎

Lines changed: 20 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -64,40 +64,45 @@ struct Config {
6464
EP_HOST_ASSERT(num_max_rdma_chunked_send_tokens <= num_max_rdma_chunked_recv_tokens / 2);
6565
}
6666

67-
size_t get_nvl_buffer_size_hint(size_t hidden_bytes, int num_ranks) const {
67+
size_t get_nvl_buffer_size_hint(size_t hidden_bytes, int num_ranks, int nvl_group_size = NUM_MAX_NVL_PEERS) const {
6868
// Below are some assumptions
6969
// TODO: add assertions
7070
constexpr int kNumMaxTopK = 128;
7171
constexpr int kNumMaxScales = 128;
72-
EP_HOST_ASSERT(num_ranks < NUM_MAX_NVL_PEERS or num_ranks % NUM_MAX_NVL_PEERS == 0);
73-
EP_HOST_ASSERT(num_ranks <= NUM_MAX_NVL_PEERS or num_sms % 2 == 0);
74-
const auto num_rdma_ranks = std::max(num_ranks / NUM_MAX_NVL_PEERS, 1);
75-
const auto num_nvl_ranks = std::min(num_ranks, NUM_MAX_NVL_PEERS);
72+
EP_HOST_ASSERT(nvl_group_size > 0 and nvl_group_size <= NUM_MAX_NVL_PEERS and NUM_MAX_NVL_PEERS % nvl_group_size == 0);
73+
EP_HOST_ASSERT(num_ranks < nvl_group_size or num_ranks % nvl_group_size == 0);
74+
EP_HOST_ASSERT(num_ranks <= nvl_group_size or num_sms % 2 == 0);
75+
const auto num_rdma_ranks = std::max(num_ranks / nvl_group_size, 1);
76+
// HT kernels keep fixed-size NVL scratch layouts with NUM_MAX_NVL_PEERS
77+
// as the stride. nvl_group_size controls active peer mapping, not the
78+
// maximum scratch lane count needed by the buffer layout.
79+
const auto num_nvl_buffer_lanes = std::min(num_ranks, NUM_MAX_NVL_PEERS);
7680
const int num_channels = num_sms / 2;
7781

7882
size_t num_bytes = 0;
79-
num_bytes += num_channels * num_nvl_ranks * (2 * num_rdma_ranks + 3) * sizeof(int);
80-
num_bytes += num_channels * num_nvl_ranks * num_max_nvl_chunked_recv_tokens * hidden_bytes;
81-
num_bytes += num_channels * num_nvl_ranks * num_max_nvl_chunked_recv_tokens * ht::get_source_meta_bytes();
82-
num_bytes += num_channels * num_nvl_ranks * num_max_nvl_chunked_recv_tokens * kNumMaxTopK * sizeof(int64_t);
83-
num_bytes += num_channels * num_nvl_ranks * num_max_nvl_chunked_recv_tokens * kNumMaxTopK * sizeof(float);
84-
num_bytes += num_channels * num_nvl_ranks * num_max_nvl_chunked_recv_tokens * kNumMaxScales * sizeof(float);
83+
num_bytes += num_channels * num_nvl_buffer_lanes * (2 * num_rdma_ranks + 3) * sizeof(int);
84+
num_bytes += num_channels * num_nvl_buffer_lanes * num_max_nvl_chunked_recv_tokens * hidden_bytes;
85+
num_bytes += num_channels * num_nvl_buffer_lanes * num_max_nvl_chunked_recv_tokens * ht::get_source_meta_bytes();
86+
num_bytes += num_channels * num_nvl_buffer_lanes * num_max_nvl_chunked_recv_tokens * kNumMaxTopK * sizeof(int64_t);
87+
num_bytes += num_channels * num_nvl_buffer_lanes * num_max_nvl_chunked_recv_tokens * kNumMaxTopK * sizeof(float);
88+
num_bytes += num_channels * num_nvl_buffer_lanes * num_max_nvl_chunked_recv_tokens * kNumMaxScales * sizeof(float);
8589
num_bytes = ((num_bytes + 127) / 128) * 128;
8690
return num_bytes;
8791
}
8892

89-
size_t get_rdma_buffer_size_hint(int64_t hidden_bytes, int num_ranks) const {
93+
size_t get_rdma_buffer_size_hint(int64_t hidden_bytes, int num_ranks, int nvl_group_size = NUM_MAX_NVL_PEERS) const {
9094
// Legacy mode
91-
if (num_ranks <= NUM_MAX_NVL_PEERS)
95+
EP_HOST_ASSERT(nvl_group_size > 0 and nvl_group_size <= NUM_MAX_NVL_PEERS and NUM_MAX_NVL_PEERS % nvl_group_size == 0);
96+
if (num_ranks <= nvl_group_size)
9297
return 0;
9398

9499
// Below are some assumptions
95100
// TODO: add assertions
96101
constexpr int kNumMaxTopK = 128;
97102
constexpr int kNumMaxScales = 128;
98-
EP_HOST_ASSERT(num_ranks % NUM_MAX_NVL_PEERS == 0);
103+
EP_HOST_ASSERT(num_ranks % nvl_group_size == 0);
99104
EP_HOST_ASSERT(num_sms % 2 == 0);
100-
const int num_rdma_ranks = num_ranks / NUM_MAX_NVL_PEERS;
105+
const int num_rdma_ranks = num_ranks / nvl_group_size;
101106
const int num_channels = num_sms / 2;
102107

103108
size_t num_bytes = 0;

‎examples/device/ep/csrc/kernels/api.cuh‎

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,21 +41,59 @@ struct gpu_nixl_ctx {
4141
nixlMemViewH barrier_mvh;
4242
nixlMemViewH remote_mvh;
4343
nixlMemViewH ht_barrier_mvh;
44+
nixlMemViewH *remote_mvh_by_rank; // [global_rank] -> one-element remote memview for RDMA buffer
45+
nixlMemViewH *barrier_mvh_by_rank; // [global_rank] -> one-element remote memview for sync/barrier buffer
46+
int *remote_mvh_indices; // [global_rank] -> compact remote/barrier memview index, -1 when disconnected
4447
int *sync_buffer_ptr; // [src_rank]
4548
int *sync_count_ptr; // [dst_rank]
4649
uint64_t *last_ht_barrier_counter;
4750
uint64_t *local_ht_barrier_counter_ptr;
4851
void *rdma_buffer_ptr;
52+
uint64_t ht_barrier_counter_offset;
4953
void **p2p_ptrs;
5054
int max_num_ranks;
5155
int num_rdma_ranks;
56+
int num_nvl_ranks;
57+
int nvl_group_size;
5258
int rank;
5359

5460
__device__ inline uint64_t offset_get(uint64_t ptr) {
5561
return ptr - reinterpret_cast<uint64_t>(rdma_buffer_ptr);
5662
}
5763
};
5864

65+
__device__ __forceinline__ nixlMemViewH remote_mvh_for_rank(const gpu_nixl_ctx& ctx, int global_rank) {
66+
EP_DEVICE_ASSERT(global_rank >= 0 and global_rank < ctx.max_num_ranks);
67+
if (ctx.remote_mvh_by_rank != nullptr) {
68+
auto mvh = ctx.remote_mvh_by_rank[global_rank];
69+
EP_DEVICE_ASSERT(mvh != nullptr);
70+
return mvh;
71+
}
72+
EP_DEVICE_ASSERT(ctx.remote_mvh != nullptr);
73+
return ctx.remote_mvh;
74+
}
75+
76+
__device__ __forceinline__ nixlMemViewH barrier_mvh_for_rank(const gpu_nixl_ctx& ctx, int global_rank) {
77+
EP_DEVICE_ASSERT(global_rank >= 0 and global_rank < ctx.max_num_ranks);
78+
if (ctx.barrier_mvh_by_rank != nullptr) {
79+
auto mvh = ctx.barrier_mvh_by_rank[global_rank];
80+
EP_DEVICE_ASSERT(mvh != nullptr);
81+
return mvh;
82+
}
83+
EP_DEVICE_ASSERT(ctx.barrier_mvh != nullptr);
84+
return ctx.barrier_mvh;
85+
}
86+
87+
__device__ __forceinline__ size_t remote_mvh_index(const gpu_nixl_ctx& ctx, int global_rank) {
88+
if (ctx.remote_mvh_by_rank != nullptr) {
89+
return 0;
90+
}
91+
EP_DEVICE_ASSERT(ctx.remote_mvh_indices != nullptr);
92+
EP_DEVICE_ASSERT(global_rank >= 0 and global_rank < ctx.max_num_ranks);
93+
const auto compact_index = ctx.remote_mvh_indices[global_rank];
94+
EP_DEVICE_ASSERT(compact_index >= 0);
95+
return static_cast<size_t>(compact_index);
96+
}
5997

6098
// Layout kernels
6199
namespace layout {
@@ -69,6 +107,7 @@ void get_dispatch_layout(const topk_idx_t* topk_idx,
69107
int num_topk,
70108
int num_ranks,
71109
int num_experts,
110+
int nvl_group_size,
72111
cudaStream_t stream);
73112

74113
} // namespace layout

‎examples/device/ep/csrc/kernels/launch.cuh‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -75,6 +75,8 @@ cfg.dynamicSmemBytes = smem_size;
7575

7676
#define SWITCH_NVL_RANKS(case_macro) \
7777
switch (num_nvl_ranks) { \
78+
case 1: \
79+
case_macro(1); \
7880
case 2: \
7981
case_macro(2); \
8082
case 4: \
@@ -87,9 +89,12 @@ cfg.dynamicSmemBytes = smem_size;
8789
while (false)
8890

8991
#define SWITCH_RDMA_RANKS(case_macro) \
90-
switch (num_ranks / NUM_MAX_NVL_PEERS) { \
92+
switch (num_rdma_ranks) { \
9193
case 2: case_macro(2); \
94+
case 3: case_macro(3); \
9295
case 4: case_macro(4); \
96+
case 5: case_macro(5); \
97+
case 6: case_macro(6); \
9398
case 8: case_macro(8); \
9499
case 16: case_macro(16); \
95100
case 18: case_macro(18); \

‎examples/device/ep/csrc/kernels/layout.cu‎

Lines changed: 13 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,8 @@ template <int kNumThreads, int kNumExpertsPerSM, int kNumRanksPerSM>
3232
__global__ void get_dispatch_layout(const topk_idx_t* topk_idx,
3333
int* num_tokens_per_rank, int* num_tokens_per_rdma_rank,
3434
int* num_tokens_per_expert, bool* is_token_in_rank,
35-
int num_tokens, int num_topk, int num_ranks, int num_experts) {
35+
int num_tokens, int num_topk, int num_ranks, int num_experts,
36+
int nvl_group_size) {
3637
auto sm_id = static_cast<int>(blockIdx.x);
3738
auto thread_id = static_cast<int>(threadIdx.x);
3839

@@ -69,15 +70,18 @@ __global__ void get_dispatch_layout(const topk_idx_t* topk_idx,
6970
}
7071

7172
if (num_tokens_per_rdma_rank != nullptr)
72-
EP_DEVICE_ASSERT(num_ranks % NUM_MAX_NVL_PEERS == 0 and num_ranks > NUM_MAX_NVL_PEERS);
73+
EP_DEVICE_ASSERT(nvl_group_size > 0 and nvl_group_size <= NUM_MAX_NVL_PEERS and
74+
kNumRanksPerSM % nvl_group_size == 0 and
75+
num_ranks % nvl_group_size == 0 and num_ranks > nvl_group_size);
7376

7477
// Count rank statistics
75-
constexpr int kNumRDMARanksPerSM = kNumRanksPerSM / NUM_MAX_NVL_PEERS;
78+
constexpr int kNumRDMARanksPerSM = kNumRanksPerSM;
7679
__shared__ int num_tokens_per_rank_per_thread[kNumThreads][kNumRanksPerSM];
7780
__shared__ int num_tokens_per_rdma_rank_per_thread[kNumThreads][kNumRDMARanksPerSM];
7881
auto sm_begin = (num_experts + kNumExpertsPerSM - 1) / kNumExpertsPerSM;
7982
int rank_begin_idx = (sm_id - sm_begin) * kNumRanksPerSM, rank_end_idx = min(rank_begin_idx + kNumRanksPerSM, num_ranks);
80-
int rdma_rank_begin_idx = rank_begin_idx / NUM_MAX_NVL_PEERS, rdma_rank_end_idx = rank_end_idx / NUM_MAX_NVL_PEERS;
83+
int rdma_rank_begin_idx = rank_begin_idx / nvl_group_size;
84+
int rdma_rank_end_idx = (rank_end_idx + nvl_group_size - 1) / nvl_group_size;
8185
if (rank_begin_idx < rank_end_idx) {
8286
const auto num_expert_per_rank = num_experts / num_ranks;
8387
auto expert_begin = rank_begin_idx * num_expert_per_rank;
@@ -100,7 +104,7 @@ __global__ void get_dispatch_layout(const topk_idx_t* topk_idx,
100104
if (expert_begin <= expert_idx and expert_idx < expert_end) {
101105
// Count single rank
102106
rank_idx = expert_idx / num_expert_per_rank - rank_begin_idx;
103-
is_in_rank[rank_idx] ++, is_in_rdma_rank[rank_idx / NUM_MAX_NVL_PEERS] ++;
107+
is_in_rank[rank_idx] ++, is_in_rdma_rank[rank_idx / nvl_group_size] ++;
104108
}
105109
}
106110

@@ -141,15 +145,16 @@ void get_dispatch_layout(const topk_idx_t* topk_idx,
141145
int* num_tokens_per_rank, int* num_tokens_per_rdma_rank,
142146
int* num_tokens_per_expert, bool* is_token_in_rank,
143147
int num_tokens, int num_topk, int num_ranks, int num_experts,
144-
cudaStream_t stream) {
148+
int nvl_group_size, cudaStream_t stream) {
145149
constexpr int kNumThreads = 256, kNumExpertsPerSM = 4, kNumRanksPerSM = 8;
146150
int num_sms = ((num_experts + kNumExpertsPerSM - 1) / kNumExpertsPerSM) + (num_ranks + kNumRanksPerSM - 1) / kNumRanksPerSM;
147-
EP_STATIC_ASSERT(kNumRanksPerSM % NUM_MAX_NVL_PEERS == 0, "Invalid number of ranks per SM");
151+
EP_HOST_ASSERT(nvl_group_size > 0 and nvl_group_size <= NUM_MAX_NVL_PEERS and
152+
kNumRanksPerSM % nvl_group_size == 0);
148153

149154
SETUP_LAUNCH_CONFIG(num_sms, kNumThreads, stream);
150155
LAUNCH_KERNEL(&cfg, (get_dispatch_layout<kNumThreads, kNumExpertsPerSM, kNumRanksPerSM>),
151156
topk_idx, num_tokens_per_rank, num_tokens_per_rdma_rank, num_tokens_per_expert, is_token_in_rank,
152-
num_tokens, num_topk, num_ranks, num_experts);
157+
num_tokens, num_topk, num_ranks, num_experts, nvl_group_size);
153158
}
154159

155160
} // namespace layout

0 commit comments

Comments
 (0)