   assert(blockDim.x % 32 == 0);
 
   T stride_max = -1e20f;
-  for (int i = 0; i < seq_len; i += blockDim.x) {
+  for (int i = 0; threadIdx.x + i < seq_len; i += blockDim.x) {
     stride_max = qk_buf[threadIdx.x + i + qk_offset] +
                              bias_qk[threadIdx.x + i + qk_offset] >
                          stride_max

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

  for (int i = 0; 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]);

  assert(blockDim.x % 32 == 0);

  float2 stride_max = make_float2(-1e20f, -1e20f);
  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]);
