Thanks for sending an requirement! Please fill in the following template to help quickly solve your problem.

Backgroud（背景信息）

aclnnBSASelectBlockMask 是 BSA（BlockSparseAttention）的前置算子，负责根据 Query 和 Key 的内容动态生成 blockSparseMask，使 BSA 的调用链从“手动提供掩码”变为“根据 Q/K 内容自适应选择稀疏模式”。

Benefit / Necessity （价值/作用）

BlockSparseAttention（BSA）算子通过外部传入的 blockSparseMask 实现稀疏注意力计算，减少长序列场景下的无效计算。当前该掩码由用户手动构造，缺乏根据 Q/K 内容自适应生成稀疏模式的能力。

Design（设计方案）

核心思想

对 Q 和 K 进行均值池化压缩，计算压缩后的 attention score（含 softmax 归一化），再在整个 head 内做 topK 选择，选出最值得计算的块位置，生成与 BSA 输入格式直接对接的 uint8 掩码。

计算公式​：

设 blockShapeOptional = [blockShapeX, blockShapeY]，压缩后块数：

$$Xblocks = \lceil Sq / blockShapeX \rceil,\quad Yblocks = \lceil Skv / blockShapeY \rceil$$

Step1：均值池化压缩 (Mean Pooling Compression)

当 actualBlockLenQueryOptional / actualBlockLenKeyOptional 为 null 时（完整压缩）：

$$q\_compressed[b, n, x, d] = \frac{1}{blockShapeX} \sum_{i=0}^{blockShapeX-1} query[b, n, x \cdot blockShapeX + i, d]$$

$$k\_compressed[b, n, y, d] = \frac{1}{blockShapeY} \sum_{j=0}^{blockShapeY-1} key[b, n, y \cdot blockShapeY + j, d]$$

当 actualBlockLenQueryOptional / actualBlockLenKeyOptional 非 null 时（部分压缩），仅对每个 block 内前 actualBlockLen 个 token 取均值：

$$q\_compressed[b, n, x, d] = \frac{1}{actualBlockLenQ[b,x]} \sum_{i=0}^{actualBlockLenQ[b,x]-1} query[b, n, x \cdot blockShapeX + i, d]$$

$$k\_compressed[b, n, y, d] = \frac{1}{actualBlockLenK[b,y]} \sum_{j=0}^{actualBlockLenK[b,y]-1} key[b, n, y \cdot blockShapeY + j, d]$$

Step2a：QK Matmul

$$score[b, n, x, y] = scale \cdot \sum_{d=0}^{D-1} q\_compressed[b, n, x, d] \cdot k\_compressed[b, n, y, d]$$

Step2b：Softmax

$$attn\_score[b, n, x, y] = softmax(score[b, n, x, :])[y] = \frac{\exp(score[b, n, x, y] - m_{final})}{l_{final}}$$

Step3：TopK 选择生成索引

$$topk\_value = \text{round}(sparsity \times Xblocks \times Yblocks)$$

$$\mathcal{indices}= \text{TopK}\left(attn\_score[b, n, x, y],\; topK\_value\right)$$

其中 indices 为 attn_score[b, n, x, y] 中 topk_value 个最大值对应的索引集合。

Step4：生成 BlockSparseMask

$$blockSparseMaskOut[b, n, x, y] = \begin{cases} 1 & (b, n, x, y) \in \mathcal{indices} \\ 0 & (b, n, x, y) \notin \mathcal{indices} \end{cases}$$