import torch
from flash_attn import flash_attn_varlen_func

# 1. 定义参数：2N拆分为前N（x_t）和后N（x_0）
N = 512
batch_size = 2
d_model = 1024
block_size = 64  # 左上Block Diagonal的块大小
causal_offset = 32  # 右上Offset Block Causal的偏移量
window_size = 128  # 右下Block Causal的梯形窗口大小

# 2. 构造输入（batch_size, 2N, d_model）
q = torch.randn(batch_size, 2*N, d_model, device="cuda", dtype=torch.bfloat16)
k = torch.randn(batch_size, 2*N, d_model, device="cuda", dtype=torch.bfloat16)
v = torch.randn(batch_size, 2*N, d_model, device="cuda", dtype=torch.bfloat16)

# 3. 分块标记：前N为query_segment 0，后N为query_segment 1；key按Offset/Block规则标记
query_lens = [N, N] * batch_size  # 前N（x_t）+后N（x_0）
key_lens = [N, N] * batch_size
# 屏蔽左下区域：让前N个query不关注后N个key的左下部分（仅关注右上Offset）
kv_padding_mask = torch.ones(batch_size, 2*N, device="cuda")
kv_padding_mask[:, N:] = torch.where(
    torch.arange(2*N)[N:] < (N - causal_offset), 0, 1  # 仅保留右上Offset区域的key
)

# 4. 调用FlashAttention Varlen API，配置各区域掩码规则
output = flash_attn_varlen_func(
    q, k, v,
    query_lens=query_lens,
    key_lens=key_lens,
    kv_padding_mask=kv_padding_mask,  # 屏蔽左下无注意力区域
    causal=[True, True],  # 右上/右下均启用因果，通过offset/window区分
    causal_offset=causal_offset,  # 右上Offset Block Causal的偏移量
    window_size=[None, window_size],  # 右下Block Causal的梯形窗口
    block_size=block_size,  # 左上Block Diagonal的块大小
    softmax_scale=1.0 / (d_model ** 0.5)
)