@triton.jit
def _rescale_kernel(peer_m, m, peer_l, l, peer_o, o, L, stride_oz,
    stride_oh, stride_om, stride_on, Z, H, N_CTX, BLOCK_M: tl.constexpr,
    BLOCK_DMODEL: tl.constexpr, BLOCK_N: tl.constexpr, LAST_STEP: tl.constexpr
    ):
    start_m = tl.program_id(0)
    off_hz = tl.program_id(1)
    o_offset = off_hz * stride_oh
    peer_o_block_ptr = tl.make_block_ptr(base=peer_o + o_offset, shape=(
        N_CTX, BLOCK_DMODEL), strides=(stride_om, stride_on), offsets=(
        start_m * BLOCK_M, 0), block_shape=(BLOCK_M, BLOCK_DMODEL), order=(
        1, 0))
    o_block_ptr = tl.make_block_ptr(base=o + o_offset, shape=(N_CTX,
        BLOCK_DMODEL), strides=(stride_om, stride_on), offsets=(start_m *
        BLOCK_M, 0), block_shape=(BLOCK_M, BLOCK_DMODEL), order=(1, 0))
    offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
    offs_n = tl.arange(0, BLOCK_N)
    peer_m_ptrs = peer_m + off_hz * N_CTX + offs_m
    m_ptrs = m + off_hz * N_CTX + offs_m
    peer_l_ptrs = peer_l + off_hz * N_CTX + offs_m
    l_ptrs = l + off_hz * N_CTX + offs_m
    peer_m_i = tl.load(peer_m_ptrs)
    peer_m_i = peer_m_i.to(tl.float32)
    m_i = tl.load(m_ptrs)
    m_i = m_i.to(tl.float32)
    peer_l_i = tl.load(peer_l_ptrs)
    peer_l_i = peer_l_i.to(tl.float32)
    l_i = tl.load(l_ptrs)
    l_i = l_i.to(tl.float32)
    peer_acc = tl.load(peer_o_block_ptr)
    peer_acc = peer_acc.to(tl.float32)
    acc = tl.load(o_block_ptr)
    acc = acc.to(tl.float32)
    lo = 0
    hi = N_CTX
    m_i_sync = tl.maximum(m_i, peer_m_i)
    alpha = tl.math.exp2(m_i - m_i_sync)
    peer_alpha = tl.math.exp2(peer_m_i - m_i_sync)
    acc_scale = l_i * 0 + alpha
    peer_acc_scale = peer_l_i * 0 + peer_alpha
    acc *= acc_scale[:, None]
    peer_acc *= peer_acc_scale[:, None]
    acc += peer_acc
    l_i = l_i * acc_scale + peer_l_i * peer_acc_scale
    tl.store(m_ptrs, m_i_sync)
    tl.store(l_ptrs, l_i)
    if LAST_STEP:
        acc = acc / l_i[:, None]
        L_ptrs = L + off_hz * N_CTX + offs_m
        tl.store(L_ptrs, m_i_sync / 1.44269504 + tl.math.log(l_i))
    tl.store(o_block_ptr, acc.to(tl.bfloat16))
