[PR #527] [CLOSED] Fix the MultiStages for intranode combine and the offset_int4 for int… #613

New issue

Closed

opened 2026-10-27 21:04:32 +00:00 by bsakaguchi · 0 comments

bsakaguchi commented

2026-10-27 21:04:32 +00:00

Owner

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepEP/pull/527
Author: @cyhdmjzzy
Created: 12/19/2025
Status: ❌ Closed

Base: main ← Head: main

📝 Commits (1)

e5e9222 Fix the MultiStages for intranode combine and the offset_int4 for internode_ll combine

📊 Changes

2 files changed (+10 additions, -5 deletions)

View changed files

📝 csrc/kernels/internode_ll.cu (+1 -1)
📝 csrc/kernels/intranode.cu (+9 -4)

📄 Description

Fix the offset_int4 for internode_ll combine

internode_ll.cu:

const auto& offset_int4 = i + 32 * kNumSendUnrolls
            = lane_id * kNumSendUnrolls + iter_idx * 32 * kNumSendUnrolls + 32 * kNumSendUnrolls
            = lane_id * kNumSendUnrolls + (iter_idx + 1) * 32 * kNumSendUnrolls

Wouldn't the lane_id * kNumSendUnrolls part cause a discrepancy?
elect_one_sync() always elects lane 0, so it's equivalent to:
const auto& offset_int4 = (iter_idx + 1) * 32 * kNumSendUnrolls
https://github.com/deepseek-ai/DeepEP/issues/525
https://github.com/deepseek-ai/DeepEP/issues/359

Fix the MultiStages for intranode combine
Each invocation of tma_store_1d issues a TMA bulk_group, and multiple TMA bulk_groups issued by the same thread execute serially (see NVIDIA PTX documentation: 9.7.9.27.2.2. Data Movement and Conversion Instructions: cp.async.bulk.wait_group). Because elect_one_sync() always selects lane 0, the tma_store_1d inside if (elect_one_sync()) will be issued serially by the same lane, preventing a true MultiStage pipeline. Assigning the first kNumStages lanes to distinct stages enables genuine MultiStage concurrency.

Furthermore, the inline assembly in tma_store_wait — asm volatile("cp.async.bulk.wait_group.read %0;" ::"n"(N) : "memory") — waits for the number of remaining TMA bulk_group read-to-shared-memory transactions initiated by the current executing thread to drop to N. If the first kNumStages lanes each own one stage, then a lane that issues tma_store_1d only needs to wait for the single TMA bulk_group it issued in the previous iteration that writes its corresponding stage in shared memory.

Finally, the __syncwarp() immediately after tma_store_wait() is unnecessary, because tma_store_fence() followed by __syncwarp() is sufficient for visibility.

_{🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.}

## 📋 Pull Request Information **Original PR:** https://github.com/deepseek-ai/DeepEP/pull/527 **Author:** [@cyhdmjzzy](https://github.com/cyhdmjzzy) **Created:** 12/19/2025 **Status:** ❌ Closed **Base:** `main` ← **Head:** `main` --- ### 📝 Commits (1) - [`e5e9222`](https://github.com/deepseek-ai/DeepEP/commit/e5e92229f8e452f5a8146fba363122bddccd763a) Fix the MultiStages for intranode combine and the offset_int4 for internode_ll combine ### 📊 Changes **2 files changed** (+10 additions, -5 deletions) <details> <summary>View changed files</summary> 📝 `csrc/kernels/internode_ll.cu` (+1 -1) 📝 `csrc/kernels/intranode.cu` (+9 -4) </details> ### 📄 Description 1. Fix the offset_int4 for internode_ll combine internode_ll.cu: ``` const auto& offset_int4 = i + 32 * kNumSendUnrolls = lane_id * kNumSendUnrolls + iter_idx * 32 * kNumSendUnrolls + 32 * kNumSendUnrolls = lane_id * kNumSendUnrolls + (iter_idx + 1) * 32 * kNumSendUnrolls ``` Wouldn't the lane_id * kNumSendUnrolls part cause a discrepancy? `elect_one_sync()` always elects lane 0, so it's equivalent to: `const auto& offset_int4 = (iter_idx + 1) * 32 * kNumSendUnrolls` https://github.com/deepseek-ai/DeepEP/issues/525 https://github.com/deepseek-ai/DeepEP/issues/359 2. Fix the MultiStages for intranode combine Each invocation of `tma_store_1d` issues a TMA bulk_group, and multiple TMA bulk_groups issued by the same thread execute serially (see [NVIDIA PTX documentation: 9.7.9.27.2.2. Data Movement and Conversion Instructions: cp.async.bulk.wait_group](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#data-movement-and-conversion-instructions-cp-async-bulk-wait-group)). Because `elect_one_sync()` always selects lane 0, the `tma_store_1d` inside `if (elect_one_sync())` will be issued serially by the same lane, preventing a true MultiStage pipeline. Assigning the first kNumStages lanes to distinct stages enables genuine MultiStage concurrency. Furthermore, the inline assembly in `tma_store_wait` — `asm volatile("cp.async.bulk.wait_group.read %0;" ::"n"(N) : "memory") `— waits for the number of remaining TMA bulk_group read-to-shared-memory transactions initiated by the current executing thread to drop to N. If the first kNumStages lanes each own one stage, then a lane that issues `tma_store_1d` only needs to wait for the single TMA bulk_group it issued in the previous iteration that writes its corresponding stage in shared memory. Finally, the `__syncwarp()` immediately after `tma_store_wait()` is unnecessary, because `tma_store_fence()` followed by `__syncwarp()` is sufficient for visibility. --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>

bsakaguchi

2026-10-27 21:04:32 +00:00