@triton.jit
def jagged_dense_flash_attention_fwd_kernel(q_ptr, k_ptr, v_ptr, ab_ptr,
    o_ptr, lse_ptr, jagged_offsets_ptr, max_seq_len, stride_ql, stride_qd,
    stride_kb, stride_kd, stride_kt, stride_vn, stride_vd, stride_ab_b,
    stride_ab_n, stride_ab_t, stride_ob, stride_ot, stride_od, D: tl.
    constexpr, T: tl.constexpr, allow_tf32: tl.constexpr, BLOCK_T: tl.
    constexpr, BLOCK_L: tl.constexpr, BLOCK_D: tl.constexpr):
    pid_t = tl.program_id(0)
    pid_batch = tl.program_id(1)
    begin = tl.load(jagged_offsets_ptr + pid_batch)
    end = tl.load(jagged_offsets_ptr + pid_batch + 1)
    length = end - begin
    length = tl.minimum(length, max_seq_len)
    if length == 0:
        return
    q_start_ptr = q_ptr + begin * stride_ql
    k_start_ptr = k_ptr + pid_batch * stride_kb
    ab_start_ptr = ab_ptr + pid_batch * stride_ab_b
    v_start_ptr = v_ptr + begin * stride_vn
    offs_t = pid_t * BLOCK_T + tl.arange(0, BLOCK_T)
    offs_d = tl.arange(0, BLOCK_D)
    ki_ptrs = k_start_ptr + offs_d[:, None] * stride_kd + offs_t[None, :
        ] * stride_kt
    ki = tl.load(ki_ptrs, mask=(offs_d[:, None] < D) & (offs_t[None, :] < T
        ), other=0.0)
    mi = tl.zeros([BLOCK_T], dtype=tl.float32) - float('inf')
    li = tl.zeros([BLOCK_T], dtype=tl.float32)
    oi = tl.zeros([BLOCK_T, BLOCK_D], dtype=tl.float32)
    for start_l in range(0, length, BLOCK_L):
        offs_l = start_l + tl.arange(0, BLOCK_L)
        qj_ptrs = q_start_ptr + offs_l[:, None] * stride_ql + offs_d[None, :
            ] * stride_qd
        qj = tl.load(qj_ptrs, mask=(offs_l[:, None] < length) & (offs_d[
            None, :] < D), other=0.0)
        qk = tl.dot(qj, ki, allow_tf32=allow_tf32)
        ab_ptrs = ab_start_ptr + offs_l[:, None] * stride_ab_n + offs_t[None, :
            ] * stride_ab_t
        abij = tl.load(ab_ptrs, mask=(offs_l[:, None] < length) & (offs_t[
            None, :] < T), other=0.0)
        qk = qk + abij
        mij_hat = tl.max(qk, axis=0)
        mi_new = tl.maximum(mi, mij_hat)
        pij_hat = tl.exp(qk - mi_new[None, :])
        pij_hat = tl.where((offs_l[:, None] < length) & (offs_t[None, :] <
            T), pij_hat, 0.0)
        lij_hat = tl.sum(pij_hat, axis=0)
        alpha = tl.exp(mi - mi_new)
        li_new = alpha * li + lij_hat
        oi = alpha[:, None] * oi
        vj_ptrs = v_start_ptr + offs_l[:, None] * stride_vn + offs_d[None, :
            ] * stride_vd
        vj = tl.load(vj_ptrs, mask=(offs_l[:, None] < length) & (offs_d[
            None, :] < D), other=0.0)
        pij_hat = pij_hat.to(v_ptr.dtype.element_ty)
        oi = oi + tl.dot(tl.trans(pij_hat), vj, allow_tf32=allow_tf32)
        mi = mi_new
        li = li_new
    oi = oi / li[:, None]
    lse_ptrs = lse_ptr + pid_batch * T + offs_t
    lse_i = mi + tl.log(li)
    tl.store(lse_ptrs, lse_i, mask=offs_t < T)
    attn_out_ptrs = o_ptr + pid_batch * stride_ob + offs_t[:, None
        ] * stride_ot + offs_d[None, :] * stride_od
    tl.store(attn_out_ptrs, oi, mask=(offs_t[:, None] < T) & (offs_d[None,
        :] < D))
