ByteDance-Seed / ByteDance-Seed/Triton-distributed
[BUG] `reduce_1d_persistent_kernel`: `tl.atomic_add` into a `torch.empty`
- Dominant language
- Python
- Stars
- 1.5k
- Forks
- 172
- PR merge metrics
- No merged PRs in 30d
Description
# Description
`reduce_1d_persistent_kernel` (`python/triton_dist/kernels/nvidia/memory_ops.py:624`) accumulates into its output cell with an atomic add:
```python
for tile_id in range(pid, num_tiles, NUM_SMS):
offs = tile_id * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
mask = offs < reduce_dim
data = tl.load(src_ptr + offs * stride_reduce, mask=mask).to(tl.float32)
accum = tl.sum(data).to(dst_ptr.dtype.element_ty)
tl.atomic_add(dst_ptr, accum) # read-modify-write, so the cell must start at 0
```
But the host wrapper `reduce_tensor` allocates that cell un-zeroed (`memory_ops.py:732`):
```python
output = torch.empty([1], dtype=tensor.dtype, device=tensor.device)
BLOCK_SIZE = 1024
reduce_1d_persistent_kernel[(num_sms, )](tensor, output, reduce_dim, stride_reduce, BLOCK_SIZE)
```
`@1512a81`
Contributor guide
Research direction
Start in python/triton_dist/kernels/nvidia/memory_ops.py at reduce_1d_persistent_kernel around line 624, then inspect the reduce_tensor wrapper around line 732. Verify that the output cell used by tl.atomic_add is initialized before the kernel runs and that the one-dimensional reduction produces deterministic correct results.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- performance
- Issue type
- Bug
- Difficulty
- 2/5
- Estimated time
- 1-3 hours
- Activity status
- Quiet
- Clarity
- Clearly specified
- Newbie friendliness
- 76/100