Project: 069_FlexPrefill_A_Context-Aware_Sparse_Attention_Mecha | Function: flex_prefill_attention (iclr)

system: You are an expert Python programmer implementing algorithms from scientific papers.
Complete the masked function based on the algorithm description.
Follow Python best practices and maintain consistent code style.
Return **only** the complete Python function code.

user: Complete this Python function:

**DESCRIPTION:** Main FlexPrefill algorithm implementing the complete context-aware sparse attention mechanism with dynamic pattern selection and computational budget allocation for efficient long-sequence inference

**KEY TERMS:** {
  "query": "Q query matrix",
  "key": "K key matrix",
  "value": "V value matrix",
  "pattern_threshold": "τ threshold for pattern selection",
  "coverage_threshold": "γ attention coverage requirement",
  "sparse_pattern": "selected sparse attention pattern"
}

**CONTEXT CODE:**     v_ptr,  # V: b x n x h x d
    o_ptr,
    # shape
    BATCH_SIZE,
    NUM_HEADS,
    NUM_KV_HEADS,
    NUM_SHARE_Q_HEADS,
    Q_LEN,
    K_LEN,
    HEAD_DIM: tl.constexpr,
    # softmax_scale
    softmax_scale,
    # causal
    causal,
    # gqa
    gqa_interleave,
    # stride
    stride_qb,
    stride_qn,
    stride_qh,
    stride_qd,
    stride_kb,
    stride_kn,
    stride_kh,
    stride_kd,
    stride_vb,
    stride_vn,
    stride_vh,
    stride_vd,
    stride_ob,
    stride_on,
    stride_oh,
    stride_od,
    # META parameters
    BLOCK_SIZE_Q: tl.constexpr,  # q block size
    BLOCK_SIZE_K: tl.constexpr,  # k block size
):
    raise NotImplementedError("Function implementation has been masked for testing")



from transformers.utils import is_flash_attn_2_available

if is_flash_attn_2_available():
    from flash_attn import flash_attn_func
else:
    flash_attn_func = triton_flash_attention


@torch.no_grad()
def flex_prefill_attention(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    gamma: float,
    tau: float = 0,
    min_budget: int = None,
    max_budget: int = None,
    gqa_interleave: bool = False,
    softmax_scale: Optional[float] = None,
    block_size: int = 128,
    return_computational_ratio: bool = False,
) -> Union[torch.Tensor, Tuple[torch.Tensor, float]]:
    """Flex Prefill sparse attention function. If query length is 1, will use flash decoding attention.

    Args:
        q (torch.Tensor): query tensor, shape [batch_size, q_len, num_q_heads, head_dim]
        k (torch.Tensor): key tensor, shape [batch_size, kv_len, num_kv_heads, head_dim]
        v (torch.Tensor): value tensor, shape [batch_size, kv_len, num_kv_heads, head_dim]
        gamma (float): attention coverage ratio, (0, 1).
        tau (float, optional): query aware head threshold, [0, 1]. Defaults to 0.
        min_budget (int, optional): minimum number of calculated tokens. Defaults to None.
        max_budget (int, optional): maximum number of calculated tokens. Defaults to None.
        gqa_interleave (bool, optional): GQA pattern. Defaults to False.
        softmax_scale (Optional[float], optional): softmax scale. Defaults to None, which means sqrt(head_dim).
        block_size (int, optional): block size. Defaults to 128.
        return_computational_ratio (bool, optional): whether to return computation ratio. Defaults to False.

    Returns:
        Union[torch.Tensor, Tuple[torch.Tensor, float]]: if return_computational_ratio is True, return attention output, else return attention output and computation ratio.
    """
    raise NotImplementedError("Function implementation has been masked for testing")


if __name__ == "__main__":
    torch.manual_seed(0)
    B, N, H, D = 1, 64000, 32, 128
    gamma = 0.9
    tau = 0.1

    q = torch.randn(B, N, H, D, device="cuda", dtype=torch.bfloat16) / 0.5
    k = torch.randn(B, N, H // 4, D, device="cuda", dtype=torch.bfloat16) / 0.5
    v = torch.randn(B, N, H // 4, D, device="cuda", dtype=torch.bfloat16)

    flex_prefill_output, computational_ratio = flex_prefill_attention(
        q,
        k,
        v,
        gamma,
        tau,
        min_budget=1024,
        max_budget=None,
        gqa_interleave=False,
        block_size=128,
        return_computational_ratio=True,
    )
    print("attention output norm:", flex_prefill_output.norm())
    print(f"computational ratio: {computational_ratio*100:.2f}%")


Complete the function implementation replacing "raise NotImplementedError".