ByteDance-Seed / ByteDance-Seed/Triton-distributed

`calc_gather_scatter_index_v2_kernel`: unmasked histogram reduction counts padding lanes into expert 0 when `ntokens * topk % 1024 != 0`

Open
#181 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

## Summary

The per-expert histogram loop in `calc_gather_scatter_index_v2_kernel` (`python/triton_dist/kernels/nvidia/moe_utils.py`) loads the final block with a mask but reduces without one:

```python
offs = n * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
mask = offs < M
expert_idx = tl.load(choosed_experts_ptr + offs, mask=mask) # no `other=`: masked-off lanes read as 0
val += tl.cast(tl.sum(expert_idx == pid), tl.int32) # reduction is NOT masked
```

On the final block, lanes with `offs >= M` read as 0, so `tl.sum(expert_idx == pid)` counts them for `pid == 0`. The wrapper `calc_gather_scatter_index_v2_triton` hardcodes `BLOCK_SIZE=1024`, so whenever `M = ntokens * topk` is not a multiple of 1024, `ntokens_by_expert[0]` is inflated by exactly `ceil(M / 1024) * 1024 - M`. In real inference `ntokens = batch * seq` is arbitrary, so this fires for almost every workload.

## Why the existing test misses it

`test/nvidia/test_moe_utils.py` uses `ntokens ∈ {1024, 2048, 4096, 8192}`, so `M` is always a multiple of 1024 and the partial final block never occurs. It also checks per-expert ranges derived from the kernel's own `ntokens_by_expert` rather than comparing the histogram against an independent reference.

## Minimal reproduction (single GPU, no distributed)

```python
import torch
from triton_dist.kernels.nvidia.moe_utils import calc_gather_scatter_index_v2_triton

torch.manual_seed(0)
ntokens, topk, nexperts = 2047, 2, 64 # M = 4094 -> the last 1024-wide block has 2 masked-off lanes
chosen = torch.randint(0, nexperts, (ntokens, topk), device="cuda", dtype=torch.int32)
M = ntokens * topk

hist, scatter_index, gather_index, expert_index, M_pad = \
calc_gather_scatter_index_v2_triton(chosen, nexperts, 1)
ref = torch.bincount(chosen.flatten().long(), minlength=nexperts).to(torch.int32)

print("experts with wrong count:", (hist - ref).nonzero().flatten().tolist())
print("hist[0] - bincount[0] =", int(hist[0] - ref[0]))
print("M_pad =", int(M_pad.item()), " M =", M)
print("scatter_index is a permutation of [0, M):",
bool(torch.equal(scatter_index.flatten().sort()[0],
torch.arange(M, dtype=torch.int32, device="cuda"))))
```

Output on `main` (1512a81, H100, torch 2.8.0+cu128):

```
experts with wrong count: [0]
hist[0] - bincount[0] = 2
M_pad = 4096 M = 4094
scatter_index is a permutation of [0, M): False
```

The over-count is buffer-independent (the masked-off lanes deterministically read the masked-load default 0, so the excess always lands on expert 0) and equals the padding-lane count.

## Reachability

`TP_MoE.run_moe_reduce_rs` (`python/triton_dist/layers/nvidia/tp_moe.py:266`) → `run_moe_reduce_rs` → `moe_reduce_rs.py:932` → `calc_gather_scatter_index_v2_triton`. Any token count that is not a multiple of `1024 / topk` corrupts the gather/scatter plan of the fused MoE reduce-scatter path.

## Environment

- Triton-distributed `main` @ 1512a81 (also inspected: the unmasked reduction is present at current HEAD)
- 1x H100, CUDA 12 driver 550.90.12, torch 2.8.0+cu128, python 3.11
- Single-GPU repro; no NVSHMEM or multi-GPU setup needed

Contributor guide

Open the contributing guide

Research direction

Start with calc_gather_scatter_index_v2_kernel and calc_gather_scatter_index_v2_triton in python/triton_dist/kernels/nvidia/moe_utils.py, then inspect test/nvidia/test_moe_utils.py. Run the single-GPU reproduction with a non-multiple-of-1024 M and compare the histogram with an independent torch.bincount reference. Done means padding lanes no longer change expert counts and scatter_index remains a permutation of [0, M).

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
distributed-systems, testing-qa
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Quiet
Clarity
Clearly specified
Newbie friendliness
65/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.