fix

ROCm · Jan 29, 2025 · 1cbec84 · 1cbec84
1 parent e80594a
commit 1cbec84
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/python/perf-kernels/MLA_decode.py b/python/perf-kernels/MLA_decode.py
@@ -273,7 +273,7 @@ def _fwd_kernel_stage2(
         mask_out = offs_n + n * BLOCK_SIZE_N < Lv
         w_vc = tl.load(w_kv_prts, mask=mask_v, other=0.0) # dc, head is parallelized (128, 512)
 
-        result += tl.sum(w_vc * acc[None, :], 1)
+        result = tl.sum(w_vc * acc[None, :], 1)
 
         w_kv_prts += BLOCK_SIZE_N * stride_w_vcd