   assert(blockDim.x % 32 == 0);
 
   float2 stride_max = make_float2(-1e20f, -1e20f);
-  for (int i = 0; i < seq_len; i += blockDim.x) {
+  for (int i = 0; threadIdx.x + i < seq_len; i += blockDim.x) {
     float2 cur =
         phi::funcs::ToFloat2<half2>(qk_buf_[threadIdx.x + i + qk_offset] +
                                     bias_qk_[threadIdx.x + i + qk_offset]);

      phi::funcs::blockReduceMax<float>(max(stride_max.x, stride_max.y), mask);

  float2 stride_sum = make_float2(0.f, 0.f);
  for (int i = 0; i < seq_len; i += blockDim.x) {
    float2 cur = phi::funcs::ToFloat2<T>(qk_buf_[threadIdx.x + i + qk_offset] +
                                         bias_qk_[threadIdx.x + i + qk_offset]);
    stride_sum.x += __expf(cur.x - max_val);

      phi::funcs::blockReduceSum<float>(stride_sum.x + stride_sum.y, mask) +
      1e-6f;

  for (int i = 0; threadIdx.x + i < seq_len; i += blockDim.x) {
    float2 cur = phi::funcs::ToFloat2<T>(qk_buf_[threadIdx.x + i + qk_offset] +
                                         bias_qk_[threadIdx.x + i + qk_offset]);
    qk_buf_[threadIdx.x + i + qk_offset] = phi::funcs::FloatsToPair<T>(
