Provide well-formed TIR in moe_misc

Required for TVM compatibility after apache/tvm#16569
octoml · Mar 27, 2024 · 8e0135f · 8e0135f
1 parent 865434a
commit 8e0135f
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/python/mlc_chat/op/moe_misc.py b/python/mlc_chat/op/moe_misc.py
@@ -101,7 +101,7 @@ def topk_softmax_func(
                         with T.block("output"):
                             vj = T.axis.remap("S", [j])
                             out[vi, vj] = T.cast(
-                                T.exp(local_top_k_f32[j] - local_top_k_max[0])
+                                T.exp(local_top_k_f32[vj] - local_top_k_max[0])
                                 / (
                                     T.exp(local_top_k_f32[0] - local_top_k_max[0])
                                     + T.exp(local_top_k_f32[1] - local_top_k_max[0])