add example

Signed-off-by: Yaoyao Ding <[email protected]>

Author: yaoyaoding

examples/attention/flash_attention_v1.py

@@ -183,7 +183,7 @@ def __call__(
             mask = self.register_tensor(
                 dtype=boolean,
                 shape=[self.block_q, self.block_kv],
-                init=lambda ij: ij[0] + q_offset >= ij[1] + kv_offset,
+                init=lambda i, j: i + q_offset >= i + kv_offset,
             )
             score = score + self.where(mask, x=0.0, y=-1e6)