iRoPE (#4063)

Summary:

X-link: facebookresearch/FBGEMM#1149

When the max_seq_len is larger than 8192, one input sample will be divided into multiple sequences. Such as:
When bs = 2, and seqlen = 7, we will have seq_lens = [0, 7, 7, 7, 7, 14, 14, 14, 14] in the prefill attention. In decoding, it won't as it's handled by the gappy bias.

Reviewed By: sijiac

Differential Revision: D73833204

Author: y-sq

fbgemm_gpu/experimental/gen_ai/src/kv_cache/kv_cache.cu

@@ -3003,7 +3003,8 @@ at::Tensor quantize_qkv_per_head(
   dim3 block_size(kThreadsPerWarp, kWarpsPerBlock);
   dim3 grid_size(cuda_calc_xblock_count(num_warps, kWarpsPerBlock));
 
-  auto scale_q = at::zeros({B, N_KVH_L}, XQ_O.options().dtype(at::kFloat));
+  auto scale_q =
+      at::zeros({cache_K.size(0), N_KVH_L}, XQ_O.options().dtype(at::kFloat));
   float* const scale_q_ptr = scale_q.data_ptr<float>();
   // Launch the kernel
   // TODO: Launch the kernel with B_T * N_H_L blocks only in case of decode.