@triton.jit
def FleetAttention_with_decay_kernel(B_ptr, C_ptr, V_ptr, gamma_ptr,
    ans_ptr, heads: tl.constexpr, seqlen: tl.constexpr, dim: tl.constexpr,
    rank: tl.constexpr, stride_vbh: tl.constexpr, stride_bbh: tl.constexpr,
    dim_BLOCK: tl.constexpr):
    rank_idx = tl.program_id(axis=0)
    bz = tl.program_id(axis=1)
    dim_block_idx = tl.program_id(axis=2)
    off_b = tl.arange(0, 1)
    off_dim = tl.arange(0, dim_BLOCK)
    off_gamma = tl.full((1,), bz % heads, dtype=tl.int32)
    cv = tl.zeros([1, dim_BLOCK], dtype=tl.float32)
    o = tl.zeros([1, dim_BLOCK], dtype=tl.float32)
    gamma = tl.load(gamma_ptr + off_gamma, mask=off_gamma < heads, other=0)
    for seq_idx in range(seqlen):
        offs_bc = bz * stride_bbh + seq_idx * rank + rank_idx + off_b[None, :]
        offs_v = (bz * stride_vbh + seq_idx * dim + dim_block_idx *
            dim_BLOCK + off_dim[None, :])
        ans_ptrs = (ans_ptr + bz * stride_vbh + seq_idx * dim + 
            dim_block_idx * dim_BLOCK + off_dim[None, :])
        v_ptrs = V_ptr + offs_v
        b_ptr = B_ptr + offs_bc
        c_ptr = C_ptr + offs_bc
        b = tl.load(b_ptr, mask=off_b[None, :] < 1, other=0)
        c = tl.load(c_ptr, mask=off_b[None, :] < 1, other=0)
        v = tl.load(v_ptrs, mask=off_dim[None, :] < dim, other=0)
        cv = c * v + cv * gamma
        o = b * cv
        ans = tl.load(ans_ptrs, mask=off_dim[None, :] < dim, other=0)
        tl.store(ans_ptrs, ans + o, mask=off_dim[None, :] < dim)
