ByteDance-Seed / ByteDance-Seed/Triton-distributed
copy_tensor: the persistent 2-D arm computes row offsets in int32, so a copy whose row extent passes 2 GiB reads and writes out of bounds
- Dominant language
- Python
- Stars
- 1.5k
- Forks
- 172
- PR merge metrics
- No merged PRs in 30d
Description
## Summary
`copy_tensor` has two 2-D kernels. `copy_2d_kernel` (`persistent=False`) casts the row index to int64 before scaling it by the row stride; `copy_2d_persistent_kernel` does not. Once `(M - 1) * stride(0)` reaches `2**31` the int32 product wraps negative, and the copy reads 2 GiB before the source and — when the destination is large too — writes 2 GiB before the destination. The uncast arm is the default: `copy_tensor(..., persistent=True)` is the signature default and five of the six in-tree call sites take it.
This needs either a big tensor (any 2-D tensor with more than `2**31` elements) or a large row stride; the 2-D branch has no contiguity assert, so a row slice of a large tensor also gets there.
https://github.com/ByteDance-Seed/Triton-distributed/blob/8260bc34398c2b8f36dc840fd22f741ca9294584/python/triton_dist/kernels/nvidia/memory_ops.py#L378-L383
## Reproduction
Both cases carve their tensors out of one allocation with a 2 GiB guard band in front, so the wrapped accesses land in memory the script owns and the result is a printed value rather than an illegal-access crash.
```python
import torch
from triton_dist.kernels.nvidia.memory_ops import copy_tensor
GiB, PAD = 1 << 30, 1 << 31
NUM_SMS = torch.cuda.get_device_properties(0).multi_processor_count
# case 1 -- 3 rows of 8 int8 taken 2**30 elements apart: row 2 is at offset 2**31.
buf = torch.empty(PAD + 2 * GiB + 8, dtype=torch.int8, device="cuda")
buf[:PAD].fill_(7) # guard band, 2 GiB before src
src = torch.as_strided(buf, size=(3, 8), stride=(GiB, 1), storage_offset=PAD)
src.fill_(42)
for persistent in (True, False):
dst = torch.full((3, 8), -1, dtype=torch.int8, device="cuda")
copy_tensor(dst, src, num_sms=NUM_SMS, persistent=persistent)
torch.cuda.synchronize()
print(f"persistent={persistent!s:5s} dst[2] = {dst[2].tolist()}")
del buf; torch.cuda.empty_cache()
# case 2 -- a plain contiguous 2.15 GB tensor: rows 262144.. wrap on BOTH the load and the store.
M, N = 262400, 8192
elems = M * N
buf = torch.empty(PAD + elems + PAD + elems, dtype=torch.int8, device="cuda")
buf[:PAD].fill_(7)
src = buf[PAD:PAD + elems].view(M, N); src.fill_(42)
for persistent in (True, False):
buf[PAD + elems:PAD + elems + PAD].fill_(0) # guard band, 2 GiB before dst
dst = buf[PAD + elems + PAD:].view(M, N); dst.fill_(-1)
copy_tensor(dst, src, num_sms=NUM_SMS, persistent=persistent)
torch.cuda.synchronize()
print(f"persistent={persistent!s:5s} wrong rows {int((dst != 42).any(dim=1).sum())}/{M}"
f" bytes written into the guard band before dst: {int((buf[PAD + elems:PAD + elems + PAD] != 0).sum())}")
```
```
persistent=True dst[2] = [7, 7, 7, 7, 7, 7, 7, 7]
persistent=False dst[2] = [42, 42, 42, 42, 42, 42, 42, 42]
persistent=True wrong rows 256/262400 bytes written into the guard band before dst: 2097152
persistent=False wrong rows 0/262400 bytes written into the guard band before dst: 0
```
The `7`s in case 1 come from 2 GiB before the source. In case 2 the 256 wrong rows are exactly the rows from 262144 (`2**31 // 8192`) on, and 2097152 = 256 x 8192 is exactly those rows written to the wrong place.
## Environment
- Repo commit tested: `8260bc34398c2b8f36dc840fd22f741ca9294584` (main)
- GPU: NVIDIA B200 (sm_100), driver 595.71.05, CUDA 13.2
- torch 2.12.1+cu130, the repo's vendored Triton 3.4.0, Python 3.11
Contributor guide
Research direction
Start in python/triton_dist/kernels/nvidia/memory_ops.py at the persistent 2-D kernel around lines 378-383, and compare its row-offset calculation with copy_2d_kernel. Use the reproduction in the issue to exercise both persistent modes on large-stride and large contiguous tensors. Done means persistent=True produces the same correct rows as persistent=False without reading or writing into the guard bands.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- backend, performance
- Issue type
- Bug
- Difficulty
- 2/5
- Estimated time
- 1-2 days
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 78/100