flagos-ai / flagos-ai/FlagTree

[BUG][ascend] uint16 and uint32 `>>` are lowered as arithmetic shifts, silently corrupting every value with the top bit set

Open
#1,121 3 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
350
Forks
149
Avg merge
2d 4h
Merged PRs (30d)
81

Description

## Summary

On the Ascend backend, `>>` on a `tl.uint16` or `tl.uint32` value is lowered as
an **arithmetic** (sign-propagating) shift instead of a logical one. Every input
whose top bit is set comes back sign-extended rather than shifted. There is no
error and no warning -- the result is simply wrong.

`tl.uint8` on the same card shifts logically, and every input whose top bit is
*clear* is correct at all three widths, so the unsigned shift is not missing
outright: it looks like the value is sign-extended into a wider register before
the shift, and the shift then propagates a bit that only exists because of that
extension.

## Environment

| | |
|---|---|
| triton | 3.5.1 (`/usr/local/lib/python3.11/site-packages/triton`) |
| torch / torch_npu | 2.8.0+cpu / 2.8.0.post4.dev20260409 |
| CANN | 9.0.0 |
| device | Ascend910B4-1 |

## Reproducer

Self-contained: torch, torch_npu and triton only.

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

B, SH = 128, 5

@triton.jit
def k_u8(x_ptr, out_ptr, SH: tl.constexpr, BLOCK: tl.constexpr):
lane = tl.arange(0, BLOCK)
bits = tl.load(x_ptr + lane).to(tl.uint8, bitcast=True)
tl.store(out_ptr + lane, (bits >> SH).to(tl.int32))

@triton.jit
def k_u16(x_ptr, out_ptr, SH: tl.constexpr, BLOCK: tl.constexpr):
lane = tl.arange(0, BLOCK)
bits = tl.load(x_ptr + lane).to(tl.uint16, bitcast=True)
tl.store(out_ptr + lane, (bits >> SH).to(tl.int32))

@triton.jit
def k_u32(x_ptr, out_ptr, SH: tl.constexpr, BLOCK: tl.constexpr):
lane = tl.arange(0, BLOCK)
bits = tl.load(x_ptr + lane).to(tl.uint32, bitcast=True)
tl.store(out_ptr + lane, (bits >> SH).to(tl.int32))

cases = [("uint8 ", k_u8, torch.int8, 0xFE, 8),
("uint16", k_u16, torch.int16, 0xFFFE, 16),
("uint32", k_u32, torch.int32, 0xFFFFFFFE, 32)]

for name, kernel, dt, raw, width in cases:
x = torch.full((B,), -2, dtype=dt, device="npu") # top bit set
out = torch.zeros(B, dtype=torch.int32, device="npu")
kernel[(1,)](x, out, SH=SH, BLOCK=B)
torch.npu.synchronize()
print(name, hex(raw), "got", int(out[0]),
"| logical", raw >> SH, "| arithmetic", (raw - (1 << width)) >> SH)
```

## Actual output

```
top bit SET -- this is where the two shifts differ
dtype input got logical arith verdict
------------------------------------------------------------------------------
uint8 0xfe 7 7 -1 LOGICAL (correct)
uint16 0xfffe -1 2047 -1 ARITHMETIC (defect)
uint32 0xfffffffe -1 134217727 -1 ARITHMETIC (defect)

top bit CLEAR -- control, every line must read correct
dtype input got expected verdict
--------------------------------------------------------------
uint8 0x7e 3 3 correct
uint16 0x7ffe 1023 1023 correct
uint32 0x7ffffffe 67108863 67108863 correct
```

Expected: `got == logical` on all three lines of the first block.

## Why it matters

The natural way to build a radix bucket out of a float is to bitcast to the
unsigned integer of the same width, flip the non-sign half so the ordering
becomes monotone, and shift. For fp16:

```python
h = x.to(tl.float16)
bits = h.to(tl.uint16, bitcast=True)
sign_set = (bits & tl.full(bits.shape, 0x8000, tl.uint16)) != 0
inv = (~bits) & tl.full(bits.shape, 0x7FFF, tl.uint16)
mapped = tl.where(sign_set, bits, inv)
bin_idx = (mapped >> 5).to(tl.uint32) # <-- destroyed here
```

For a *negative* input `mapped` keeps the raw bits, whose top bit is set by
definition. Measured on this backend, with `x = -1.5` (fp16 bits `0xbe00`):

```
fp16 radix bucket of -1.5 (bits 0xbe00) : got -528, expected 1520
```

`0xbe00` read as int16 is `-16896`, and `-16896 >> 5 == -528` exactly. So every
negative value collapses out of its bucket while every positive one is fine --
a top-k selection that is exactly right on half its input and silently wrong on
the other half.

The same construction at 32 bits (bitcast fp32 to `uint32`, then `bits >> 21`)
is the standard first radix pass, and it has no natural workaround.

## Workaround

Widen first, then mask off the sign extension before shifting:

```python
bin_idx = (mapped.to(tl.int32) & 0xFFFF) >> 5
```

Verified on this stack: `got 2047, expected 2047`.

## Note

The 32-bit case may be a regression rather than a long-standing defect: we have
a production kernel whose first radix pass is `bits >> 21` on a genuine
`tl.uint32` with no workaround, and it passes its correctness suite on a
triton 3.2.0 / CANN 8.5.0 Ascend stack -- which it could not do if the shift
were arithmetic there. We have not yet been able to run this reproducer on that
older stack to confirm; we will follow up in a comment when we have.

Contributor guide

Open the contributing guide

Research direction

Start by running the self-contained Triton reproducer on the Ascend910B4 environment described in the issue, comparing uint8, uint16, and uint32 results for set and clear top bits. Trace the Ascend backend lowering of unsigned right shift and verify that the corrected behavior matches the logical values, including the fp16 radix-bucket example and the existing control cases.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
compilers, machine-learning
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Active
Clarity
Mostly clear
Newbie friendliness
48/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.