fix invalid instruction error in 2shot ar kernel on Ampere

yilin-void · yilin-void · commit 3d01255faa27 · 2025-11-28T14:33:09.000+08:00
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
@@ -134,11 +134,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
`@@ -134,11 +134,17 @@ public:`
`134`	`134`	`// corresponding CTA has not been launched.`
`135`	`135`	`for (int flag_idx = blockIdx.x; flag_idx < kBarrierFlagCount; flag_idx += gridDim.x)`
`136`	`136`	`{`
	`137`	`+#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))`
`137`	`138`	`asm volatile(`
`138`	`139`	`"st.global.relaxed.sys.b32 [%1], %0;" ::"r"(m_flag_value), "l"(m_target_flag + flag_idx * NRanks));`
	`140`	`+#else`
	`141`	`+ st_flag(m_target_flag + flag_idx * NRanks, m_flag_value);`
	`142`	`+#endif`
`139`	`143`	`}`
	`144`	`+#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))`
`140`	`145`	`// Single release fence`
`141`	`146`	`asm volatile("fence.release.sys;");`
	`147`	`+#endif`
`142`	`148`
`143`	`149`	`while (ld_flag(m_current_flag) == prev_flag(m_flag_value))`
`144`	`150`	`{`