[None][perf] optimize DSV4 FP8 o_a decode epilogue

Signed-off-by: Mingyang Hao <200044211+mingyangHao@users.noreply.github.com>

Author: mingyangHao

tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py

@@ -3821,6 +3821,9 @@ def forward_fp8out(
             sf_n = ceil_div(n, 128)
             aligned_mn = pad_up(m, 4)
             n_tiles_per_group = ceil_div(n, 128)
+            # The cooperative SMEM epilogue wins only for the single-token
+            # decode shape; retain the register epilogue for every other M.
+            use_fp8_smem_epilogue = m == 1
 
             if c_tensor.dtype != torch.float8_e4m3fn or c_tensor.shape != (
                     batch_size, m, n):
@@ -3840,6 +3843,7 @@ def forward_fp8out(
                 mma_tiler_mn,
                 cluster_shape_mn,
                 self.use_tvm_ffi,
+                use_fp8_smem_epilogue,
             )
             if cache_key not in self.__class__.kernel_cache:
                 a_ptr = make_ptr(cutlass.Float8E4M3FN,
@@ -3894,6 +3898,7 @@ def forward_fp8out(
                     sf_out_ptr,
                     aligned_mn,
                     n_tiles_per_group,
+                    use_fp8_smem_epilogue,
                     options="--opt-level 2 --enable-tvm-ffi",
                 )
                 self.__class__.kernel_cache[cache_key] = compiled_gemm