@triton.autotune(list(filter(keep, configsTmaWSPersistent)), key=['N_CTX'])
@triton.jit
def _attn_fwd_tma_ws_persistent(Q, K, V, sm_scale, M, Out, desc_q, desc_k,
    desc_v, desc_o, stride_qz, stride_qh, stride_qm, stride_qk, stride_kz,
    stride_kh, stride_kn, stride_kk, stride_vz, stride_vh, stride_vk,
    stride_vn, stride_oz, stride_oh, stride_om, stride_on, Z, H, N_CTX,
    BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, HEAD_DIM: tl.constexpr,
    STAGE: tl.constexpr, ENABLE_TMA: tl.constexpr, LOOP_SCHEDULE: tl.
    constexpr, ENABLE_WS: tl.constexpr, GRID_MULTIPLE: tl.constexpr):
    tl.static_assert(BLOCK_N <= HEAD_DIM)
    n_tile_num = tl.cdiv(N_CTX, BLOCK_M)
    prog_id = tl.program_id(0)
    num_progs = tl.num_programs(0)
    total_tiles = n_tile_num * Z * H
    tiles_per_sm = total_tiles // num_progs
    if prog_id < total_tiles % num_progs:
        tiles_per_sm += 1
    tile_idx = prog_id
    for _ in range(0, tiles_per_sm):
        pid = tile_idx % n_tile_num
        off_hz = tile_idx // n_tile_num
        _attn_fwd_compute_ws(Q, K, V, sm_scale, M, Out, desc_q, desc_k,
            desc_v, desc_o, stride_qz, stride_qh, stride_qm, stride_qk,
            stride_kz, stride_kh, stride_kn, stride_kk, stride_vz,
            stride_vh, stride_vk, stride_vn, stride_oz, stride_oh,
            stride_om, stride_on, off_hz, pid, Z, H, N_CTX, BLOCK_M,
            BLOCK_N, HEAD_DIM, STAGE, ENABLE_TMA, LOOP_SCHEDULE)
        tile_idx += num_progs
