Address review comments.

Author: timlee0212

csrc/trtllm_mnnvl_allreduce.cu

@@ -103,8 +103,7 @@ void trtllm_mnnvl_allreduce_fusion(TensorView input, int64_t multicast_buffer_pt
       status = twoshotAllreduceFusionDispatch<c_type>(params);
     }
     TVM_FFI_ICHECK(status == cudaSuccess)
-        << "twoshot_allreduce_dispatch_world_size failed with error code "
-        << cudaGetErrorString(status);
+        << "trtllm_mnnvl_allreduce_fusion failed with error code " << cudaGetErrorString(status);
   });
 }
 

flashinfer/comm/trtllm_mnnvl_ar.py

@@ -192,9 +192,6 @@ def trtllm_mnnvl_allreduce_fusion(
             gamma: Gamma tensor (if rmsnorm)
             epsilon: Epsilon value (if rmsnorm)
         """
-        print(
-            f"[Rank {rank}] Inside Kernel: multicast_buffer_ptr: {multicast_buffer_ptr:x}, buffer_ptrs_dev: {buffer_ptrs_dev:x}, buffer_ptr_local: {buffer_ptr_local:x}, buffer_flags_mnnvl: {buffer_flags_mnnvl}"
-        )
         module.trtllm_mnnvl_allreduce_fusion(
             input,
             multicast_buffer_ptr,

tests/comm/test_trtllm_mnnvl_allreduce.py

@@ -1,5 +1,6 @@
 # Check torch version:
-from typing import Tuple, Optional
+import traceback
+from typing import Tuple
 
 import pytest
 import torch
@@ -451,8 +452,6 @@ def run_mnnvl_ar_full(
         rank_failed = True
         failure_message = f"FAILED[rank={rank}]: seq_lens={seq_lens}, fusion={fusion}, dtype={dtype} failed: {e}"
         print(failure_message)
-        import traceback
-
         print(traceback.format_exc())
 
         # Gather failure status from all ranks for logging