[Bug][metax][TLE] smem support is incomplete: scalar-indexed `local_ptr` asserts at LoadStoreOpToLLVM.cpp:427, smem atomics are unsupported, and smem+histogram silently corrupts data

- Repo: https://github.com/flagos-ai/FlagTree
- Labels: `bug`, `metax`, `tle`
- Environment: flagtree main `2cd4cc6a` (wheel `0.6.0+metax.git2cd4cc6a`, 2026-09-09) on MetaX C550
- All reproducers are inlined in the Appendix; every listed kernel is self-contained (torch + flagtree only)
- Related: #742 (Add Metax TLE support), #971 (Metax TLE support local pointer)

## Summary

Metax TLE shared-memory support is only partially implemented, and the supported/unsupported split is asymmetric. Five problems, each with a minimal deterministic reproducer:

| # | Problem | Kind | Listing |
|---|---|---|---|
| 1 | Scalar-indexed `local_ptr` + offset access asserts at `LoadStoreOpToLLVM.cpp:427` | compile-time | 1, 2 |
| 2 | `tt.atomic_rmw` on smem pointers fails front-end verification | compile-time | 3 |
| 3 | `local_ptr` indices depending on a dynamic loop variable crash | compile-time | 4, 5, 6 |
| 4 | TLE smem buffer overlaps `tl.histogram`'s smem scratch (large tile) | runtime | 7, 8 |
| 5 | Masked smem load ignores `other` | wrong value (minor) | 9 |
| – | `buffered_tensor` cannot cross function boundaries | limitation | 10 |
| – | smem capacity: 32 KB + histogram scratch fits, 64 KB does not | informational | — |

All affected patterns are used by production TLE kernels on NVIDIA (FlagGems `top_k_per_row_decode` / `bin_topk`, FlagGems-vllm `persistent_topk`) and/or covered by upstream TLE tests (`python/test/tle/unit/test_tle_gpu_local_ptr.py`). The metax CI (`metax3.6-build-and-test.yml`) does not run `python/test/tle/`, so none of them are exercised on metax.

The merged metax TLE local-pointer PR #971 explicitly declares support for scalar, tensor, static and dynamic `local_ptr` indices, for masked loads and stores through local pointers, and for Triton atomic operations on shared memory. Problems 1, 2, 3 and 5 below are in direct contradiction with those declarations.

## Environment

| | |
|---|---|
| flagtree | main `2cd4cc6a` (2026-09-09), wheel `0.6.0+metax.git2cd4cc6a`, triton 3.6.0 |
| build | `FLAGTREE_BACKEND=metax`, `-DFLAGTREE_TLE=ON`, `LLVM_SYSPATH=metax-llvm19` (reports LLVM 22.0.0git) |
| torch / MACA | 2.8.0+metax3.7.1.4 / 3.7.1.5 |
| device | MetaX C550 (x86_64) |

---

## Problem 1: scalar-indexed `local_ptr` + offsets assert at LoadStoreOpToLLVM.cpp:427

**Pattern** (FlagGems `top_k_per_row_decode.py` L691-712, runs on NVIDIA):

```python
s_radix_count_ptr = tle.gpu.local_ptr(s_radix_counts, (0,)) # scalar index
radix_count_vec_ptr = s_radix_count_ptr + bins # scalar base + tensor offset
tl.store(s_radix_count_ptr + lane, 0, mask=lane < RADIX)
```

**Minimal repro** (Listing 1): scalar index + `base + lane` access.

**Actual**: compile-time crash during `ConvertTritonGPUToLLVM`:

```
.../third_party/metax/plugin/lib/TritonMETAXGPUToLLVM/LoadStoreOpToLLVM.cpp:427:
... LoadOpConversion::matchAndRewrite(...):
Assertion `wordNElems * nWords * numVecs == numElems' failed.
```

Reproduced for BLOCK ∈ {256, 512, 1024}, num_warps ∈ {4, 8, 16} (sizePerThread [1] and [2] both fail). The same assertion exists in `StoreOpConversion`.

**The equivalent tensor-index form works** (Listing 2): the identical access with a *tensor* index passes. So the missing piece is specifically the scalar-pointer (`tt.splat`) lowering path; full-view and tensor-index pointer tensors both work.

**Impact**: the scalar-index form is the documented usage used by production kernels; it works on NVIDIA.

---

## Problem 2: `tt.atomic_rmw` on smem pointers is unsupported

**Pattern** (FlagGems `top_k_per_row_decode.py` L728, runs on NVIDIA):

```python
tl.atomic_add(s_radix_count_ptr + digit, ones, mask=take,
sem="relaxed", scope="cta")
```

**Minimal repro** (Listing 3): the exact production form.

**Actual**: parse-time failure:

```
error: 'tt.atomic_rmw' op failed to verify that ptr type matches value type
%3 = "tt.atomic_rmw"(...) : (tensor<256x!tt.ptr>, tensor<256xi32>, tensor<256xi1>) -> ...
```

Reproduced for: no offset (direct full-view pointers), `scope="gpu"`, scalar value, and without `sem`/`scope` args. The failure happens at front-end verification, so no `mctle` handling is reached.

**Impact**: smem histogram accumulation (the standard fast path for radix/top-k kernels) must fall back to global atomics, which is exactly the slow path these kernels are designed to avoid.

---

## Problem 3: `local_ptr` indices depending on a dynamic loop variable crash

A dynamic loop variable inside the `local_ptr` indices crashes `ConvertTritonGPUToLLVM` (427 assertion or pass failure, depending on layout). The loop context itself is fine. The trigger boundary is defined by three variants (Listings 4-6):

```python
# Listing 4 (PASS): local_ptr inside tl.range, indices fixed
for t in tl.range(0, NTILE):
tl.store(tle.gpu.local_ptr(buf, (flat,)), ...)

# Listing 5 (CRASH): same loop, indices depend on the loop variable
for t in tl.range(0, NTILE):
idx = t * BLOCK + flat
tl.store(tle.gpu.local_ptr(buf, (idx,)), tl.load(in_ptr + idx))

# Listing 6 (PASS): static unroll with the same expression
for t in tl.static_range(0, NTILE):
idx = t * BLOCK + flat
tl.store(tle.gpu.local_ptr(buf, (idx,)), tl.load(in_ptr + idx))
```

**Impact**: this pattern is covered by the upstream TLE test `test_tle_gpu_local_ptr.py::_local_pointer_looped_elementwise_kernel` (`for slice_idx in range(SLICES): local_ptr(smem, (block_offset + slice_indices,))`) and is required for tile-wise processing of buffers larger than one tile.

---

## Problem 4: TLE smem + `tl.histogram` silently corrupts data (large tile)

A plain smem round-trip with a histogram in between returns wrong values, stably. Minimal repro (Listing 7, TILE=4096, num_warps=8, `local_ptr` used in the documented tensor-index form): convert a float tile, store it to smem, run a histogram, read the buffer back. The kernel writes the read-back values to the output; a sentinel-initialized output distinguishes "wrong value" from "never written".

**Actual**: stable FAIL across runs — 480-544 of 4096 outputs wrong; all slots are written, the *values* are wrong; the error count varies between runs (non-deterministic magnitude).

**Root cause (localized)**: the TLE buffer and the shared-memory scratch used by `tl.histogram` are allocated on top of each other. Evidence from the metax compile of Listing 7:

- the kernel's `shared` size is 16384 B — exactly the TLE buffer (TILE=4096 × 4 B); the histogram's scratch is not accounted for;
- in the generated LLIR, the histogram's shared-memory accumulator is zero-initialized and updated with `atomicrmw add` at `@global_smem + 0` (256 bins, i.e. bytes [0, 1024));
- the TLE buffer is addressed from a runtime offset (`offset = v << 4`, in bytes) covering `[offset, offset + 16384)`, overlapping the histogram's `[0, 1024)` region.

The histogram's zero-init and atomics therefore corrupt the TLE buffer, and the read-back returns wrong values.

**NVIDIA cross-check**: the identical kernels (Listings 7 and 8) pass on an H800 (sm90, FlagTree 0.6.0+gite04b0cb4, triton 3.6.0) at num_warps 8 and 16, 6 repetitions (24/24 checks PASS). This is a metax-specific allocation defect, not a TLE usage limitation.

**Excluded causes (verified)**:

| Variant | Result |
|---|---|
| Same kernel at BLOCK=256 (histogram / cumsum / 4 rounds / row loop) | PASS |
| Multi-CTA (2) + cooperative spin barrier + histogram | PASS |
| smem round-trip without histogram, sizes 256-8192 | PASS |
| smem round-trip with runtime masks, across sizes | PASS |
| Same computation on **global** pointers, 4 combinations | PASS |
| Listing 8 (no histogram/atomic) | PASS |
| Listing 7 (histogram/atomic) | **FAIL** |

**Layout sensitivity**: semantically equivalent variants flip between PASS and FAIL depending on unrelated compile-time structure:

- a variant with store and load masks enabled (PASS) vs an equivalent variant with no histogram and no output mask (FAIL) — identical smem accesses, different compile-time structure
- the same variant without an output mask (FAIL) vs with one (PASS)

This is consistent with the allocation overlap: whether `tl.histogram` emits shared-memory atomics (and where the TLE buffer lands) depends on the generated layout.

**Impact**: this silently corrupts results of the production kernel this was found in (a 4-round radix top-k): 130/512 selected indices wrong, uniformly spread; output indices all unique and in range, so the corruption is easy to miss. The identical algorithm with global pointers passes the full test suite (136 tests).

---

## Problem 5 (minor): masked smem load ignores `other`

Listing 9 shows the masked-out tail coming back as 0 (the buffer content) instead of the requested `other`. The same pattern is used by upstream `test_tle_gpu_local_ptr.py` (lines 96 and 159).

---

## Limitation: `buffered_tensor` cannot cross function boundaries

Passing a `tle.gpu.alloc` result to another `@triton.jit` function (Listing 10) fails with `AttributeError: 'triton._C.libtriton.ir.builder' object has no attribute 'get_memdesc_type'`. Upstream TLE has no documented/tested usage of passing a buffered tensor as an argument, so this is reported as a limitation rather than a bug — it does force all smem users into one function.

## Informational: smem capacity

Including histogram/cumsum scratch: `CHUNK = 8192` uint32 (32 KB) compiles (total `shared = 49152`); `CHUNK = 16384` (64 KB) fails with `out of resource: shared memory, Required: 69632, Hardware limit: 65536`.

---

## Reproducer index

| Listing | Problem | Note | Observed |
|---|---|---|---|
| 1 | 1 | scalar index | CRASH `LoadStoreOpToLLVM.cpp:427` |
| 2 | 1 | tensor index (equivalent access) | PASS |
| 3 | 2 | production form | CRASH `tt.atomic_rmw` verifier |
| 4 | 3 | fixed indices | PASS |
| 5 | 3 | indices depend on loop variable | CRASH 427 |
| 6 | 3 | static unroll | PASS |
| 7 | 4 | with histogram/atomic | FAIL (stable corruption) |
| 8 | 4 | without histogram/atomic | PASS |
| 9 | 5 | masked load | BUG (`other` ignored) |
| 10 | limitation | buffered tensor argument | CRASH (compilation) |

## Analysis / suspected locations

1. **Problem 1 (C vs 2c asymmetry)**: both lower to `tt.addptr(, )` followed by `tl.load`/ `tl.store` on shared-memory pointers. In the working full-view/tensor-index form the pointer tensor comes from `tle.gpu.local_pointers` (keeps the layout/encoding); in the failing scalar-index form it is `tt.splat(scalar_local_ptr)` broadcast before `tt.addptr`. The vectorization decomposition at `LoadStoreOpToLLVM.cpp:427` handles the former but not the latter, although NVIDIA supports both and they are semantically equivalent — the splat path appears to be missing.
2. **Problem 2**: `tt.atomic_rmw` on `!tt.ptr<..., 3>` fails at front-end verification (no `mctle` handling is reached).
3. **Problem 4**: the TLE buffer and `tl.histogram`'s shared-memory scratch are allocated in the same range (see "Root cause" above). The metax allocation pass does not reserve space for the histogram scratch when a TLE `gpu.alloc` is present (the total `shared` accounts only for the TLE buffer). The layout sensitivity is a consequence: whether the histogram emits shared-memory atomics depends on the chosen layout.

---

# Appendix: reproducer listings

## Listing 1 — Problem 1: scalar-indexed `local_ptr` + offset

```python
import torch
import triton
import triton.language as tl
import triton.experimental.tle.language as tle

BLOCK = 256

@triton.jit
def kernel(out_ptr, BLOCK: tl.constexpr):
buf = tle.gpu.alloc([BLOCK], dtype=tl.int32, layout=None,
scope=tle.gpu.smem, nv_mma_shared_layout=False)
base = tle.gpu.local_ptr(buf, (0,)) # scalar index
lane = tl.arange(0, BLOCK)
tl.store(base + lane, lane) # -> assertion
tl.debug_barrier()
tl.store(out_ptr + lane, tl.load(base + lane)) # -> assertion

out = torch.zeros(BLOCK, dtype=torch.int32, device="cuda")
kernel[(1,)](out, BLOCK, num_warps=4) # metax: compile-time 427 assertion
```

## Listing 2 — Problem 1: equivalent tensor-index form (works)

```python
import torch
import triton
import triton.language as tl
import triton.experimental.tle.language as tle

BLOCK = 256

@triton.jit
def kernel(out_ptr, BLOCK: tl.constexpr):
buf = tle.gpu.alloc([BLOCK], dtype=tl.int32, layout=None,
scope=tle.gpu.smem, nv_mma_shared_layout=False)
zero = tl.zeros([BLOCK], dtype=tl.int32) # tensor index (documented form)
base = tle.gpu.local_ptr(buf, (zero,)) # base[i] = &slot[0]
lane = tl.arange(0, BLOCK)
tl.store(base + lane, lane) # -> slot[lane]
tl.debug_barrier()
tl.store(out_ptr + lane, tl.load(base + lane))

out = torch.zeros(BLOCK, dtype=torch.int32, device="cuda")
kernel[(1,)](out, BLOCK, num_warps=4)
assert torch.equal(out, torch.arange(BLOCK, dtype=torch.int32, device="cuda"))
```

## Listing 3 — Problem 2: smem atomic in production form

```python
import torch
import triton
import triton.language as tl
import triton.experimental.tle.language as tle

RADIX = 256
BLOCK = 256

@triton.jit
def kernel(out_ptr, BLOCK: tl.constexpr, RADIX: tl.constexpr):
buf = tle.gpu.alloc([RADIX], dtype=tl.int32, layout=None,
scope=tle.gpu.smem, nv_mma_shared_layout=False)
base = tle.gpu.local_ptr(buf, (0,))
lane = tl.arange(0, BLOCK)
ones = tl.full([BLOCK], 1, tl.int32)
digit = (lane * 7) % RADIX
take = lane < RADIX
tl.atomic_add(base + digit, ones, mask=take,
sem="relaxed", scope="cta") # -> verifier failure
tl.debug_barrier()
tl.store(out_ptr + tl.arange(0, RADIX), tl.load(base + tl.arange(0, RADIX)))

out = torch.zeros(RADIX, dtype=torch.int32, device="cuda")
kernel[(1,)](out, BLOCK, RADIX, num_warps=4)
```

## Listing 4 — Problem 3: dynamic loop with fixed indices (works)

```python
import torch
import triton
import triton.language as tl
import triton.experimental.tle.language as tle

BLOCK = 256
NITER = 4

@triton.jit
def kernel(in_ptr, out_ptr, BLOCK: tl.constexpr, NITER: tl.constexpr):
buf = tle.gpu.alloc([BLOCK], dtype=tl.int32, layout=None,
scope=tle.gpu.smem, nv_mma_shared_layout=False)
flat = tl.arange(0, BLOCK)
for i in tl.range(NITER): # dynamic loop
ptrs = tle.gpu.local_ptr(buf, (flat,)) # indices do not depend on i
tl.store(ptrs, tl.load(in_ptr + i * BLOCK + flat))
tl.debug_barrier()
tl.store(out_ptr + i * BLOCK + flat, tl.load(ptrs))

n = BLOCK * NITER
x = torch.arange(n, dtype=torch.int32, device="cuda")
out = torch.zeros(n, dtype=torch.int32, device="cuda")
kernel[(1,)](x, out, BLOCK, NITER, num_warps=4)
assert torch.equal(out, x)
```

## Listing 5 — Problem 3: dynamic loop, indices depend on the loop variable

```python
import torch
import triton
import triton.language as tl
import triton.experimental.tle.language as tle

BLOCK = 256
NITER = 4

@triton.jit
def kernel(in_ptr, out_ptr, BLOCK: tl.constexpr, NITER: tl.constexpr):
buf = tle.gpu.alloc([BLOCK * NITER], dtype=tl.int32, layout=None,
scope=tle.gpu.smem, nv_mma_shared_layout=False)
flat = tl.arange(0, BLOCK)
for i in tl.range(NITER): # dynamic loop
idx = i * BLOCK + flat # depends on loop var
ptrs = tle.gpu.local_ptr(buf, (idx,)) # -> crash
tl.store(ptrs, tl.load(in_ptr + idx))
tl.debug_barrier()
tl.store(out_ptr + idx, tl.load(ptrs))

n = BLOCK * NITER
x = torch.arange(n, dtype=torch.int32, device="cuda")
out = torch.zeros(n, dtype=torch.int32, device="cuda")
kernel[(1,)](x, out, BLOCK, NITER, num_warps=4) # metax: crash
```

## Listing 6 — Problem 3: static unroll with the same expression (works)

```python
import torch
import triton
import triton.language as tl
import triton.experimental.tle.language as tle

BLOCK = 256
NITER = 4

@triton.jit
def kernel(in_ptr, out_ptr, BLOCK: tl.constexpr, NITER: tl.constexpr):
buf = tle.gpu.alloc([BLOCK * NITER], dtype=tl.int32, layout=None,
scope=tle.gpu.smem, nv_mma_shared_layout=False)
flat = tl.arange(0, BLOCK)
for i in tl.static_range(NITER): # unrolled; i is constexpr
idx = i * BLOCK + flat
ptrs = tle.gpu.local_ptr(buf, (idx,))
tl.store(ptrs, tl.load(in_ptr + idx))
tl.debug_barrier()
tl.store(out_ptr + idx, tl.load(ptrs))

n = BLOCK * NITER
x = torch.arange(n, dtype=torch.int32, device="cuda")
out = torch.zeros(n, dtype=torch.int32, device="cuda")
kernel[(1,)](x, out, BLOCK, NITER, num_warps=4)
assert torch.equal(out, x)
```

## Listing 7 — Problem 4: smem round-trip + histogram corrupts data

```python
import torch
import triton
import triton.language as tl
import triton.experimental.tle.language as tle

TILE = 4096
NUM_WARPS = 8
SENTINEL = -7

@triton.jit
def _convert_to_uint32_v2(x):
bits = x.to(tl.uint32, bitcast=True)
mask = ((bits >> 31) * 0x7FFFFFFF) | 0x80000000
return bits ^ mask

@triton.jit
def kernel(in_ptr, out_ptr, LEN, TILE: tl.constexpr):
buf = tle.gpu.alloc([TILE], dtype=tl.uint32, layout=None,
scope=tle.gpu.smem, nv_mma_shared_layout=False)
flat = tl.arange(0, TILE)
valid = flat < LEN
x = tl.load(in_ptr + flat, mask=valid, other=float("-inf"))
bits = _convert_to_uint32_v2(x)
tl.store(tle.gpu.local_ptr(buf, (flat,)), bits, mask=valid)
h = tl.histogram((bits >> 24).to(tl.int32), 256, mask=valid)
tl.atomic_add(out_ptr + TILE + tl.arange(0, 256), h, mask=h > 0)
tl.debug_barrier()
v = tl.load(tle.gpu.local_ptr(buf, (flat,)), mask=valid, other=0)
tl.store(out_ptr + flat, v.to(tl.int32), mask=valid)

def host_convert(x):
bits = x.view(torch.int32).to(torch.int64) & 0xFFFFFFFF
sign = bits >> 31
mask = ((sign * 0x7FFFFFFF) | 0x80000000) & 0xFFFFFFFF
return (bits ^ mask).to(torch.int64).to(torch.int32)

torch.manual_seed(0)
x = torch.randn(TILE, dtype=torch.float32, device="cuda")
out = torch.full((TILE + 256,), SENTINEL, dtype=torch.int32, device="cuda")
kernel[(1,)](x, out, TILE, TILE, num_warps=NUM_WARPS)
torch.cuda.synchronize()
got = out[:TILE]
exp = host_convert(x)
# metax: FAIL, e.g. 480-544 of 4096 values wrong, 0 never-written
print("wrong:", ((got != exp) & (got != SENTINEL)).sum().item())
```

## Listing 8 — Problem 4: same kernel without histogram/atomic (works)

Identical to Listing 7 with the two histogram lines removed:

```python
# h = tl.histogram((bits >> 24).to(tl.int32), 256, mask=valid)
# tl.atomic_add(out_ptr + TILE + tl.arange(0, 256), h, mask=h > 0)
```

Result: PASS (0 wrong values).

## Listing 9 — Problem 5: masked smem load ignores `other`

```python
import torch
import triton
import triton.language as tl
import triton.experimental.tle.language as tle

BLOCK = 256
TAIL = 10

@triton.jit
def kernel(out_ptr, BLOCK: tl.constexpr, TAIL: tl.constexpr):
buf = tle.gpu.alloc([BLOCK], dtype=tl.int32, layout=None,
scope=tle.gpu.smem, nv_mma_shared_layout=False)
ptrs = tle.gpu.local_ptr(buf)
flat = tl.arange(0, BLOCK)
tl.store(ptrs, flat + 3)
tl.debug_barrier()
mask = flat < (BLOCK - TAIL)
v = tl.load(ptrs, mask=mask, other=-1) # tail should be -1
tl.store(out_ptr + flat, v)

out = torch.zeros(BLOCK, dtype=torch.int32, device="cuda")
kernel[(1,)](out, BLOCK, TAIL, num_warps=4)
# metax: out[BLOCK-TAIL:] == 0 (expected -1), i.e. `other` is ignored
```

## Listing 10 — Limitation: `buffered_tensor` as a function argument

```python
import torch
import triton
import triton.language as tl
import triton.experimental.tle.language as tle

BLOCK = 256

@triton.jit
def worker(buf, out_ptr, BLOCK: tl.constexpr): # buf = tle.gpu.alloc(...)
ptrs = tle.gpu.local_ptr(buf)
flat = tl.arange(0, BLOCK)
tl.store(ptrs, flat)
tl.debug_barrier()
tl.store(out_ptr + flat, tl.load(ptrs))

@triton.jit
def kernel(out_ptr, BLOCK: tl.constexpr):
buf = tle.gpu.alloc([BLOCK], dtype=tl.int32, layout=None,
scope=tle.gpu.smem, nv_mma_shared_layout=False)
worker(buf, out_ptr, BLOCK) # AttributeError: ir.builder has no
# attribute 'get_memdesc_type'

out = torch.zeros(BLOCK, dtype=torch.int32, device="cuda")
kernel[(1,)](out, BLOCK, num_warps=4)
```

Start by running the self-contained Appendix listings, then inspect LoadStoreOpToLLVM.cpp:427 and the upstream python/test/tle/unit/test_tle_gpu_local_ptr.py cases. Compare scalar and tensor local-pointer lowering, atomic verification, and TLE allocation alongside tl.histogram. Done means the reported compile failures, corruption, and masked-load behavior are covered by passing reproducers without regressing working cases.

Written by the indexing model from the issue text.