From 0443e3179fa9feada0b90e03be8567218fedad47 Mon Sep 17 00:00:00 2001
From: William Hu <67835823+willhu-jpg@users.noreply.github.com>
Date: Wed, 16 Sep 2026 21:37:19 -0400
Subject: [PATCH] fix(kda_prefill): fence shared writes before async proxy
 reads (#39124)

---
 python/sglang/kernels/jit/csrc/attention/kda_prefill.cu | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/python/sglang/kernels/jit/csrc/attention/kda_prefill.cu b/python/sglang/kernels/jit/csrc/attention/kda_prefill.cu
index f1b467680..3b84b658a 100644
--- a/python/sglang/kernels/jit/csrc/attention/kda_prefill.cu
+++ b/python/sglang/kernels/jit/csrc/attention/kda_prefill.cu
@@ -150,6 +150,9 @@ static __device__ __forceinline__ uint32_t ld_acq_b32(const uint32_t* ptr) {
 static __device__ __forceinline__ void fence_async_global() {
   asm volatile("fence.proxy.async.global;");
 }
+static __device__ __forceinline__ void fence_async_shared() {
+  asm volatile("fence.proxy.async.shared::cta;" ::: "memory");
+}
 
 // ---- tcgen05 (TMEM lifecycle, ld/st, MMA, fences) ----
 static __device__ __forceinline__ void tcgen05_alloc(uint32_t smem_addr_for_taddr, uint32_t n_cols) {
@@ -1524,6 +1527,7 @@ __device__ static void chain_body(
       // after the TMA through mb_pre)
       ptx::mbar_wait_parity(&S.mb_pre, q2 & 1);
       widen_map_tf32(bL, pL, vc, ch);
+      ptx::fence_async_shared();
       ptx::tcgen05_fence_before_thread_sync();
       __syncthreads();
       if (tid == 0) {