fix: fix fmha_v2 sync bug

Signed-off-by: Mingyang Jiang <[email protected]>

Author: jmydurant

cpp/kernels/fmha_v2/src/fmha/hopper/fragment.h

@@ -486,7 +486,7 @@ struct Softmax_saver_tma
         int lane = threadIdx.x % Cta_tile::THREADS_PER_WARP;
         if (lane % 4 < 2)
         {
-            values = p_sum[lane % 2] == 0.f ? 1.f : 1.0f / p_sum[lane % 2];
+            values = p_sum[lane % 2];
         }
         else
         {

cpp/tensorrt_llm/kernels/mlaChunkedPrefill.cu

@@ -159,8 +159,8 @@ inline __device__ void dequantCopy(
 // merged_softmax_sum [q_total_len, H, 2] (float, max/sum)
 template <typename T>
 __global__ void mergeAttnWithSoftmaxKernel(T* merged_attn, float2* merged_softmax_stats, T const* pre_attn,
-    float2 const* pre_softmax_stats, T const* curr_attn, float2 const* curr_softmax_stats, float bmm1_scale,
-    int64_t const* cu_q_seq_len, int64_t const* merge_op, int const num_heads, int const head_size)
+    float2 const* pre_softmax_stats, T const* curr_attn, float2 const* curr_softmax_stats, int64_t const* cu_q_seq_len,
+    int64_t const* merge_op, int const num_heads, int const head_size)
 {
     using KT = MergeSoftmaxTraits<T>;
     int const batch_idx = static_cast<int>(blockIdx.y);

tensorrt_llm/_torch/attention_backend/trtllm.py

@@ -1449,7 +1449,6 @@ def merge_attention_for_mla(
         temp_attn: torch.Tensor,
         softmax_stats: torch.Tensor,
         temp_softmax_stats: torch.Tensor,
-        bmm1_scale: float,
         merge_op: torch.Tensor,
         metadata: TrtllmAttentionMetadata,
     ) -> None: