ByteDance-Seed / ByteDance-Seed/Triton-distributed

Possible incorrect barrier offset in `tile_kernel_dispatch_token_intra_node` zero-token expert path

Open
#169 0 comments 0 reactions 0 assignees View on GitHub
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.