def METHOD_NAME(ctx, q, k, v, mask, causal, q_bucket_size, k_bucket_size):
    """Algorithm 2 in the paper"""
    device = q.device
    dtype = q.dtype
    max_neg_value = -torch.finfo(q.dtype).max
    qk_len_diff = max(k.shape[-2] - q.shape[-2], 0)
    o = torch.zeros_like(q)
    all_row_sums = torch.zeros((*q.shape[:-1], 1), dtype=dtype, device=device)
    all_row_maxes = torch.full(
        (*q.shape[:-1], 1), max_neg_value, dtype=dtype, device=device
    )
    scale = q.shape[-1] ** -0.5
    if mask is None:
        mask = (None,) * math.ceil(q.shape[-2] / q_bucket_size)
    else:
        mask = rearrange(mask, "b n -> b 1 1 n")
        mask = mask.split(q_bucket_size, dim=-1)
    row_splits = zip(
        q.split(q_bucket_size, dim=-2),
        o.split(q_bucket_size, dim=-2),
        mask,
        all_row_sums.split(q_bucket_size, dim=-2),
        all_row_maxes.split(q_bucket_size, dim=-2),
    )
    for ind, (qc, oc, row_mask, row_sums, row_maxes) in enumerate(row_splits):
        q_start_index = ind * q_bucket_size - qk_len_diff
        col_splits = zip(
            k.split(k_bucket_size, dim=-2),
            v.split(k_bucket_size, dim=-2),
        )
        for k_ind, (kc, vc) in enumerate(col_splits):
            k_start_index = k_ind * k_bucket_size
            attn_weights = (
                torch.einsum("... i d, ... j d -> ... i j", qc, kc) * scale
            )
            if row_mask is not None:
                attn_weights.masked_fill_(~row_mask, max_neg_value)
            if causal and q_start_index < (k_start_index + k_bucket_size - 1):
                causal_mask = torch.ones(
                    (qc.shape[-2], kc.shape[-2]), dtype=torch.bool, device=device
                ).triu(q_start_index - k_start_index + 1)
                attn_weights.masked_fill_(causal_mask, max_neg_value)
            block_row_maxes = attn_weights.amax(dim=-1, keepdims=True)
            attn_weights -= block_row_maxes
            exp_weights = torch.exp(attn_weights)
            if row_mask is not None:
                exp_weights.masked_fill_(~row_mask, 0.0)
            block_row_sums = exp_weights.sum(dim=-1, keepdims=True).clamp(
                min=EPSILON
            )
            new_row_maxes = torch.maximum(block_row_maxes, row_maxes)
            exp_values = torch.einsum(
                "... i j, ... j d -> ... i d", exp_weights, vc
            )
            exp_row_max_diff = torch.exp(row_maxes - new_row_maxes)
            exp_block_row_max_diff = torch.exp(block_row_maxes - new_row_maxes)
            new_row_sums = (
                exp_row_max_diff * row_sums
                + exp_block_row_max_diff * block_row_sums
            )
            oc.mul_((row_sums / new_row_sums) * exp_row_max_diff).add_(
                (exp_block_row_max_diff / new_row_sums) * exp_values
            )
            row_maxes.copy_(new_row_maxes)
            row_sums.copy_(new_row_sums)
    ctx.args = (causal, scale, mask, q_bucket_size, k_bucket_size)
    ctx.save_for_backward(q, k, v, o, all_row_sums, all_row_maxes)
    return o