fix

yaoyaoding · yaoyaoding · commit 6ae3aeee0a7f · 2025-08-26T23:24:53.000-04:00
Signed-off-by: Yaoyao Ding &lt;dingyaoyao.cs@gmail.com&gt;
diff --git a/examples/attention/flash_attention_v1.py b/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 i, j: i + q_offset >= i + kv_offset,
+                init=lambda i, j: i + q_offset >= j + kv_offset,
             )
             score = score + self.where(mask, x=0.0, y=-1e6)
 

Original file line number	Diff line number	Diff line change
`@@ -183,7 +183,7 @@ def __call__(`
`183`	`183`	`mask = self.register_tensor(`
`184`	`184`	`dtype=boolean,`
`185`	`185`	`shape=[self.block_q, self.block_kv],`
`186`		`- init=lambda i, j: i + q_offset >= i + kv_offset,`
	`186`	`+ init=lambda i, j: i + q_offset >= j + kv_offset,`
`187`	`187`	`)`
`188`	`188`	`score = score + self.where(mask, x=0.0, y=-1e6)`
`189`	`189`