style: fix sm90 megamoe dispatch indentation
This commit is contained in:
Xinyi Liu
@@ -245,215 +245,215 @@ sm90_fp8_mega_moe_impl(void* y,
|
||||
kDispatchGridSyncIndex, kAfterWorkspaceCleanBarrierTag>(
|
||||
workspace, sym_buffer, sm_idx, thread_idx, dispatch_sync);
|
||||
|
||||
for (uint32_t i = thread_idx; i < kNumExperts; i += kNumDispatchThreads)
|
||||
smem_expert_count[i] = 0;
|
||||
ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx);
|
||||
for (uint32_t i = thread_idx; i < kNumExperts; i += kNumDispatchThreads)
|
||||
smem_expert_count[i] = 0;
|
||||
ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx);
|
||||
|
||||
const auto read_topk_idx = [&](const auto& process) {
|
||||
#pragma unroll 1
|
||||
for (uint32_t i = (sm_idx * kNumDispatchWarps + warp_idx) * kNumTokensPerWarp;
|
||||
i < num_tokens;
|
||||
i += kNumSMs * kNumDispatchWarps * kNumTokensPerWarp) {
|
||||
const uint32_t token_offset = lane_idx / kNumTopk;
|
||||
const uint32_t token_idx = i + token_offset;
|
||||
int expert_idx = -1;
|
||||
if (token_idx < num_tokens and lane_idx < kNumActivateLanes) {
|
||||
expert_idx = static_cast<int>(
|
||||
__ldg(input_topk_idx_buffer.get_base_ptr<int64_t>() + i * kNumTopk + lane_idx));
|
||||
if (expert_idx >= 0 and expert_idx < static_cast<int>(kNumExperts))
|
||||
process(i * kNumTopk + lane_idx, static_cast<uint32_t>(expert_idx));
|
||||
const auto read_topk_idx = [&](const auto& process) {
|
||||
#pragma unroll 1
|
||||
for (uint32_t i = (sm_idx * kNumDispatchWarps + warp_idx) * kNumTokensPerWarp;
|
||||
i < num_tokens;
|
||||
i += kNumSMs * kNumDispatchWarps * kNumTokensPerWarp) {
|
||||
const uint32_t token_offset = lane_idx / kNumTopk;
|
||||
const uint32_t token_idx = i + token_offset;
|
||||
int expert_idx = -1;
|
||||
if (token_idx < num_tokens and lane_idx < kNumActivateLanes) {
|
||||
expert_idx = static_cast<int>(
|
||||
__ldg(input_topk_idx_buffer.get_base_ptr<int64_t>() + i * kNumTopk + lane_idx));
|
||||
if (expert_idx >= 0 and expert_idx < static_cast<int>(kNumExperts))
|
||||
process(i * kNumTopk + lane_idx, static_cast<uint32_t>(expert_idx));
|
||||
}
|
||||
__syncwarp();
|
||||
}
|
||||
};
|
||||
|
||||
// Count local outgoing token-topk entries per global expert.
|
||||
read_topk_idx([&](const uint32_t&, const uint32_t& expert_idx) {
|
||||
atomicAdd_block(smem_expert_count + expert_idx, 1);
|
||||
});
|
||||
ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx);
|
||||
|
||||
// Convert per-SM counts into global per-rank offsets. High 32 bits count arrived SMs.
|
||||
for (uint32_t i = thread_idx; i < kNumExperts; i += kNumDispatchThreads) {
|
||||
const uint64_t send_value = (1ull << 32) | static_cast<uint64_t>(smem_expert_count[i]);
|
||||
smem_expert_count[i] = static_cast<uint32_t>(
|
||||
ptx::atomic_add(workspace.get_expert_send_count_ptr(i), send_value));
|
||||
}
|
||||
ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx);
|
||||
|
||||
// Write source token-topk indices into destination ranks' local-expert tables.
|
||||
read_topk_idx([&](const uint32_t& token_topk_idx, const uint32_t& expert_idx) {
|
||||
const uint32_t dst_rank_idx = expert_idx / kNumExpertsPerRank;
|
||||
const uint32_t dst_local_expert_idx = expert_idx - dst_rank_idx * kNumExpertsPerRank;
|
||||
const uint32_t dst_slot_idx = atomicAdd_block(smem_expert_count + expert_idx, 1);
|
||||
const auto dst_ptr = workspace.get_src_token_topk_idx_ptr(
|
||||
dst_local_expert_idx, sym_buffer.rank_idx, dst_slot_idx);
|
||||
*sym_buffer.map(dst_ptr, dst_rank_idx) = token_topk_idx;
|
||||
});
|
||||
|
||||
// Wait until all local SMs have finished filling local send-count/source-index data.
|
||||
comm::grid_sync<kNumSMs, kDispatchGridSyncIndex>(
|
||||
workspace, sm_idx, thread_idx, dispatch_sync);
|
||||
|
||||
// Publish per-rank expert counts and summed ready/count words to destination ranks.
|
||||
if (sm_idx == 0) {
|
||||
for (uint32_t i = thread_idx; i < kNumExperts; i += kNumDispatchThreads) {
|
||||
const uint32_t dst_rank_idx = i / kNumExpertsPerRank;
|
||||
const uint32_t dst_local_expert_idx = i - dst_rank_idx * kNumExpertsPerRank;
|
||||
const uint64_t expert_status = *workspace.get_expert_send_count_ptr(i);
|
||||
*sym_buffer.map(
|
||||
workspace.get_expert_recv_count_ptr(sym_buffer.rank_idx, dst_local_expert_idx),
|
||||
dst_rank_idx) = expert_status & 0xffffffffu;
|
||||
ptx::atomic_add_sys(
|
||||
sym_buffer.map(workspace.get_expert_recv_count_sum_ptr(dst_local_expert_idx), dst_rank_idx),
|
||||
expert_status);
|
||||
}
|
||||
}
|
||||
ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx);
|
||||
|
||||
// All ranks must observe the complete count/source-index tables before pulling.
|
||||
comm::nvlink_barrier<kNumRanks, kNumSMs, kNumDispatchThreads,
|
||||
kDispatchGridSyncIndex, kBeforeDispatchPullBarrierTag>(
|
||||
workspace, sym_buffer, sm_idx, thread_idx, dispatch_sync,
|
||||
/* sync_prologue */ false, /* sync_epilogue */ true);
|
||||
|
||||
auto scheduler = sched::SM90MegaMoEScheduler<
|
||||
BLOCK_M, BLOCK_N, BLOCK_K,
|
||||
L1_SHAPE_N, L1_SHAPE_K,
|
||||
L2_SHAPE_N, L2_SHAPE_K,
|
||||
kNumExpertsPerRank,
|
||||
kNumExpertsPerWave,
|
||||
kNumSMs, kNumRanks,
|
||||
kUseNMajorL2>(workspace);
|
||||
scheduler.fetch_expert_recv_count();
|
||||
|
||||
if (sm_idx == 0 and cumulative_local_expert_recv_stats != nullptr) {
|
||||
for (uint32_t i = thread_idx; i < kNumExpertsPerRank; i += kNumDispatchThreads) {
|
||||
const auto num_recv_tokens = static_cast<uint32_t>(*workspace.get_expert_recv_count_sum_ptr(i));
|
||||
ptx::red_add(cumulative_local_expert_recv_stats + i, static_cast<int>(num_recv_tokens));
|
||||
}
|
||||
}
|
||||
ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx);
|
||||
|
||||
constexpr uint32_t kNumRanksPerLane = math::constexpr_ceil_div(kNumRanks, 32u);
|
||||
int current_expert_idx = -1;
|
||||
uint32_t stored_rank_count[kNumRanksPerLane] = {};
|
||||
uint32_t expert_start_idx = 0, expert_end_idx = 0;
|
||||
uint32_t expert_pool_block_offset = 0;
|
||||
|
||||
constexpr uint32_t kNumGlobalDispatchWarps = kNumSMs * kNumDispatchWarps;
|
||||
for (uint32_t token_idx = sm_idx * kNumDispatchWarps + warp_idx; ; token_idx += kNumGlobalDispatchWarps) {
|
||||
int old_expert_idx = current_expert_idx;
|
||||
while (token_idx >= expert_end_idx) {
|
||||
if (++ current_expert_idx >= static_cast<int>(kNumExpertsPerRank))
|
||||
break;
|
||||
|
||||
expert_pool_block_offset += math::ceil_div(expert_end_idx - expert_start_idx, BLOCK_M);
|
||||
expert_start_idx = expert_end_idx;
|
||||
expert_end_idx += scheduler.get_num_tokens(static_cast<uint32_t>(current_expert_idx));
|
||||
}
|
||||
|
||||
if (current_expert_idx >= static_cast<int>(kNumExpertsPerRank))
|
||||
break;
|
||||
|
||||
if (old_expert_idx != current_expert_idx) {
|
||||
old_expert_idx = current_expert_idx;
|
||||
#pragma unroll
|
||||
for (uint32_t i = 0; i < kNumRanksPerLane; ++ i) {
|
||||
const uint32_t j = i * 32 + lane_idx;
|
||||
stored_rank_count[i] = j < kNumRanks ?
|
||||
static_cast<uint32_t>(*workspace.get_expert_recv_count_ptr(j, static_cast<uint32_t>(current_expert_idx))) : 0;
|
||||
}
|
||||
}
|
||||
|
||||
uint32_t current_rank_in_expert_idx = 0;
|
||||
uint32_t remaining[kNumRanksPerLane];
|
||||
#pragma unroll
|
||||
for (uint32_t i = 0; i < kNumRanksPerLane; ++ i)
|
||||
remaining[i] = stored_rank_count[i];
|
||||
|
||||
uint32_t offset = 0;
|
||||
const uint32_t token_idx_in_expert = token_idx - expert_start_idx;
|
||||
uint32_t slot_idx = token_idx_in_expert;
|
||||
uint32_t token_idx_in_rank = 0;
|
||||
while (true) {
|
||||
uint32_t num_actives_in_lane = 0;
|
||||
uint32_t min_in_lane = 0xffffffffu;
|
||||
#pragma unroll
|
||||
for (uint32_t i = 0; i < kNumRanksPerLane; ++ i) {
|
||||
num_actives_in_lane += remaining[i] > 0;
|
||||
if (remaining[i] > 0)
|
||||
min_in_lane = cute::min(min_in_lane, remaining[i]);
|
||||
}
|
||||
const uint32_t num_active_ranks = __reduce_add_sync(0xffffffff, num_actives_in_lane);
|
||||
const uint32_t length = __reduce_min_sync(0xffffffff, min_in_lane);
|
||||
const uint32_t num_round_tokens = length * num_active_ranks;
|
||||
|
||||
if (slot_idx < num_round_tokens) {
|
||||
const uint32_t slot_idx_in_round = slot_idx % num_active_ranks;
|
||||
uint32_t num_seen_ranks = 0;
|
||||
#pragma unroll
|
||||
for (uint32_t i = 0; i < kNumRanksPerLane; ++ i) {
|
||||
const uint32_t mask = __ballot_sync(0xffffffff, remaining[i] > 0);
|
||||
const uint32_t num_active_lanes = __popc(mask);
|
||||
if (slot_idx_in_round >= num_seen_ranks and slot_idx_in_round < num_seen_ranks + num_active_lanes)
|
||||
current_rank_in_expert_idx = i * 32 + __fns(mask, 0, slot_idx_in_round - num_seen_ranks + 1);
|
||||
num_seen_ranks += num_active_lanes;
|
||||
}
|
||||
token_idx_in_rank = offset + slot_idx / num_active_ranks;
|
||||
break;
|
||||
}
|
||||
|
||||
slot_idx -= num_round_tokens;
|
||||
offset += length;
|
||||
#pragma unroll
|
||||
for (uint32_t i = 0; i < kNumRanksPerLane; ++ i)
|
||||
remaining[i] -= cute::min(remaining[i], length);
|
||||
}
|
||||
|
||||
const uint32_t src_token_topk_idx = *workspace.get_src_token_topk_idx_ptr(
|
||||
static_cast<uint32_t>(current_expert_idx), current_rank_in_expert_idx, token_idx_in_rank);
|
||||
const uint32_t src_token_idx = src_token_topk_idx / kNumTopk;
|
||||
const uint32_t src_topk_idx = src_token_topk_idx - src_token_idx * kNumTopk;
|
||||
const uint32_t pool_token_idx = expert_pool_block_offset * BLOCK_M + token_idx_in_expert;
|
||||
const uint32_t sf_pool_token_idx =
|
||||
(expert_pool_block_offset + token_idx_in_expert / BLOCK_M) * SF_BLOCK_M +
|
||||
(token_idx_in_expert % BLOCK_M);
|
||||
|
||||
const auto remote_token_ptr = sym_buffer.map(
|
||||
input_token_buffer.get_data_buffer(src_token_idx).get_base_ptr<uint4>(),
|
||||
current_rank_in_expert_idx);
|
||||
const auto local_token_ptr = l1_token_buffer.get_data_buffer(pool_token_idx).get_base_ptr<uint4>();
|
||||
#pragma unroll
|
||||
for (uint32_t i = lane_idx; i < kNumTokenUint4; i += 32)
|
||||
local_token_ptr[i] = remote_token_ptr[i];
|
||||
|
||||
const auto remote_sf_ptr = sym_buffer.map(
|
||||
input_sf_buffer.get_data_buffer(src_token_idx).get_base_ptr<float>(),
|
||||
current_rank_in_expert_idx);
|
||||
const auto local_sf_ptr = l1_sf_buffer.get_base_ptr<float>();
|
||||
#pragma unroll
|
||||
for (uint32_t i = lane_idx; i < kNumSFValues; i += 32)
|
||||
local_sf_ptr[i * kNumPaddedSFPoolTokens + sf_pool_token_idx] = remote_sf_ptr[i];
|
||||
|
||||
__syncwarp();
|
||||
if (lane_idx == 0) {
|
||||
const float weight = *sym_buffer.map(
|
||||
input_topk_weights_buffer.get_base_ptr<float>() + src_token_topk_idx,
|
||||
current_rank_in_expert_idx);
|
||||
*l1_topk_weights_buffer.get_data_buffer(pool_token_idx).get_base_ptr<float>() = weight;
|
||||
*workspace.get_token_src_metadata_ptr(pool_token_idx) = {
|
||||
current_rank_in_expert_idx, src_token_idx, src_topk_idx};
|
||||
}
|
||||
__syncwarp();
|
||||
__threadfence();
|
||||
__syncwarp();
|
||||
if (lane_idx == 0) {
|
||||
ptx::red_add_rel(
|
||||
workspace.get_l1_arrival_count_ptr(expert_pool_block_offset + token_idx_in_expert / BLOCK_M), 1);
|
||||
}
|
||||
__syncwarp();
|
||||
}
|
||||
};
|
||||
|
||||
// Count local outgoing token-topk entries per global expert.
|
||||
read_topk_idx([&](const uint32_t&, const uint32_t& expert_idx) {
|
||||
atomicAdd_block(smem_expert_count + expert_idx, 1);
|
||||
});
|
||||
ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx);
|
||||
|
||||
// Convert per-SM counts into global per-rank offsets. High 32 bits count arrived SMs.
|
||||
for (uint32_t i = thread_idx; i < kNumExperts; i += kNumDispatchThreads) {
|
||||
const uint64_t send_value = (1ull << 32) | static_cast<uint64_t>(smem_expert_count[i]);
|
||||
smem_expert_count[i] = static_cast<uint32_t>(
|
||||
ptx::atomic_add(workspace.get_expert_send_count_ptr(i), send_value));
|
||||
}
|
||||
ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx);
|
||||
|
||||
// Write source token-topk indices into destination ranks' local-expert tables.
|
||||
read_topk_idx([&](const uint32_t& token_topk_idx, const uint32_t& expert_idx) {
|
||||
const uint32_t dst_rank_idx = expert_idx / kNumExpertsPerRank;
|
||||
const uint32_t dst_local_expert_idx = expert_idx - dst_rank_idx * kNumExpertsPerRank;
|
||||
const uint32_t dst_slot_idx = atomicAdd_block(smem_expert_count + expert_idx, 1);
|
||||
const auto dst_ptr = workspace.get_src_token_topk_idx_ptr(
|
||||
dst_local_expert_idx, sym_buffer.rank_idx, dst_slot_idx);
|
||||
*sym_buffer.map(dst_ptr, dst_rank_idx) = token_topk_idx;
|
||||
});
|
||||
|
||||
// Wait until all local SMs have finished filling local send-count/source-index data.
|
||||
comm::grid_sync<kNumSMs, kDispatchGridSyncIndex>(
|
||||
workspace, sm_idx, thread_idx, dispatch_sync);
|
||||
|
||||
// Publish per-rank expert counts and summed ready/count words to destination ranks.
|
||||
if (sm_idx == 0) {
|
||||
for (uint32_t i = thread_idx; i < kNumExperts; i += kNumDispatchThreads) {
|
||||
const uint32_t dst_rank_idx = i / kNumExpertsPerRank;
|
||||
const uint32_t dst_local_expert_idx = i - dst_rank_idx * kNumExpertsPerRank;
|
||||
const uint64_t expert_status = *workspace.get_expert_send_count_ptr(i);
|
||||
*sym_buffer.map(
|
||||
workspace.get_expert_recv_count_ptr(sym_buffer.rank_idx, dst_local_expert_idx),
|
||||
dst_rank_idx) = expert_status & 0xffffffffu;
|
||||
ptx::atomic_add_sys(
|
||||
sym_buffer.map(workspace.get_expert_recv_count_sum_ptr(dst_local_expert_idx), dst_rank_idx),
|
||||
expert_status);
|
||||
}
|
||||
}
|
||||
ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx);
|
||||
|
||||
// All ranks must observe the complete count/source-index tables before pulling.
|
||||
comm::nvlink_barrier<kNumRanks, kNumSMs, kNumDispatchThreads,
|
||||
kDispatchGridSyncIndex, kBeforeDispatchPullBarrierTag>(
|
||||
workspace, sym_buffer, sm_idx, thread_idx, dispatch_sync,
|
||||
/* sync_prologue */ false, /* sync_epilogue */ true);
|
||||
|
||||
auto scheduler = sched::SM90MegaMoEScheduler<
|
||||
BLOCK_M, BLOCK_N, BLOCK_K,
|
||||
L1_SHAPE_N, L1_SHAPE_K,
|
||||
L2_SHAPE_N, L2_SHAPE_K,
|
||||
kNumExpertsPerRank,
|
||||
kNumExpertsPerWave,
|
||||
kNumSMs, kNumRanks,
|
||||
kUseNMajorL2>(workspace);
|
||||
scheduler.fetch_expert_recv_count();
|
||||
|
||||
if (sm_idx == 0 and cumulative_local_expert_recv_stats != nullptr) {
|
||||
for (uint32_t i = thread_idx; i < kNumExpertsPerRank; i += kNumDispatchThreads) {
|
||||
const auto num_recv_tokens = static_cast<uint32_t>(*workspace.get_expert_recv_count_sum_ptr(i));
|
||||
ptx::red_add(cumulative_local_expert_recv_stats + i, static_cast<int>(num_recv_tokens));
|
||||
}
|
||||
}
|
||||
ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx);
|
||||
|
||||
constexpr uint32_t kNumRanksPerLane = math::constexpr_ceil_div(kNumRanks, 32u);
|
||||
int current_expert_idx = -1;
|
||||
uint32_t stored_rank_count[kNumRanksPerLane] = {};
|
||||
uint32_t expert_start_idx = 0, expert_end_idx = 0;
|
||||
uint32_t expert_pool_block_offset = 0;
|
||||
|
||||
constexpr uint32_t kNumGlobalDispatchWarps = kNumSMs * kNumDispatchWarps;
|
||||
for (uint32_t token_idx = sm_idx * kNumDispatchWarps + warp_idx; ; token_idx += kNumGlobalDispatchWarps) {
|
||||
int old_expert_idx = current_expert_idx;
|
||||
while (token_idx >= expert_end_idx) {
|
||||
if (++ current_expert_idx >= static_cast<int>(kNumExpertsPerRank))
|
||||
break;
|
||||
|
||||
expert_pool_block_offset += math::ceil_div(expert_end_idx - expert_start_idx, BLOCK_M);
|
||||
expert_start_idx = expert_end_idx;
|
||||
expert_end_idx += scheduler.get_num_tokens(static_cast<uint32_t>(current_expert_idx));
|
||||
}
|
||||
|
||||
if (current_expert_idx >= static_cast<int>(kNumExpertsPerRank))
|
||||
break;
|
||||
|
||||
if (old_expert_idx != current_expert_idx) {
|
||||
old_expert_idx = current_expert_idx;
|
||||
#pragma unroll
|
||||
for (uint32_t i = 0; i < kNumRanksPerLane; ++ i) {
|
||||
const uint32_t j = i * 32 + lane_idx;
|
||||
stored_rank_count[i] = j < kNumRanks ?
|
||||
static_cast<uint32_t>(*workspace.get_expert_recv_count_ptr(j, static_cast<uint32_t>(current_expert_idx))) : 0;
|
||||
}
|
||||
}
|
||||
|
||||
uint32_t current_rank_in_expert_idx = 0;
|
||||
uint32_t remaining[kNumRanksPerLane];
|
||||
#pragma unroll
|
||||
for (uint32_t i = 0; i < kNumRanksPerLane; ++ i)
|
||||
remaining[i] = stored_rank_count[i];
|
||||
|
||||
uint32_t offset = 0;
|
||||
const uint32_t token_idx_in_expert = token_idx - expert_start_idx;
|
||||
uint32_t slot_idx = token_idx_in_expert;
|
||||
uint32_t token_idx_in_rank = 0;
|
||||
while (true) {
|
||||
uint32_t num_actives_in_lane = 0;
|
||||
uint32_t min_in_lane = 0xffffffffu;
|
||||
#pragma unroll
|
||||
for (uint32_t i = 0; i < kNumRanksPerLane; ++ i) {
|
||||
num_actives_in_lane += remaining[i] > 0;
|
||||
if (remaining[i] > 0)
|
||||
min_in_lane = cute::min(min_in_lane, remaining[i]);
|
||||
}
|
||||
const uint32_t num_active_ranks = __reduce_add_sync(0xffffffff, num_actives_in_lane);
|
||||
const uint32_t length = __reduce_min_sync(0xffffffff, min_in_lane);
|
||||
const uint32_t num_round_tokens = length * num_active_ranks;
|
||||
|
||||
if (slot_idx < num_round_tokens) {
|
||||
const uint32_t slot_idx_in_round = slot_idx % num_active_ranks;
|
||||
uint32_t num_seen_ranks = 0;
|
||||
#pragma unroll
|
||||
for (uint32_t i = 0; i < kNumRanksPerLane; ++ i) {
|
||||
const uint32_t mask = __ballot_sync(0xffffffff, remaining[i] > 0);
|
||||
const uint32_t num_active_lanes = __popc(mask);
|
||||
if (slot_idx_in_round >= num_seen_ranks and slot_idx_in_round < num_seen_ranks + num_active_lanes)
|
||||
current_rank_in_expert_idx = i * 32 + __fns(mask, 0, slot_idx_in_round - num_seen_ranks + 1);
|
||||
num_seen_ranks += num_active_lanes;
|
||||
}
|
||||
token_idx_in_rank = offset + slot_idx / num_active_ranks;
|
||||
break;
|
||||
}
|
||||
|
||||
slot_idx -= num_round_tokens;
|
||||
offset += length;
|
||||
#pragma unroll
|
||||
for (uint32_t i = 0; i < kNumRanksPerLane; ++ i)
|
||||
remaining[i] -= cute::min(remaining[i], length);
|
||||
}
|
||||
|
||||
const uint32_t src_token_topk_idx = *workspace.get_src_token_topk_idx_ptr(
|
||||
static_cast<uint32_t>(current_expert_idx), current_rank_in_expert_idx, token_idx_in_rank);
|
||||
const uint32_t src_token_idx = src_token_topk_idx / kNumTopk;
|
||||
const uint32_t src_topk_idx = src_token_topk_idx - src_token_idx * kNumTopk;
|
||||
const uint32_t pool_token_idx = expert_pool_block_offset * BLOCK_M + token_idx_in_expert;
|
||||
const uint32_t sf_pool_token_idx =
|
||||
(expert_pool_block_offset + token_idx_in_expert / BLOCK_M) * SF_BLOCK_M +
|
||||
(token_idx_in_expert % BLOCK_M);
|
||||
|
||||
const auto remote_token_ptr = sym_buffer.map(
|
||||
input_token_buffer.get_data_buffer(src_token_idx).get_base_ptr<uint4>(),
|
||||
current_rank_in_expert_idx);
|
||||
const auto local_token_ptr = l1_token_buffer.get_data_buffer(pool_token_idx).get_base_ptr<uint4>();
|
||||
#pragma unroll
|
||||
for (uint32_t i = lane_idx; i < kNumTokenUint4; i += 32)
|
||||
local_token_ptr[i] = remote_token_ptr[i];
|
||||
|
||||
const auto remote_sf_ptr = sym_buffer.map(
|
||||
input_sf_buffer.get_data_buffer(src_token_idx).get_base_ptr<float>(),
|
||||
current_rank_in_expert_idx);
|
||||
const auto local_sf_ptr = l1_sf_buffer.get_base_ptr<float>();
|
||||
#pragma unroll
|
||||
for (uint32_t i = lane_idx; i < kNumSFValues; i += 32)
|
||||
local_sf_ptr[i * kNumPaddedSFPoolTokens + sf_pool_token_idx] = remote_sf_ptr[i];
|
||||
|
||||
__syncwarp();
|
||||
if (lane_idx == 0) {
|
||||
const float weight = *sym_buffer.map(
|
||||
input_topk_weights_buffer.get_base_ptr<float>() + src_token_topk_idx,
|
||||
current_rank_in_expert_idx);
|
||||
*l1_topk_weights_buffer.get_data_buffer(pool_token_idx).get_base_ptr<float>() = weight;
|
||||
*workspace.get_token_src_metadata_ptr(pool_token_idx) = {
|
||||
current_rank_in_expert_idx, src_token_idx, src_topk_idx};
|
||||
}
|
||||
__syncwarp();
|
||||
__threadfence();
|
||||
__syncwarp();
|
||||
if (lane_idx == 0) {
|
||||
ptx::red_add_rel(
|
||||
workspace.get_l1_arrival_count_ptr(expert_pool_block_offset + token_idx_in_expert / BLOCK_M), 1);
|
||||
}
|
||||
__syncwarp();
|
||||
}
|
||||
} else if (thread_idx < kNumDispatchThreads + kNumTMAThreads) {
|
||||
if constexpr (BLOCK_M == 128) {
|
||||
if (warp_idx == kNumDispatchWarps) {
|
||||
|
||||
Reference in New Issue
Block a user