[ascend] tl.atomic_add returns non-unique per-lane values while the accumulation is correct

[cheersluvs]

Summary

On the Ascend backend, tl.atomic_add accumulates correctly but hands most
lanes the wrong return value. Lanes are grouped in eights, and every lane in a
group receives that group's base rather than the value immediately before its
own addition.

The contract is that N lanes each adding 1 to one address receive N distinct
values. Here 512 lanes receive 65 distinct values, and the counter is
nevertheless exactly 512.

Environment

triton	3.2.0 (triton/backends/ascend)
torch / torch_npu	2.6.0+cpu / 2.6.0rc1
CANN	8.5.0
device	Ascend910B4-1, driver 25.2.0, aarch64

Reproducer

Self-contained: torch, torch_npu and triton only.

import torch, torch_npu, triton
import triton.language as tl

N = 512

@triton.jit
def all_lanes(CNT, POS, N: tl.constexpr):
    lane = tl.arange(0, N)
    p = tl.atomic_add(CNT + tl.zeros([N], tl.int32),
                      tl.full([N], 1, tl.int32), sem="relaxed")
    tl.store(POS + lane, p)

@triton.jit
def half_lanes(CNT, POS, N: tl.constexpr):
    lane = tl.arange(0, N)
    m = (lane % 2) == 0
    p = tl.atomic_add(CNT + tl.zeros([N], tl.int32),
                      tl.full([N], 1, tl.int32), mask=m, sem="relaxed")
    tl.store(POS + lane, p, mask=m)

for kernel, n_active, label in ((all_lanes, N, "all lanes"),
                                (half_lanes, N // 2, "every other lane")):
    cnt = torch.zeros(1, dtype=torch.int32, device="npu")
    pos = torch.full((N,), -1, dtype=torch.int32, device="npu")
    kernel[(1,)](cnt, pos, N=N)
    torch.npu.synchronize()
    vals = [v for v in pos.cpu().tolist() if v != -1]
    print(label, "counter", int(cnt[0]), "expected", n_active,
          "| distinct", len(set(vals)), "expected", n_active,
          "| first", sorted(set(vals))[:6])

Actual output

all 512 lanes
  participating lanes : 512
  counter after       : 512   (expected 512)
  distinct returns    : 65    (expected 512)
  first distinct      : [1, 2, 9, 17, 25, 33]
  lanes per distinct  : 7.88

every other lane (256 active)
  participating lanes : 256
  counter after       : 256   (expected 256)
  distinct returns    : 64    (expected 256)
  first distinct      : [1, 5, 9, 13, 17, 21]
  lanes per distinct  : 4.00

Expected in both cases: distinct == participating lanes.

Analysis

Measured facts:

• the accumulation is correct in both configurations
• the returned values form an arithmetic progression whose step equals the
  number of active lanes per group of eight -- 8 when all lanes participate,
  4 when every other one does
• lanes per distinct is exactly 4.00 in the masked case

Consistent with a group-aggregated atomic: one atomic per eight lanes adding the
group's total, correct, but the group's base returned to every lane in the group
instead of base plus the intra-group prefix. Offered as a hypothesis; the
measurements above are the claim.

Impact

Stream compaction is the standard use of the return value:

pos = tl.atomic_add(counter, 1, mask=take)   # claim a slot
tl.store(out + pos, value, mask=take)        # write into it

With eight lanes handed the same slot the stores collide, and a masked store
with duplicate lane addresses appears to be dropped on this backend, so roughly
seven eighths of the output is lost. Because the count is correct, every
downstream bound check based on it passes and nothing detects the loss.

We hit this in top_k_per_row_prefill in FlagGems-vllm, where the two atomics
behave oppositely and localise it precisely:

use	return value	result
histogram binning	discarded	complete and correct, all 2048 bins
output compaction	used as a store address	7 of 64 entries written, each valid

vLLM's own CUDA kernel for this operator uses the same idiom
(int dstIdx = atomicAdd(&counter, 1)), so a faithful port inherits the problem.

Scope

We have replaced the compaction with a prefix sum that does not use the return
value, and verified on Moore Threads that both paths are equivalent (20/20 on
the operator's test suite each). The operator is still incorrect on Ascend
after that change, so there is at least one further problem beyond this one.
This issue covers only the atomic return value.