gemini-code-assist[bot] commented on code in PR #18861:
URL: https://github.com/apache/tvm/pull/18861#discussion_r2869487247


##########
src/runtime/contrib/nvshmem/kv_transfer.cu:
##########
@@ -258,34 +265,36 @@ int _KVTransfer(DLTensor* remote_pages, DLTensor* k, 
DLTensor* v, DLTensor* remo
 int _KVTransferPageToPage(DLTensor* remote_pages, DLTensor* local_pages,
                           DLTensor* remote_position_map, DLTensor* 
local_position_map,
                           DLTensor* remote_tp_group_pe_offset, TVMStreamHandle 
transfer_stream) {
-  CHECK_EQ(remote_pages->device.device_type, kDLCUDA)
+  TVM_FFI_CHECK_EQ(remote_pages->device.device_type, kDLCUDA, ValueError)
       << "The device of remote_pages matrix must be CUDA.";
-  CHECK_EQ(local_pages->device.device_type, kDLCUDA) << "The device of k 
matrix must be CUDA.";
-  CHECK_EQ(remote_position_map->device.device_type, kDLCUDA)
+  TVM_FFI_CHECK_EQ(local_pages->device.device_type, kDLCUDA, ValueError)
+      << "The device of k matrix must be CUDA.";

Review Comment:
   ![medium](https://www.gstatic.com/codereviewagent/medium-priority.svg)
   
   The error message here refers to `k matrix`, but the check is being 
performed on `local_pages`. This seems to be a copy-paste error and could be 
misleading. Please update the message to refer to `local_pages`.
   
   ```
     TVM_FFI_CHECK_EQ(local_pages->device.device_type, kDLCUDA, ValueError)
         << "The device of local_pages matrix must be CUDA.";
   ```



##########
src/runtime/contrib/nvshmem/kv_transfer.cu:
##########
@@ -258,34 +265,36 @@ int _KVTransfer(DLTensor* remote_pages, DLTensor* k, 
DLTensor* v, DLTensor* remo
 int _KVTransferPageToPage(DLTensor* remote_pages, DLTensor* local_pages,
                           DLTensor* remote_position_map, DLTensor* 
local_position_map,
                           DLTensor* remote_tp_group_pe_offset, TVMStreamHandle 
transfer_stream) {
-  CHECK_EQ(remote_pages->device.device_type, kDLCUDA)
+  TVM_FFI_CHECK_EQ(remote_pages->device.device_type, kDLCUDA, ValueError)
       << "The device of remote_pages matrix must be CUDA.";
-  CHECK_EQ(local_pages->device.device_type, kDLCUDA) << "The device of k 
matrix must be CUDA.";
-  CHECK_EQ(remote_position_map->device.device_type, kDLCUDA)
+  TVM_FFI_CHECK_EQ(local_pages->device.device_type, kDLCUDA, ValueError)
+      << "The device of k matrix must be CUDA.";
+  TVM_FFI_CHECK_EQ(remote_position_map->device.device_type, kDLCUDA, 
ValueError)
       << "The device of remote_position_map matrix must be CUDA.";
   size_t dev_id = remote_pages->device.device_id;
-  CHECK_EQ(local_pages->device.device_id, dev_id)
+  TVM_FFI_CHECK_EQ(local_pages->device.device_id, dev_id, ValueError)
       << "The device id of remote_pages and k matrix doesn't match.";

Review Comment:
   ![medium](https://www.gstatic.com/codereviewagent/medium-priority.svg)
   
   Similar to the previous comment, the error message here refers to `k matrix` 
while checking `local_pages`. Please update it for clarity to refer to 
`local_pages`.
   
   ```
     TVM_FFI_CHECK_EQ(local_pages->device.device_id, dev_id, ValueError)
         << "The device id of remote_pages and local_pages matrix doesn't 
match.";
   ```



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]


---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to