ByteDance-Seed / ByteDance-Seed/Triton-distributed

[BUG] `reduce_1d_persistent_kernel`: `tl.atomic_add` into a `torch.empty`

Open Beginner friendly
#182 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

# 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

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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.