@triton.jit
def rotate_half_kernel(qk_seq_ptr, position_ids_ptr, qk_seq_stride,
    position_ids_batch_stride, seq_len, HEAD_DIM: tl.constexpr,
    BLOCK_HEIGHT: tl.constexpr, BLOCK_WIDTH: tl.constexpr, INV_BASE: tl.
    constexpr):
    HALF_HEAD: tl.constexpr = HEAD_DIM // 2
    STEPS_PER_ROW: tl.constexpr = HALF_HEAD // BLOCK_WIDTH
    batch_seq = tl.program_id(axis=0)
    row_blk_x_col_blk = tl.program_id(axis=1)
    row_blk = row_blk_x_col_blk // STEPS_PER_ROW
    row = row_blk * BLOCK_HEIGHT
    if BLOCK_WIDTH < HALF_HEAD:
        col_blk = row_blk_x_col_blk % STEPS_PER_ROW
        col = col_blk * BLOCK_WIDTH
    else:
        col: tl.constexpr = 0
    batch = batch_seq // seq_len
    seq = batch_seq % seq_len
    position_id = tl.load(position_ids_ptr + batch *
        position_ids_batch_stride + seq)
    freq = tl.libdevice.exp((col + tl.arange(0, BLOCK_WIDTH)).to(tl.float32
        ) * INV_BASE) * position_id
    cos = tl.cos(freq).to(tl.float32)
    sin = tl.sin(freq).to(tl.float32)
    col_offsets: tl.constexpr = tl.arange(0, BLOCK_WIDTH)
    embed_offsets = row * HEAD_DIM + col + col_offsets
    x_ptrs = qk_seq_ptr + batch_seq * qk_seq_stride + embed_offsets
    for k in range(0, BLOCK_HEIGHT):
        x = tl.load(x_ptrs).to(tl.float32)
        y = tl.load(x_ptrs + HALF_HEAD).to(tl.float32)
        out_x = x * cos - y * sin
        tl.store(x_ptrs, out_x)
        out_y = x * sin + y * cos
        tl.store(x_ptrs + HALF_HEAD, out_y)
        x_ptrs += HEAD_DIM
