[https://nvbugs/5655885][fix] fix invalid instruction error in 2shot ar kernel on Ampere (NVIDIA#9394)

yilin-void · sherry-1001 · commit 87b1a4c9549f · 2025-12-16T08:36:04.000Z
Signed-off-by: Yilin Zhang &lt;18275976+yilin-void@users.noreply.github.com&gt;
diff --git a/cpp/tensorrt_llm/kernels/communicationKernels/allReduceFusionKernels.cu b/cpp/tensorrt_llm/kernels/communicationKernels/allReduceFusionKernels.cu
@@ -137,11 +137,17 @@ public:
             // corresponding CTA has not been launched.
             for (int flag_idx = blockIdx.x; flag_idx < kBarrierFlagCount; flag_idx += gridDim.x)
             {
+#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
                 asm volatile(
                     "st.global.relaxed.sys.b32 [%1], %0;" ::"r"(m_flag_value), "l"(m_target_flag + flag_idx * NRanks));
+#else
+                st_flag(m_target_flag + flag_idx * NRanks, m_flag_value);
+#endif
             }
+#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
             // Single release fence
             asm volatile("fence.release.sys;");
+#endif
 
             while (ld_flag(m_current_flag) == prev_flag(m_flag_value))
             {

Original file line number	Diff line number	Diff line change
`@@ -137,11 +137,17 @@ public:`
`137`	`137`	`// corresponding CTA has not been launched.`
`138`	`138`	`for (int flag_idx = blockIdx.x; flag_idx < kBarrierFlagCount; flag_idx += gridDim.x)`
`139`	`139`	`{`
	`140`	`+#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))`
`140`	`141`	`asm volatile(`
`141`	`142`	`"st.global.relaxed.sys.b32 [%1], %0;" ::"r"(m_flag_value), "l"(m_target_flag + flag_idx * NRanks));`
	`143`	`+#else`
	`144`	`+ st_flag(m_target_flag + flag_idx * NRanks, m_flag_value);`
	`145`	`+#endif`
`142`	`146`	`}`
	`147`	`+#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))`
`143`	`148`	`// Single release fence`
`144`	`149`	`asm volatile("fence.release.sys;");`
	`150`	`+#endif`
`145`	`151`
`146`	`152`	`while (ld_flag(m_current_flag) == prev_flag(m_flag_value))`
`147`	`153`	`{`