OSS CI fixes (pytorch#4077)

Summary:

X-link: facebookresearch/FBGEMM#1161

Fixing:
https://github.com/pytorch/FBGEMM/actions/runs/14771631771/job/41472512079#step:17:1902
https://github.com/pytorch/FBGEMM/actions/runs/14772049760

Reviewed By: q10, spcyppt

Differential Revision: D74196655

Author: Aya-ZIbra

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

@@ -227,7 +227,13 @@ DEVICE_INLINE void quantize_fp8_kv(
     __half2* qparam = nullptr,
     bool do_norm = false);
 
-DEVICE_INLINE void per_row_norm(fx4& dst);
+DEVICE_INLINE void per_row_norm(fx4& dst) {
+  float sum = fx4_dot(dst, dst);
+  // Warp reduce sum
+  sum = warpReduceSum(sum);
+  float rsqr = rsqrtf(sum / D_H);
+  dst = fx4_scale(dst, rsqr);
+}
 DEVICE_INLINE void per_row_amax(fx4& dst, float* amax);
 DEVICE_INLINE void per_head_amax(fx4& dst, float* amax);
 __global__ void nope_qkv_varseq_prefill_kernel(
@@ -2850,7 +2856,6 @@ at::Tensor quantize_qkv_per_head(
     // HH += N_KVH_L * 2;
     qparam_k_ptr = qparam_k.value().data_ptr<float>();
     qparam_v_ptr = qparam_v.value().data_ptr<float>();
-    CHECK_EQ(HH, 7);
   }
   auto num_warps = B_T * HH;
   dim3 block_size(kThreadsPerWarp, kWarpsPerBlock);
@@ -2883,13 +2888,6 @@ at::Tensor quantize_qkv_per_head(
   C10_CUDA_KERNEL_LAUNCH_CHECK();
   return scale_q;
 }
-DEVICE_INLINE void per_row_norm(fx4& dst) {
-  float sum = fx4_dot(dst, dst);
-  // Warp reduce sum
-  sum = warpReduceSum(sum);
-  float rsqr = rsqrtf(sum / D_H);
-  dst = fx4_scale(dst, rsqr);
-}
 
 DEVICE_INLINE void per_head_amax(fx4& dst, float* amax) {
   dst = fx4_abs(dst);
@@ -2998,5 +2996,21 @@ std::tuple<at::Tensor, at::Tensor> dequantize_fp8_cache(
   throw std::runtime_error(
       "CUDA version is older than 12.0"); // requires CUDA>=12
 }
+at::Tensor quantize_qkv_per_head(
+    at::Tensor xqkv_amax_row, // [B_T, HH]
+    at::Tensor xqkv, // [B_T, HH, D_H]
+    at::Tensor varseq_seqpos, // [B_T]
+    std::optional<at::Tensor> varseq_batch, // [B_T]
+    at::Tensor q_seqstarts, // [B+1]
+    at::Tensor cache_K, // [B][MAX_T][N_KVH][D_H]
+    at::Tensor cache_V, // [B][MAX_T][N_KVH][D_H]
+    at::Tensor XQ_O, // [B_T][N_H][D]
+    int64_t max_seq_length, // Length of the sequence
+    std::optional<at::Tensor> qparam_k,
+    std::optional<at::Tensor> qparam_v) {
+  throw std::runtime_error(
+      "CUDA version is older than 12.0"); // requires CUDA>=12
+}
+
 #endif
 } // namespace fbgemm_gpu