ct.load 4-D strided slice (shape=(1,1,k,1)) silently returns wrong values; scalar, 2-D row, 1-D flat, and ct.gather forms of the same elements are correct #116

Summary

ct.load of a 4-D strided slice — ct.load(A, index=(h, m, base, 0), shape=(1, 1, k, 1)).reshape((k,)) on an int32 tensor — returns wrong values (silently: the kernel compiles and runs with no error). Reading the exact same elements through any other documented form — k scalar (1,1,1,1) loads, a 2-D (1,k) row load from a zero-copy view, a 1-D flat load, or ct.gather — is correct. The defect appears confined to this specific 4-D strided-slice lowering path, and its manifestation differs by arch and vector width.

Environment

• cuda.tile (pip cuda-tile) 1.7.0, tileiras 13.4.92, CUDA 13.0 / 13.4
• GPUs: NVIDIA GB10 (sm_121, CUDA 13.4 container, aarch64) and NVIDIA GB200 (sm_100, CUDA 13.0, aarch64)

Observed results

Loading k consecutive elements of an int32 tensor A of shape (BNH, NM, SLOTS, 1) at A[h, m, base : base+k, 0]:

On sm_121 (GB10): the 4-D form is wrong at k=2 (returns values from a different offset — e.g. [204, 205] where the truth is [198, 199]) and returns zeros at k=4 and k=8; correct at k=8 on this arch in one run, but see sm_100.

On sm_100 (GB200): the 4-D form is wrong at k=2, k=4 AND k=8 (huge relative diffs downstream; verified bit-level at the load in a probe that compares all five forms in a single kernel launch).

Two properties that make this nasty in practice:

1. It is silent — no compile or runtime error; downstream math just produces garbage.
2. Consumers of the miscompiled vector (broadcast arithmetic, ct.extract) are all equally wrong, which initially looks like a consumer-side bug; comparing against scalar loads of the same addresses in the same launch pins it to the load.

Minimal repro

One kernel, one launch, five constructions of the same k consecutive elements, results compared on the host (the "vec4d" slot is the broken form; the other non-reference slots are the working alternatives):

import cuda.tile as ct
import torch

BNH, N_M, MAX_SLOT = 2, 4, 32
KS = (2, 4, 8)
N_CON = 5  # ref, vec4d, flat1d, row2d, gather


@ct.kernel
def probe(ACTIVE_IDX, flat, a2d, OUT,
          h: ct.Constant[ct.int32], m: ct.Constant[ct.int32],
          grp: ct.Constant[ct.int32], K: ct.Constant[ct.int32],
          MAX: ct.Constant[ct.int32], N_M_C: ct.Constant[ct.int32]):
    base: ct.constexpr = grp * K
    row: ct.constexpr = h * N_M_C + m
    # 0: reference — k scalar loads (known good)
    for j in ct.static_iter(range(K)):
        v = ct.load(ACTIVE_IDX, (h, m, base + j, 0), shape=(1, 1, 1, 1)).reshape(())
        ct.store(OUT, (0 * K + j,), v.astype(ct.int32))
    # 1: 4-D strided slice (BROKEN)
    vec4d = ct.load(ACTIVE_IDX, (h, m, base, 0), shape=(1, 1, K, 1)).reshape((K,))
    ct.store(OUT, (1 * K // K,), vec4d.astype(ct.int32))
    # 2: 1-D flat view, tile-aligned tile-space index (correct)
    j_tile = row * (MAX // K) + grp
    flat_vec = ct.load(flat, (j_tile,), shape=(K,))
    ct.store(OUT, (2,), flat_vec.astype(ct.int32))
    # 3: 2-D (rows, MAX) view, (1, K) row tile (correct)
    row_vec = ct.load(a2d, (row, grp), shape=(1, K)).reshape((K,))
    ct.store(OUT, (3,), row_vec.astype(ct.int32))
    # 4: ct.gather on the 1-D flat view (correct)
    idx = (row * MAX + base) + ct.arange(K, dtype=ct.int32)
    g_vec = ct.gather(flat, idx)
    ct.store(OUT, (4,), g_vec.astype(ct.int32))


def main():
    A = torch.arange(BNH * N_M * MAX_SLOT, dtype=torch.int32, device="cuda"
                     ).reshape(BNH, N_M, MAX_SLOT, 1).contiguous()
    flat = A.view(-1)
    a2d = A.reshape(BNH * N_M, MAX_SLOT)
    names = ["ref", "vec4d", "flat1d", "row2d", "gather"]
    for K in KS:
        OUT = torch.zeros(N_CON * K, dtype=torch.int32, device="cuda")
        h, m, grp = 1, 2, 3
        base = grp * K
        truth = A[h, m, base:base + K, 0].tolist()
        ct.launch(torch.cuda.current_stream(), (1, 1, 1), probe,
                  (A, flat, a2d, OUT, h, m, grp, K, MAX_SLOT, N_M))
        torch.cuda.synchronize()
        got = OUT.view(N_CON, K).cpu()
        for c, name in enumerate(names):
            ok = got[c].tolist() == truth
            print(f"K={K} {name}: {'PASS' if ok else 'FAIL'} "
                  f"got={got[c].tolist()} truth={truth}")


if __name__ == "__main__":
    main()

Expected (GB10, sm_121): vec4d FAIL at every K ([204, 205] vs [198, 199] at K=2; zeros at K=4/8), all other constructions PASS. On GB200 (sm_100), vec4d FAIL at every K as well.

Workaround

Use a different construction for multi-element reads: we ship the 2-D (1,k) row-load form (host passes a zero-copy view(-1, slots) of the same buffer) with latency=1, allow_tma=False on the small tile — with defaults, the tiny load additionally pays per-iteration TMA/pipeline overhead that dwarfs 8-32 bytes of traffic. With the 2-D form our kernels are bit-identical against scalar-load references on both archs.

Question

Is the 4-D strided-slice (shape=(1,1,k,1)-style) load path known-broken or untested in tileiras 13.4.92? Happy to provide full environment details, SASS dumps, or a larger repro if useful.