Skip to content

Commit ec4d1f3

Browse files
committed
[NVBUG-6250866][bugfix] fix DeepEP intranode combine fallback
Signed-off-by: Dongfeng Yu <dongfengy@nvidia.com> [NVBUG-6250866][bugfix] wire DeepEP patch into FetchContent Signed-off-by: Dongfeng Yu <dongfengy@nvidia.com>
1 parent 33b0a32 commit ec4d1f3

2 files changed

Lines changed: 31 additions & 1 deletion

File tree

3rdparty/fetch_content.json

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,8 @@
2626
"display_name": "deep_ep",
2727
"git_repository": "${github_base_url}/deepseek-ai/DeepEP",
2828
"git_tag": "5be51b228a7c82dbdb213ea58e77bffd12b38af8",
29-
"use_url": true
29+
"use_url": true,
30+
"patch_file": "patches/deep_ep_intranode_combine_fix.patch"
3031
},
3132
{
3233
"name": "deepgemm",
Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,29 @@
1+
--- a/csrc/kernels/intranode.cu
2+
+++ b/csrc/kernels/intranode.cu
3+
@@ -846,7 +846,9 @@
4+
// Wait TMA arrival
5+
if (lane_id == 0)
6+
tma_store_wait<kNumStages - 1>();
7+
- __syncwarp();
8+
+ // hidden_int4 is not always divisible by a warp. The final tile can have
9+
+ // only a subset of lanes active, so synchronize only currently active lanes.
10+
+ __syncwarp(__activemask());
11+
12+
// Write into TMA buffer
13+
auto tma_stage_idx = (i / 32) % kNumStages;
14+
@@ -855,13 +857,13 @@
15+
16+
// Issue TMA
17+
tma_store_fence();
18+
- __syncwarp();
19+
+ __syncwarp(__activemask());
20+
if (lane_id == 0) {
21+
auto tma_bytes = min(32, hidden_int4 - i) * static_cast<int>(sizeof(int4));
22+
tma_store_1d(reinterpret_cast<int4*>(tma_buffer) + tma_stage_idx * 32,
23+
recv_int4 + token_idx * hidden_int4 + i, tma_bytes, false);
24+
}
25+
- __syncwarp();
26+
+ __syncwarp(__activemask());
27+
#else
28+
recv_int4[token_idx * hidden_int4 + i] = out_int4;
29+
#endif

0 commit comments

Comments
 (0)