Skip to content

GitLab

"vscode:/vscode.git/clone" did not exist on "454caac9101da608e1582e7cb3277bf60fd3c1b3"
Commit 24324cc5 authored by Shangyan Zhou's avatar Shangyan Zhou
Browse files

Fix combine tma mbarrier

parent ee5bd170
Show whitespace changes
Inline Side-by-side
Showing with 2 additions and 2 deletions
csrc/kernels/internode.cu
View file @ 24324cc5
...@@ -1439,13 +1439,13 @@ combine(int4* combined_x, float* combined_topk_weights, ...@@ -1439,13 +1439,13 @@ combine(int4* combined_x, float* combined_topk_weights,
// TMA stuffs // TMA stuffs
extern __shared__ __align__(1024) uint8_t smem_tma_buffer[]; extern __shared__ __align__(1024) uint8_t smem_tma_buffer[];
auto tma_buffer = smem_tma_buffer + dst_nvl_rank * kNumTMABytesPerSenderWarp; auto tma_buffer = smem_tma_buffer + dst_nvl_rank * kNumTMABytesPerSenderWarp;
auto tma_mbarrier = reinterpret_cast<uint64_t*>(tma_buffer + hidden_bytes); auto tma_mbarrier = reinterpret_cast<uint64_t*>(tma_buffer + num_bytes_per_token);
uint32_t tma_phase = 0; uint32_t tma_phase = 0;
if (lane_id == 0) { if (lane_id == 0) {
mbarrier_init(tma_mbarrier, 1); mbarrier_init(tma_mbarrier, 1);
fence_view_async_shared(); fence_view_async_shared();
fence_barrier_init(); fence_barrier_init();
EP_DEVICE_ASSERT(hidden_bytes + sizeof(uint64_t) <= kNumTMABytesPerSenderWarp); EP_DEVICE_ASSERT(num_bytes_per_token + sizeof(uint64_t) <= kNumTMABytesPerSenderWarp);
} }
__syncwarp(); __syncwarp();
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment