fix: synchronize final decode async copy

- wait for the final split-KV tile before reading shared memory - cover long decode with production context capacity
2026-08-09 00:31:47 +08:00
parent 7f0e8bb8c2
commit a33ca04f60
2 changed files with 11 additions and 4 deletions
@@ -121,7 +121,10 @@ __global__ void attn_decode_split_kv_mma_kernel(AttentionParams<bf16> p) {
            load_tile(ti_begin + i, i);

        for (int it = 0; it < ntiles; it++) {
-            cp_async_wait_group<STAGES - 1>();
+            if (it + 1 == ntiles)
+                cp_async_wait_group<0>();
+            else
+                cp_async_wait_group<STAGES - 1>();
            __syncwarp();
            process_tile(it, it & (STAGES - 1));
            __syncwarp();