ByteDance-Seed / ByteDance-Seed/Triton-distributed
Possible incorrect barrier offset in `tile_kernel_dispatch_token_intra_node` zero-token expert path
- Dominant language
- Python
- Stars
- 1.5k
- Forks
- 172
- PR merge metrics
- No merged PRs in 30d
Description
```markdown
## Problem
In `python/triton_dist/kernels/nvidia/ep_all2all_fused.py`, inside `tile_kernel_dispatch_token_intra_node`, the zero-token expert notification path appears to compute the barrier offset with the global expert rank instead of the local expert id.
Current code:
```python
if pid == 0:
for i in range(thread_idx, experts_per_rank * world_size, num_warps * WARP_SIZE):
tokens_this_expert = ld(local_splits_buf + i)
if tokens_this_expert == 0:
libshmem_device.signal_op(
barriers_ptr + i // experts_per_rank * world_size + rank,
1,
libshmem_device.NVSHMEM_SIGNAL_SET,
i // experts_per_rank,
)
```
Here `i` is the global expert id. It should be decomposed as:
```python
expert_rank = i // experts_per_rank
expert_idx_intra_rank = i % experts_per_rank
```
The destination rank for the signal should be `expert_rank`, but the offset into `barriers_ptr` should use `expert_idx_intra_rank`, because `barriers_ptr` is indexed as:
```text
[experts_per_rank, world_size]
```
The non-zero-token path uses this layout correctly:
```python
barriers_ptr + expert_idx_intra_rank * world_size + rank
```
## Expected Fix
The zero-token path should likely be:
```python
libshmem_device.signal_op(
barriers_ptr + i % experts_per_rank * world_size + rank,
1,
libshmem_device.NVSHMEM_SIGNAL_SET,
i // experts_per_rank,
)
```
## Impact
With the current code, zero-token experts may signal the wrong local expert barrier slot on the destination rank. This can cause:
- grouped GEMM waiting forever for a barrier that was never signaled;
- a different expert barrier being incorrectly marked ready;
- intermittent hangs or incorrect synchronization depending on routing distribution.
This is most visible when some source rank sends zero tokens to a given expert.
```
Contributor guide
Research direction
Start in python/triton_dist/kernels/nvidia/ep_all2all_fused.py at tile_kernel_dispatch_token_intra_node and compare the zero-token notification path with the non-zero-token barrier indexing. Verify the local expert index and destination rank against the [experts_per_rank, world_size] layout, then confirm that zero-token experts signal the correct barrier without affecting other slots.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- distributed-systems
- Issue type
- Bug
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Quiet
- Clarity
- Clearly specified
- Newbie friendliness
- 72/100