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`
- 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
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