@triton.jit
def matmul_kernel_tma_persistent(a_desc_ptr, a_scale_ptr, b_desc_ptr,
    b_scale_ptr, c_desc_ptr, M, N, K, stride_a_scale_m, stride_b_scale_n,
    BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr, BLOCK_SIZE_K:
    tl.constexpr, GROUP_SIZE_M: tl.constexpr, NUM_SMS: tl.constexpr,
    output_dtype: tl.constexpr, ROW_WISE_SCALING: tl.constexpr):
    tl.inline_asm_elementwise(
        'fence.proxy.tensormap::generic.acquire.gpu [$1], 128; // $0 dummy reg'
        , '=r, l', [a_desc_ptr], dtype=tl.int32, is_pure=False, pack=1)
    tl.inline_asm_elementwise(
        'fence.proxy.tensormap::generic.acquire.gpu [$1], 128; // $0 dummy reg'
        , '=r, l', [b_desc_ptr], dtype=tl.int32, is_pure=False, pack=1)
    tl.inline_asm_elementwise(
        'fence.proxy.tensormap::generic.acquire.gpu [$1], 128; // $0 dummy reg'
        , '=r, l', [c_desc_ptr], dtype=tl.int32, is_pure=False, pack=1)
    dtype = tl.float8e4nv
    start_pid = tl.program_id(axis=0)
    num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
    num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
    k_tiles = tl.cdiv(K, BLOCK_SIZE_K)
    num_tiles = num_pid_m * num_pid_n
    tiles_per_SM = num_tiles // NUM_SMS
    if start_pid < num_tiles % NUM_SMS:
        tiles_per_SM += 1
    tile_id = start_pid - NUM_SMS
    ki = -1
    pid_m = 0
    pid_n = 0
    offs_am = 0
    offs_bn = 0
    num_pid_in_group = GROUP_SIZE_M * num_pid_n
    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
    a_scale, b_scale = load_scales(a_scale_ptr, b_scale_ptr, ROW_WISE_SCALING)
    for _ in range(0, k_tiles * tiles_per_SM):
        ki = tl.where(ki == k_tiles - 1, 0, ki + 1)
        if ki == 0:
            tile_id += NUM_SMS
            group_id = tile_id // num_pid_in_group
            first_pid_m = group_id * GROUP_SIZE_M
            group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
            pid_m = first_pid_m + tile_id % group_size_m
            pid_n = tile_id % num_pid_in_group // group_size_m
            offs_am = pid_m * BLOCK_SIZE_M
            offs_bn = pid_n * BLOCK_SIZE_N
        offs_k = ki * BLOCK_SIZE_K
        a = tl._experimental_descriptor_load(a_desc_ptr, [offs_am, offs_k],
            [BLOCK_SIZE_M, BLOCK_SIZE_K], dtype)
        b = tl._experimental_descriptor_load(b_desc_ptr, [offs_bn, offs_k],
            [BLOCK_SIZE_N, BLOCK_SIZE_K], dtype)
        accumulator = tl.dot(a, b.T, accumulator)
        if ki == k_tiles - 1:
            offs_cm = offs_am + tl.arange(0, BLOCK_SIZE_M)
            offs_cn = offs_bn + tl.arange(0, BLOCK_SIZE_N)
            accumulator = apply_scaling(accumulator, a_scale, b_scale,
                ROW_WISE_SCALING, offs_cm, offs_cn, M, N, stride_a_scale_m,
                stride_b_scale_n)
            c = accumulator.to(output_dtype)
            tl._experimental_descriptor_store(c_desc_ptr, c, [offs_am, offs_bn]
                )
            accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.
                float32)
