[Bug][tle][musa] unexpected modification of values in shared memory

[liuxiao0909c]

docker: registry.mthreads.com/mcconline/inference/vllm:v0.20.2-ph1-4.3.5-torch2.7.1-v1.1.0

flagtree: 0.6.0rc3+mthreads3.6

muInfo:

root@worker38073:~# muInfo
compiler: mcc
--------------------------------------------------------------------------------
device#                          0
Name:                            MTT S5000
Driver Version:                  4.3
compute capability:              3.1
...

testcase:

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

device='musa'
NUM_THREADS_PER_BLOCK = 512


@triton.jit
def tle_top_k_per_row_prefill(
    logits_ptr,
    out_indices_ptr,
    row_starts,
    row_ends,
    stride0,
    stride1,
    vocab_size,
    TOPK: tl.constexpr,
    TOPKP: tl.constexpr,
    BLOCK_SIZE: tl.constexpr,
    USE_RADIX_FINAL: tl.constexpr,
    ROW_OFFSET: tl.constexpr,
):
    NUM_FILNAL_ITEMS: tl.constexpr = 2048
    NUM_BINS: tl.constexpr = 2048
    VEC: tl.constexpr = 4

    row_id = tl.program_id(0) + ROW_OFFSET
    row_start = tl.load(row_starts + row_id)
    row_end = tl.load(row_ends + row_id)
    logits_ptr += row_id * stride0
    # float4 align
    x_off_mod = (row_id * stride0 + row_start) % VEC
    skip_elems = 0 if x_off_mod == 0 else VEC - x_off_mod
    out_indices_ptr += row_id * TOPK

    # used for histogram, indices in final sort and exclude prefix_sum
    s_histogram = tle.gpu.alloc(
        [NUM_BINS],
        dtype=tl.int32,
        layout=None,
        scope=tle.gpu.smem,
        nv_mma_shared_layout=False,
    )
    s_final_logits = tle.gpu.alloc(
        [NUM_FILNAL_ITEMS],
        dtype=tl.float32,
        layout=None,
        scope=tle.gpu.smem,
        nv_mma_shared_layout=False,
    )
    s_out_indices = tle.gpu.alloc(
        [TOPKP],
        dtype=tl.int32,
        layout=None,
        scope=tle.gpu.smem,
        nv_mma_shared_layout=False,
    )
    s_final_cnt = tle.gpu.alloc(
        [1],
        dtype=tl.int32,
        layout=None,
        scope=tle.gpu.smem,
        nv_mma_shared_layout=False,
    )
    s_threshold_bin_idx = tle.gpu.alloc(
        [1],
        dtype=tl.int32,
        layout=None,
        scope=tle.gpu.smem,
        nv_mma_shared_layout=False,
    )
    s_final_bin_size = tle.gpu.alloc(
        [1],
        dtype=tl.int32,
        layout=None,
        scope=tle.gpu.smem,
        nv_mma_shared_layout=False,
    )
    s_found_topk_values = tle.gpu.alloc(
        [1],
        dtype=tl.int32,
        layout=None,
        scope=tle.gpu.smem,
        nv_mma_shared_layout=False,
    )
    s_histogram_ptr = tle.gpu.local_ptr(s_histogram, (0,))
    s_final_logits_ptr = tle.gpu.local_ptr(s_final_logits, (0,))
    s_out_indices_ptr = tle.gpu.local_ptr(s_out_indices, (0,))
    s_final_cnt_ptr = tle.gpu.local_ptr(s_final_cnt, (0,))
    s_threshold_bin_idx_ptr = tle.gpu.local_ptr(s_threshold_bin_idx, (0,))
    s_final_bin_size_ptr = tle.gpu.local_ptr(s_final_bin_size, (0,))
    s_found_topk_values_ptr = tle.gpu.local_ptr(s_found_topk_values, (0,))

    tl.debug_barrier()
    tl.store(s_histogram_ptr + tl.arange(0, 2048), 0)  # set 0 for all elems in s_histogram
    tl.debug_barrier()
    lane = tl.arange(0, BLOCK_SIZE)
    vec = tl.arange(0, VEC)
    base = BLOCK_SIZE * VEC + lane * VEC
    offs = base[:, None] + vec[None, :]
    x_vec = tl.load(logits_ptr + offs)
    tl.debug_barrier()
    my_counts1 = tl.load(s_histogram_ptr + tl.arange(0, 2048)) # expect all 0
    tl.debug_barrier()
    h = x_vec.to(tl.uint32, bitcast=True)
    tl.debug_barrier()
    my_counts2 = tl.load(s_histogram_ptr + tl.arange(0, 2048)) # expect all 0
    tl.debug_barrier()
    tl.device_print("h:", h)
    tl.device_print("my_counts1:", my_counts1)
    tl.device_print("my_counts2:", my_counts2)


def top_k_per_row_prefill(
    logits, row_starts, row_ends, indices, num_rows, stride0, stride1, top_k
):
    vocab_size = logits.shape[1]
    assert num_rows == logits.shape[0]
    topkp = triton.next_power_of_2(top_k)
    tle_top_k_per_row_prefill[(num_rows,)](
        logits,
        indices,
        row_starts,
        row_ends,
        stride0,
        stride1,
        vocab_size,
        TOPK=top_k,
        TOPKP=topkp,
        BLOCK_SIZE=NUM_THREADS_PER_BLOCK,
        USE_RADIX_FINAL=False,
        ROW_OFFSET=0,
        num_warps=NUM_THREADS_PER_BLOCK // 32,
    )


torch.manual_seed(789)
num_rows = 1
vocab_size = 129280
top_k = 1024

logits = torch.randn(num_rows, vocab_size, device=device, dtype=torch.float32)
row_starts = torch.zeros(num_rows, dtype=torch.int32, device=device)
row_ends = torch.full((num_rows,), vocab_size, dtype=torch.int32, device=device)
stride0 = logits.stride(0)
stride1 = logits.stride(1)
indices_test = torch.empty((num_rows, top_k), dtype=torch.int32, device=device)

top_k_per_row_prefill(
    logits, row_starts, row_ends, indices_test, num_rows, stride0, stride1, top_k
)

testcase description: Allocate 2048 int32 elements of shared memory for s_histogram in kernel tle_top_k_per_row_prefill, initialize them to zero via tl.store. Later, all elements from s_histogram are loaded into my_counts1 and my_counts2. However, when an additional statement h = x_vec.to(tl.uint32, bitcast=True) is inserted between these two load operations, my_counts2 reveals that the data in s_histogram has been modified.

output:

[Kylin1207]

I used tl.store + torch.equal as a replacement for device_print for comparison. The error should be caused by device_print printing invalid data. The tle interface itself should be fine and can continue to be used. For future debugging, I suggest first using the torch tensor approach. I will continue to follow up on the device_print error.

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

device = 'musa'
NUM_THREADS_PER_BLOCK = 512
NUM_BINS = 2048


@triton.jit
def kernel_check_smem(
    logits_ptr,
    out_h_ptr,
    out_counts1_ptr,
    out_counts2_ptr,
    stride0,
    TOPK: tl.constexpr,
    TOPKP: tl.constexpr,
    BLOCK_SIZE: tl.constexpr,
    ROW_OFFSET: tl.constexpr,
    INIT_VALUE: tl.constexpr,
    ENABLE_PRINT: tl.constexpr,
):
    NUM_BINS: tl.constexpr = 2048
    VEC: tl.constexpr = 4

    row_id = tl.program_id(0) + ROW_OFFSET
    logits_ptr += row_id * stride0

    # Reproduce the exact same shared-memory allocation pattern as the
    # original tle_top_k_per_row_prefill kernel.
    s_histogram = tle.gpu.alloc(
        [NUM_BINS], dtype=tl.int32, layout=None,
        scope=tle.gpu.smem, nv_mma_shared_layout=False,
    )
    s_final_logits = tle.gpu.alloc(
        [NUM_BINS], dtype=tl.float32, layout=None,
        scope=tle.gpu.smem, nv_mma_shared_layout=False,
    )
    s_out_indices = tle.gpu.alloc(
        [TOPKP], dtype=tl.int32, layout=None,
        scope=tle.gpu.smem, nv_mma_shared_layout=False,
    )
    s_final_cnt = tle.gpu.alloc(
        [1], dtype=tl.int32, layout=None,
        scope=tle.gpu.smem, nv_mma_shared_layout=False,
    )
    s_threshold_bin_idx = tle.gpu.alloc(
        [1], dtype=tl.int32, layout=None,
        scope=tle.gpu.smem, nv_mma_shared_layout=False,
    )
    s_final_bin_size = tle.gpu.alloc(
        [1], dtype=tl.int32, layout=None,
        scope=tle.gpu.smem, nv_mma_shared_layout=False,
    )
    s_found_topk_values = tle.gpu.alloc(
        [1], dtype=tl.int32, layout=None,
        scope=tle.gpu.smem, nv_mma_shared_layout=False,
    )
    s_histogram_ptr = tle.gpu.local_ptr(s_histogram, (0,))

    tl.debug_barrier()
    if INIT_VALUE == -1:
        # sentinel: use a known non-zero pattern to make corruption obvious
        tl.store(s_histogram_ptr + tl.arange(0, NUM_BINS),
                 tl.arange(0, NUM_BINS).to(tl.int32) * 7 + 3)
    else:
        tl.store(s_histogram_ptr + tl.arange(0, NUM_BINS), 0)
    tl.debug_barrier()

    lane = tl.arange(0, BLOCK_SIZE)
    vec = tl.arange(0, VEC)
    base = BLOCK_SIZE * VEC + lane * VEC
    offs = base[:, None] + vec[None, :]
    x_vec = tl.load(logits_ptr + offs)
    tl.debug_barrier()

    my_counts1 = tl.load(s_histogram_ptr + tl.arange(0, NUM_BINS))
    tl.debug_barrier()

    h = x_vec.to(tl.uint32, bitcast=True)
    tl.debug_barrier()

    my_counts2 = tl.load(s_histogram_ptr + tl.arange(0, NUM_BINS))
    tl.debug_barrier()

    out_h_ptr += row_id * (BLOCK_SIZE * VEC)
    h_out_offs = lane[:, None] * VEC + vec[None, :]
    tl.store(out_h_ptr + h_out_offs, h)

    out_counts1_ptr += row_id * NUM_BINS
    out_counts2_ptr += row_id * NUM_BINS
    tl.store(out_counts1_ptr + tl.arange(0, NUM_BINS), my_counts1)
    tl.store(out_counts2_ptr + tl.arange(0, NUM_BINS), my_counts2)

    if ENABLE_PRINT:
        tl.device_print("h:", h)
        tl.device_print("my_counts1:", my_counts1)
        tl.device_print("my_counts2:", my_counts2)


def run_variant(logits, init_value, enable_print):
    num_rows = 1
    top_k = 1024

    h_out = torch.full(
        (num_rows, NUM_THREADS_PER_BLOCK * 4),
        -123456789,
        dtype=torch.int32,
        device=device,
    )
    counts1 = torch.full(
        (num_rows, NUM_BINS), -123456789, dtype=torch.int32, device=device
    )
    counts2 = torch.full(
        (num_rows, NUM_BINS), -123456789, dtype=torch.int32, device=device
    )

    kernel_check_smem[(num_rows,)](
        logits,
        h_out,
        counts1,
        counts2,
        logits.stride(0),
        TOPK=top_k,
        TOPKP=triton.next_power_of_2(top_k),
        BLOCK_SIZE=NUM_THREADS_PER_BLOCK,
        ROW_OFFSET=0,
        INIT_VALUE=init_value,
        ENABLE_PRINT=enable_print,
        num_warps=NUM_THREADS_PER_BLOCK // 32,
    )
    torch.musa.synchronize()
    return h_out.cpu(), counts1.cpu(), counts2.cpu()


def run_case(init_value, expected, label):
    num_rows = 1
    vocab_size = 129280
    vec = 4

    # Use exactly the same logits for the no-print and print-enabled variants.
    logits = torch.randn(num_rows, vocab_size, device=device, dtype=torch.float32)
    h_expected = (
        logits[:, NUM_THREADS_PER_BLOCK * vec : 2 * NUM_THREADS_PER_BLOCK * vec]
        .cpu()
        .contiguous()
        .view(torch.int32)
    )

    h_no_print, c1_no_print, c2_no_print = run_variant(
        logits, init_value, enable_print=False
    )
    h_with_print, c1_with_print, c2_with_print = run_variant(
        logits, init_value, enable_print=True
    )

    no_print_ok = (
        torch.equal(h_no_print, h_expected)
        and torch.equal(c1_no_print, expected)
        and torch.equal(c2_no_print, expected)
    )
    with_print_ok = (
        torch.equal(h_with_print, h_expected)
        and torch.equal(c1_with_print, expected)
        and torch.equal(c2_with_print, expected)
    )

    print(f"\n[{label}]")
    print(f"  no print : {'PASS' if no_print_ok else 'FAIL'}")
    print(f"  print    : {'PASS' if with_print_ok else 'FAIL'}")


def main():
    torch.manual_seed(789)

    # Case 1: zero initialization (the exact scenario from the original report)
    expected_zero = torch.zeros((1, NUM_BINS), dtype=torch.int32)
    run_case(0, expected_zero, "Case 1: zero-init (original report)")

    # Case 2: known non-zero pattern i*7+3 — makes any corruption obvious
    expected_pat = (torch.arange(NUM_BINS, dtype=torch.int32) * 7 + 3).unsqueeze(0)
    run_case(-1, expected_pat, "Case 2: known pattern i*7+3")


if __name__ == "__main__":
    main()