ByteDance-Seed / ByteDance-Seed/Triton-distributed
[NVIDIA] low_latency_all_to_all_v2: int32 buffer offsets overflow at large num_experts * max_m * hidden, corrupting dispatch/combine
- Dominant language
- Python
- Stars
- 1.5k
- Forks
- 172
- PR merge metrics
- No merged PRs in 30d
Description
## Symptom
As `num_experts * max_m * hidden` grows, the low-latency all-to-all v2
kernels compute element offsets that cross `2^31` and wrap negative. The
dispatch side trips first (its per-token message stride is larger than
`hidden`); the combine side follows slightly later. Depending on where the
wrapped address lands, the result is an illegal memory access or silently
corrupted routed tokens. There is no size check on this path, so nothing
tells the user the config is out of range.
## Root cause
At HEAD (8260bc3), five element-offset expressions in
`python/triton_dist/kernels/nvidia/low_latency_all_to_all_v2.py` are computed
in int32 -- every operand is an int32 scalar or a constexpr int, so Triton
evaluates the product in int32 before the pointer add:
| line | kernel | expression, offsets reached |
|---|---|---|
| 271 | `dispatch_kernel_v2` | `(dst_expert_local_idx * world_size * MAX_M + rank * MAX_M + dst_slot) * ELE_PER_MSG`; up to `(num_local_experts * world_size * max_m - 1) * ELE_PER_MSG` |
| 415 | `combine_kernel_v2` | `local_expert_idx * world_size * MAX_M * HIDDEN + dispatch_start * HIDDEN` |
| 416 | `combine_kernel_v2` | same form as L415, for `send_tokens_comm_buf` |
| 435-436 | `combine_kernel_v2` | `(rank * NUM_EXPERTS_PER_RANK + local_expert_idx) * MAX_M * HIDDEN + dispatch_src_idx * HIDDEN`; up to ~`num_experts * max_m * hidden` (final element accessed includes the per-token element index) |
| 480 | `combine_kernel_v2` | `(token_topk_indices * MAX_M + token_id) * HIDDEN`; same bound as L435 |
`dispatch_postprocess_kernel_v2_for_expert` in the same file already promotes
its equivalent offsets via `.to(tl.int64)`, so these five look like an
oversight rather than a design decision.
## Affected configs
For the online-FP8 path (the only mode `EPLowLatencyAllToAllLayer` accepts),
`ELE_PER_MSG = hidden + 16 + hidden/32` (fp8 payload + 16 meta bytes +
fp32 group scales at group size 128), so the dispatch side (L271) wraps once
`num_local_experts * world_size * max_m * ELE_PER_MSG` passes `2^31`:
- 256 experts, hidden 7168 (DeepSeek-V3 shape): wraps from `max_m ~= 1133`
- 256 experts, hidden 4096, max_m 2048: `256 * 2048 * 4240 = 2.22e9 > 2^31`
The combine side (L435, L480) wraps once `num_experts * max_m * hidden`
itself passes `2^31`, i.e. once the per-phase combine receive buffer grows
past 4 GiB in bf16.
The AMD variant (`python/triton_dist/kernels/amd/low_latency_all_to_all_v2.py`)
has the same five int32 expressions at lines 287, 431, 432, 451-452, 499; I
did not test on AMD hardware.
## Repro (2x H100 80GB, single node, repo test unmodified)
`-G 64 -M 4680 -N 7168` puts the L271 offset past `2^31` for every token
routed to a local-expert index of 31 (global experts 31 and 63) -- base
offset `31 * 2 * 4680 * 7408 = 2,149,505,280 > 2^31` before the slot term --
while keeping `num_experts * max_m * hidden = 2.147e9` just under `2^31`:
```bash
export NVSHMEM_SYMMETRIC_SIZE=36000000000
export NVSHMEM_DISABLE_CUDA_VMM=1
export CUDA_DEVICE_MAX_CONNECTIONS=1
# plus your usual NVSHMEM bootstrap settings; we used the UID bootstrap
# over loopback, and CUDA_VISIBLE_DEVICES to pick two GPUs
python -m torch.distributed.run --nproc_per_node=2 \
python/triton_dist/test/nvidia/test_ep_ll_a2a.py \
-M 4680 -N 7168 -G 64 --topk 8 --check --iters 1 --verify-iters 1
```
Token counts and routing are random and unseeded, so the run is not formally
deterministic, but any token routed to global expert 31 or 63 wraps; at these
token counts a passing run is vanishingly unlikely, and it failed on every
attempt here.
Observed at HEAD: both ranks abort with
`RuntimeError: Triton Error [CUDA]: an illegal memory access was encountered`,
raised at the first kernel launch after `dispatch()` returns (the async fault
comes from the dispatch kernels).
Controls:
- same command with `-M 2048` (below the boundary): passes, unpatched and patched;
- same command at `-M 4680` with the five casts applied: passes the test's
exact (`atol=0, rtol=0`) comparison against its torch reference.
Aside: the config is pinned to `-M 4680` because the test harness itself has
the same class of bug -- `_dequant_kernel` in
`python/triton_dist/test/nvidia/ep_a2a_utils.py` computes `off_m[:, None] * N`
in int32, so any config with `num_experts * max_m * hidden > 2^31` breaks the
test utility even with the library fixed. Larger configs corrupt in the same
way, they just cannot be validated with the stock test.
## Environment
- Triton-distributed @ 8260bc3, built from source (triton 3.4.0, ptxas from
the bundled wheel)
- torch 2.7.1+cu128, CUDA 12.8, NVSHMEM 3.6.5 (UID bootstrap, single node)
- 2x NVIDIA H100 80GB HBM3, driver 535.230.02, python 3.12
## Fix
Casting the leading index of each of the five expressions to `tl.int64`
(the pattern `dispatch_postprocess_kernel_v2_for_expert` already uses) fixes
the repro above. No regression observed from the casts in two A/B runs at a
below-boundary size on 8x H100 (timing ratios 1.01 and 1.00 vs uncast).
Contributor guide
Research direction
Start in python/triton_dist/kernels/nvidia/low_latency_all_to_all_v2.py and compare the five listed offset expressions with the existing promoted offsets in dispatch_postprocess_kernel_v2_for_expert. Review the corresponding expressions in python/triton_dist/kernels/amd/low_latency_all_to_all_v2.py, then run test/nvidia/test_ep_ll_a2a.py with the reported boundary configuration and control size. Done means the large configuration passes the exact reference comparison without illegal accesses, while the control remains passing.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- distributed-systems, performance
- Issue type
- Bug
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 76/100