首尾元素相同的间隔循环策略

被制作时长两年半的个人练习生2026-01-14 10:22

在看Flash Attention4的过程中，发现一个很有趣的细节。以往在写输出数组等方式时，会面临一个元素和逗号间接输出的情况，且第一个元素前和第二个元素后都不包含逗号，如[3, 4, 5, 6, 7]，写过程序的应该都见过这种处理方式，一种手段是先输出一个元素，然后循环输出逗号+元素，另一种手段是先循环输出元素+逗号，然后输出输出最后一个元素。

然后我就在内联汇编里看到了如下代码：

cpp 复制代码

f"@leader_thread tcgen05.mma.cta_group::1.kind::f16 [tmem_acc], [tmem_a], smem_desc_b, idesc, {pred_str};\n\t"
            + "".join(
                (
                    # f"add.u32 tmem_a, tmem_a, {hex(offset_a_diff[k - 1])};\n\t"
                    f"add.u32 smem_desc_b_lo, smem_desc_b_lo, {hex(offset_b_diff[k - 1])};\n\t"
                    f"mov.b64 smem_desc_b, {{smem_desc_b_lo, smem_desc_b_hi}};\n\t"
                    # f"@leader_thread tcgen05.mma.cta_group::1.kind::f16 [tmem_acc], [tmem_a], smem_desc_b, idesc, 1;\n\t"
                    f"@leader_thread tcgen05.mma.cta_group::1.kind::f16 [tmem_acc], [tmem_a + {hex(offset_a[k])}], smem_desc_b, idesc, 1;\n\t"
                )
                for k in range(1, cute.size(tCrA.shape[2]) if cutlass.const_expr(mbar_ptr is None) else cute.size(tCrA.shape[2]) // 4 * 3)
            )
            ```
注意看第一个@leader_thread和随后的循环中的@leader_thread，正是采用第一种策略。