@triton.jit
def stabilize_fi(F, I, M, F_STABILIZED, I_STABILIZED, NH: tl.constexpr, S:
    tl.constexpr):
    bh_id = tl.program_id(0)
    batch_id = bh_id // NH
    head_id = bh_id % NH
    s_range = tl.arange(0, S)
    batch_offset_fi = batch_id * NH * S + head_id * S
    f = tl.load(F + batch_offset_fi + s_range, s_range < S)
    i = tl.load(I + batch_offset_fi + s_range, s_range < S)
    f = tl.log(tl.sigmoid(f))
    _, m = tl.associative_scan((f, i), 0, stabilization_scan_op)
    m_shifted = roll(m, 0)
    i = tl.exp(i - m)
    f = tl.exp(f - m + m_shifted)
    tl.store(F_STABILIZED + batch_offset_fi + s_range, f, s_range < S)
    tl.store(I_STABILIZED + batch_offset_fi + s_range, i, s_range < S)
    tl.store(M + batch_offset_fi + s_range, m, s_range < S)
