Add dtype checks for q-kv tensors (#280)

Right now, we require q and kv tensors to have the same dtype, but that
is not enforced, which can lead to cryptic memory errors in case of a
misconfiguration. This PR adds a check to ensure that we prevent
mismatched dtypes.

Author: Yard1

python/csrc/batch_decode.cu

@@ -33,6 +33,8 @@ std::vector<torch::Tensor> batch_decode_with_padded_kv_cache(
   CHECK_SHAPE(k_padded, v_padded);
   CHECK_EQ(q.size(0), k_padded.size(0));
   CHECK_EQ(q.size(2), k_padded.size(3));
+  CHECK_EQ(q.scalar_type(), k_padded.scalar_type());
+  CHECK_EQ(q.scalar_type(), v_padded.scalar_type());
   unsigned int batch_size = q.size(0);
   unsigned int num_qo_heads = q.size(1);
   unsigned int head_dim = q.size(2);
@@ -206,6 +208,7 @@ std::vector<torch::Tensor> BatchDecodeWithPagedKVCachePyTorchWrapper::Forward(
   CHECK_DIM(1, paged_kv_last_page_len);  // (B,)
   CHECK_DIM(1, paged_kv_indptr);         // (B+1,)
   CHECK_DIM(1, paged_kv_indices);        // (nnz,)
+  CHECK_EQ(q.scalar_type(), paged_kv_data.scalar_type());
   // (num_max_pages, 2, H_kv, page_size, head_dim) for HND
   // (num_max_pages, 2, page_size, H_kv, head_dim) for NHD
   CHECK_DIM(5, paged_kv_data);

python/csrc/batch_prefill.cu

@@ -70,6 +70,7 @@ std::vector<torch::Tensor> BatchPrefillWithPagedKVCachePyTorchWrapper::Forward(
   CHECK_DIM(1, paged_kv_indptr);         // (B + 1,)
   CHECK_DIM(1, paged_kv_indices);        // (nnz_kv,)
   CHECK_DIM(1, paged_kv_last_page_len);  // (B,)
+  CHECK_EQ(q.scalar_type(), paged_kv_data.scalar_type());
   int64_t batch_size = qo_indptr.size(0) - 1;
   int64_t nnz_qo = q.size(0);
   int64_t num_qo_heads = q.size(1);
@@ -173,6 +174,7 @@ std::vector<torch::Tensor> BatchPrefillWithPagedKVCachePyTorchWrapper::ForwardCu
   CHECK_DIM(1, paged_kv_last_page_len);  // (B,)
   CHECK_DIM(1, custom_mask);             // (nnz_qk,)
   CHECK_DIM(1, qk_indptr);               // (B + 1,)
+  CHECK_EQ(q.scalar_type(), paged_kv_data.scalar_type());
   int64_t batch_size = qo_indptr.size(0) - 1;
   int64_t nnz_qo = q.size(0);
   int64_t num_qo_heads = q.size(1);
@@ -299,6 +301,8 @@ std::vector<torch::Tensor> BatchPrefillWithRaggedKVCachePyTorchWrapper::Forward(
   CHECK_DIM(3, k);          // (nnz_kv, H_kv, D) if NHD else (H_kv, nnz_kv, D)
   CHECK_DIM(3, v);          // (nnz_kv, H_kv, D) if NHD else (H_kv, nnz_kv, D)
   CHECK_DIM(1, kv_indptr);  // (B + 1,)
+  CHECK_EQ(q.scalar_type(), k.scalar_type());
+  CHECK_EQ(q.scalar_type(), v.scalar_type());
   int64_t batch_size = qo_indptr.size(0) - 1;
   int64_t nnz_qo = q.size(0);
   int64_t num_qo_heads = q.size(1);
@@ -382,6 +386,8 @@ std::vector<torch::Tensor> BatchPrefillWithRaggedKVCachePyTorchWrapper::ForwardC
   CHECK_DIM(1, kv_indptr);    // (B + 1,)
   CHECK_DIM(1, custom_mask);  // (nnz_qk,)
   CHECK_DIM(1, qk_indptr);    // (B + 1,)
+  CHECK_EQ(q.scalar_type(), k.scalar_type());
+  CHECK_EQ(q.scalar_type(), v.scalar_type());
   int64_t batch_size = qo_indptr.size(0) - 1;
   int64_t nnz_qo = q.size(0);
   int64_t num_qo_heads = q.size(1);

python/csrc/single_decode.cu

@@ -32,6 +32,8 @@ torch::Tensor single_decode_with_kv_cache(torch::Tensor q, torch::Tensor k, torc
   CHECK_DIM(3, v);
   CHECK_SHAPE(k, v);
   CHECK_EQ(q.size(1), k.size(2));
+  CHECK_EQ(q.scalar_type(), k.scalar_type());
+  CHECK_EQ(q.scalar_type(), v.scalar_type());
   unsigned int num_qo_heads = q.size(0);
   unsigned int head_dim = q.size(1);
   unsigned int kv_len, num_kv_heads;

python/csrc/single_prefill.cu

@@ -32,6 +32,8 @@ std::vector<torch::Tensor> single_prefill_with_kv_cache(
   CHECK_DIM(3, v);
   CHECK_SHAPE(k, v);
   CHECK_EQ(q.size(2), k.size(2));
+  CHECK_EQ(q.scalar_type(), k.scalar_type());
+  CHECK_EQ(q.scalar_type(), v.scalar_type());
   unsigned int head_dim = q.size(2);
   unsigned int kv_len, qo_len, num_kv_heads, num_qo_heads;
   QKVLayout kv_layout = static_cast<QKVLayout>(layout);