def flash_attn_with_kvcache(
    q,                   # (batch, seqlen_q, nheads, headdim)
    k_cache,             # (batch_cache, seqlen_cache, nheads_k, headdim) 或 Paged
    v_cache,             # 同 k_cache
    k=None,              # (batch, seqlen_new, nheads_k, headdim) 新 token
    v=None,              # 同 k
    rotary_cos=None,     # (seqlen_ro, rotary_dim/2) RoPE cos
    rotary_sin=None,     # RoPE sin
    cache_seqlens=None,  # (batch,) 或 int，当前缓存长度
    cache_batch_idx=None,# (batch,) 批次到 cache 的索引映射
    cache_leftpad=None,  # (batch,) 左侧 padding 偏移
    block_table=None,    # (batch, max_num_blocks) Paged KV 页表
    softmax_scale=None,
    causal=False,
    window_size=(-1, -1),
    softcap=0.0,
    rotary_interleaved=True,
    alibi_slopes=None,
    num_splits=0,        # 0=自动，>0=手动指定
    return_softmax_lse=False,
):

# block_table: (batch, max_num_blocks), dtype=int32
# 每个条目是一个物理页面的索引
block_table = torch.tensor([
    [0, 1, -1, -1],  # 请求 A: 页面 0, 1
    [3, -1, -1, -1], # 请求 B: 页面 3
    [2, 4, 5, 6],    # 请求 C: 页面 2, 4, 5, 6
], dtype=torch.int32, device="cuda")

# k_cache/v_cache 形状变为 (num_blocks, page_block_size, nheads_k, headdim)
k_cache = torch.randn(8, 256, 8, 128, device="cuda", dtype=torch.float16)
v_cache = torch.randn(8, 256, 8, 128, device="cuda", dtype=torch.float16)

// hopper/paged_kv.h:133-153
void load_page_table(const int n_block) {
    for (int i = 0; i < kPageEntryPerThread; ++i) {
        int const row_idx = n_block * kBlockN + row;
        int page_idx, page_offset;
        // FastDivmod: 将 row_idx 分解为页号和页内偏移
        page_idx = page_size_divmod.divmod(page_offset, row_idx + leftpad_k);
        int const page = mPageTable[page_idx];  // 查页表
        tPrPageOffset[i] = {page, page_offset};
    }
}

int num_splits_heuristic(int total_mblocks, int num_SMs, int num_n_blocks,
                         int num_m_blocks, int size_one_kv_head,
                         bool is_causal_or_local, int max_splits) {
    // 策略 1: 如果已有足够并行度，不拆分
    if (total_mblocks >= 0.8f * num_SMs) {
        // 例外：KV 超出 L2 Cache 且有足够 Q token
        int const size_l2 = 50 * 1024 * 1024;  // 假设 50MB L2
        if (size_one_kv_head > size_l2 && num_m_blocks >= num_SMs * 2) {
            return min((size_one_kv_head + size_l2 - 1) / size_l2, max_splits);
        }
        return 1;
    }

    // 策略 2: KV 块太少，拆分无意义
    if (num_n_blocks <= 4) { return 1; }

    // 策略 3: 寻找最佳 SM 利用率
    max_splits = min({max_splits, num_SMs, num_n_blocks});
    float max_efficiency = 0.f;
    for (int s = 1; s <= max_splits; s++) {
        float n_waves = float(total_mblocks * s) / num_SMs;
        float eff = n_waves / ceil(n_waves);  // wave 效率
        max_efficiency = max(max_efficiency, eff);
    }

    // 选择达到最佳效率 85% 的最小 split 数
    for (int s = 1; s <= max_splits; s++) {
        if (efficiency[s-1] >= 0.85 * max_efficiency) return s;
    }
    return 1;
}

// hopper/flash_fwd_combine_kernel.h:345-381
// 数值稳定的 online-softmax 聚合
float lse_max = max_across_splits(ts2rrLSE);
float lse_sum = 0.f;
for (int s = 0; s < num_splits; ++s) {
    float scale = exp(ts2rrLSE(s) - lse_max);
    lse_sum += scale;
    ts2rrLSE(s) = scale;  // 复用存储权重
}
float inv_sum = 1.0f / lse_sum;
for (int s = 0; s < num_splits; ++s) {
    ts2rrLSE(s) *= inv_sum;  // 归一化权重
}

from flash_attn import flash_attn_with_kvcache

# 预计算 RoPE cos/sin
rotary_dim = 128
freqs = 1.0 / (10000 ** (torch.arange(0, rotary_dim, 2).float() / rotary_dim))
t = torch.arange(max_seqlen)
freqs = torch.outer(t, freqs)
rotary_cos = torch.cos(freqs).to(device="cuda", dtype=torch.float32)
rotary_sin = torch.sin(freqs).to(device="cuda", dtype=torch.float32)

# 推理时传入
out = flash_attn_with_kvcache(
    q, k_cache, v_cache,
    k=k_new, v=v_new,
    rotary_cos=rotary_cos,
    rotary_sin=rotary_sin,
    cache_seqlens=cache_seqlens,
    causal=True,
    rotary_interleaved=False,  # GPT-NeoX 风格
)

# 大 batch（如 batch=256, heads=32）: 不需要 split
# 并行度 = 256 × 32 = 8192 >> num_SMs

# 小 batch（如 batch=1, heads=32）: 需要 split
# 并行度 = 1 × 32 = 32 < num_SMs (H100 有 132 个 SM)
# num_splits ≈ 132 / 32 ≈ 4

# 5 个活跃请求，物理 cache 可容纳 8 个
# 请求 0 → cache 位置 2
# 请求 1 → cache 位置 5
# 请求 2 → cache 位置 0
cache_batch_idx = torch.tensor([2, 5, 0], dtype=torch.int32, device="cuda")