@triton.jit
def wkv_triton_log_space_forward_kernel(w_ptr, w_s_c, u_ptr, u_s_c, k_ptr,
    k_s_b, k_s_t, k_s_c, v_ptr, v_s_b, v_s_t, v_s_c, state_ptr, state_s_b,
    state_s_abe, state_s_c, wkv_ptr, wkv_s_b, wkv_s_t, wkv_s_c,
    state_out_ptr, state_out_s_b, state_out_s_abe, state_out_s_t,
    state_out_s_c, chans, tsz, eps: tl.constexpr, log_eps: tl.constexpr,
    normalize: tl.constexpr, BLOCK_SIZE_C: tl.constexpr):
    b_idx = tl.program_id(0)
    c_idx = tl.program_id(1)
    cs = c_idx * BLOCK_SIZE_C + tl.arange(0, BLOCK_SIZE_C)
    cmask = cs < chans
    k_ptr = k_ptr + b_idx * k_s_b
    v_ptr = v_ptr + b_idx * v_s_b
    ln_alpha_p_ptr = state_ptr + b_idx * state_s_b
    ln_alpha_m_ptr = state_ptr + b_idx * state_s_b + state_s_abe
    ln_beta_ptr = state_ptr + b_idx * state_s_b + 2 * state_s_abe
    wkv_ptr = wkv_ptr + b_idx * wkv_s_b
    ln_alpha_p_out_ptr = state_out_ptr + b_idx * state_out_s_b
    ln_alpha_m_out_ptr = (state_out_ptr + b_idx * state_out_s_b +
        state_out_s_abe)
    ln_beta_out_ptr = (state_out_ptr + b_idx * state_out_s_b + 2 *
        state_out_s_abe)
    ln_alpha_p = tl.load(ln_alpha_p_ptr + cs * state_s_c, mask=cmask).to(tl
        .float32)
    ln_alpha_m = tl.load(ln_alpha_m_ptr + cs * state_s_c, mask=cmask).to(tl
        .float32)
    ln_beta = tl.load(ln_beta_ptr + cs * state_s_c, mask=cmask).to(tl.float32)
    w = tl.load(w_ptr + cs * w_s_c, mask=cmask).to(tl.float32)
    u = tl.load(u_ptr + cs * u_s_c, mask=cmask).to(tl.float32)
    for t in range(tsz):
        kt = tl.load(k_ptr + t * k_s_t + cs * k_s_c, mask=cmask).to(tl.float32)
        vt = tl.load(v_ptr + t * v_s_t + cs * v_s_c, mask=cmask).to(tl.float32)
        vt_p = tl.maximum(vt, 0) + eps
        vt_m = tl.maximum(-vt, 0) + eps
        ln_v_p = tl.log(vt_p)
        ln_v_m = tl.log(vt_m)
        if normalize:
            ln_alpha_pm = tl.minimum(ln_alpha_p, ln_alpha_m) - eps
            ln_alpha_p = logsubexp(ln_alpha_p, ln_alpha_pm, log_eps)
            ln_alpha_m = logsubexp(ln_alpha_m, ln_alpha_pm, log_eps)
        ln_wkv_p = logaddexp(u + kt + ln_v_p, ln_alpha_p) - logaddexp(u +
            kt, ln_beta)
        ln_wkv_m = logaddexp(u + kt + ln_v_m, ln_alpha_m) - logaddexp(u +
            kt, ln_beta)
        wkv = tl.exp(ln_wkv_p) - tl.exp(ln_wkv_m)
        tl.store(wkv_ptr + t * wkv_s_t + cs * wkv_s_c, wkv, mask=cmask)
        ln_alpha_p = logaddexp(w + ln_alpha_p, kt + ln_v_p)
        ln_alpha_m = logaddexp(w + ln_alpha_m, kt + ln_v_m)
        ln_beta = logaddexp(w + ln_beta, kt)
        tl.store(ln_alpha_p_out_ptr + t * state_out_s_t + cs *
            state_out_s_c, ln_alpha_p, mask=cmask)
        tl.store(ln_alpha_m_out_ptr + t * state_out_s_t + cs *
            state_out_s_c, ln_alpha_m, mask=cmask)
        tl.store(ln_beta_out_ptr + t * state_out_s_t + cs * state_out_s_c,
            ln_beta, mask=cmask)
