@triton.jit
def _permute_tokens_kernel(out_ptr, in_ptr, mapping, hidden_size,
    BLOCK_SIZE: tl.constexpr):
    token_id = tl.program_id(axis=0)
    block_id = tl.program_id(axis=1)
    target_pos = tl.load(mapping + token_id)
    src_start = token_id * hidden_size + block_id * BLOCK_SIZE
    src_offsets = src_start + tl.arange(0, BLOCK_SIZE)
    src_data = tl.load(in_ptr + src_offsets)
    target_start = target_pos * hidden_size + block_id * BLOCK_SIZE
    target_offsets = target_start + tl.arange(0, BLOCK_SIZE)
    tl.store(out_ptr + target_offsets, src_data)
