Skip to content

Commit 4da3949

Browse files
[NFC] tkw.reduction -> tkw.iterate post rebase
Signed-off-by: Nicolas Vasilache <nicolasvasilache@users.noreply.github.com>
1 parent d3e92d9 commit 4da3949

File tree

1 file changed

+1
-1
lines changed

1 file changed

+1
-1
lines changed

lit_tests/kernel/wave/moe.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -271,7 +271,7 @@ def fused_moe_kernel(
271271
res_reg = tkl.Register[TOPK, B, D2, tkl.f32](0.0)
272272

273273
# fmt: off
274-
@tkw.reduction(N, init_args=[res_reg])
274+
@tkw.iterate(N, init_args=[res_reg])
275275
def repeat(acc: tkl.Register[TOPK, B, D2, tkl.f32]) -> tkl.Register[TOPK, B, D2, tkl.f32]:
276276
###
277277
# TMP_3[TOPK, B, D2] = TMP_2[TOPK, B, N:]

0 commit comments

Comments
 (0)