Fix coderabbit nits

Author: djns99

csrc/trtllm_moe_alltoall.cu

@@ -188,7 +188,7 @@ Tuple<Array<int64_t>, Array<int64_t>, int64_t> moeA2ADispatchOp(
 
     TVM_FFI_ICHECK(totalBytesNeeded % elementSize == 0)
         << "Misaligned payload buffer " << i << " with element size " << elementSize
-        << ". Consider putting ordering payloads by minimum element size";
+        << ". Consider reordering payloads by largest to smallest element size";
   }
 
   auto* workspaceBase = static_cast<uint8_t*>(workspace.data_ptr());

flashinfer/comm/trtllm_moe_alltoall.py

@@ -5,8 +5,6 @@
 supporting multiple payloads per collective operation.
 """
 
-# TODO Review
-
 from dataclasses import dataclass
 from types import SimpleNamespace
 from typing import Optional
@@ -648,9 +646,9 @@ def get_combine_payload_tensor_in_workspace(
 
 __all__ = [
     "MoeAlltoAll",
-    "moe_a2a_initialize",
-    "moe_a2a_dispatch",
     "moe_a2a_combine",
-    "moe_a2a_sanitize_expert_ids",
+    "moe_a2a_dispatch",
     "moe_a2a_get_workspace_size_per_rank",
+    "moe_a2a_initialize",
+    "moe_a2a_sanitize_expert_ids",
 ]

tests/comm/test_mnnvl_moe_alltoall.py

@@ -38,17 +38,18 @@ def safe_run(func, *args, **kwargs):
     comm = MPI.COMM_WORLD
     try:
         func(*args, **kwargs)
-    except MPIExit as e:
-        raise e
-    except Exception as e:
+    except MPIExit:
+        raise
+    except Exception:
         traceback.print_exc()
         comm.allgather(True)
-        raise e
+        raise
 
 
 @pytest.fixture(autouse=True)
 def setup_test():
     torch.manual_seed(0x1234)
+    yield
 
 
 def compute_target_rank_id(expert_id, num_experts_per_rank):
@@ -154,7 +155,7 @@ def fake_moe(
     # Process each token
     for token_idx in range(num_tokens):
         results = []
-        # For each expert selected for this token/
+        # For each expert selected for this token
         for k in range(top_k):
             expert_id = token_selected_experts[token_idx, k].item()
             if is_ep:

tests/comm/test_trtllm_moe_alltoall.py

@@ -24,7 +24,7 @@
 
 @pytest.fixture(autouse=True, scope="session")
 def setup_test_environment():
-    """Set up test environment and warm up JIT compilation."""
+    """Set up torch seed for deterministic tests."""
     torch.manual_seed(0xD5)
     yield
 
@@ -410,7 +410,7 @@ def fake_moe(
     # Process each token
     for token_idx in range(num_tokens):
         results = []
-        # For each expert selected for this token/
+        # For each expert selected for this token
         for k in range(top_k):
             expert_id = token_selected_experts[token_idx, k].item()
             if is_ep and not (