diff --git a/DCNv4_op/DCNv4/modules/flash_deform_attn.py b/DCNv4_op/DCNv4/modules/flash_deform_attn.py
index 0c4922b..d7340de 100644
--- a/DCNv4_op/DCNv4/modules/flash_deform_attn.py
+++ b/DCNv4_op/DCNv4/modules/flash_deform_attn.py
@@ -111,7 +111,6 @@ class FlashDeformAttn(nn.Module):
         value = value.view(N, Len_in, self.n_heads, self.d_model // self.n_heads)
         sampling_offsets = self.sampling_offsets(query).view(N, Len_q, self.n_heads, self.n_levels, self.n_points, 2)
         attention_weights = self.attention_weights(query).view(N, Len_q, self.n_heads, self.n_levels * self.n_points)
-        attention_weights = F.softmax(attention_weights, -1).view(N, Len_q, self.n_heads, self.n_levels, self.n_points)
         # N, Len_q, n_heads, n_levels, n_points, 2
         if reference_points.shape[-1] == 2:
             offset_normalizer = torch.stack([input_spatial_shapes[..., 1], input_spatial_shapes[..., 0]], -1)
@@ -128,13 +127,16 @@ class FlashDeformAttn(nn.Module):
             raise ValueError(
                 "Last dim of reference_points must be 2 or 4, but get {} instead.".format(reference_points.shape[-1])
             )
+            
+        # Cat sampling_offsets and attention_weights, generate sampling_loc_attn:
+        sampling_locations = sampling_locations.flatten(-3).half()
+        sampling_loc_attn = torch.cat([sampling_locations, attention_weights], dim=-1)
 
         output = FlashDeformAttnFunction.apply(
             value,
             input_spatial_shapes,
             input_level_start_index,
-            sampling_locations,
-            attention_weights,
+            sampling_loc_attn,
             self.im2col_step,
             self.n_points           
         )