code cleanup

Signed-off-by: Amey naik <[email protected]>

Author: ameynaik-hub

cpp/tensorrt_llm/kernels/communicationKernels/allReduceFusionKernels.cu

@@ -691,11 +691,6 @@ void allreduce_fusion_kernel_launcher(AllReduceFusionParams const& params)
     TLLM_CHECK(oneshot || threads_per_block >= params.nranks);
     int block_size = threads_per_block;
 
-    // // Override block size to 1024 for AllGather operations
-    // if (params.pattern == AllReduceFusionPattern::kAllGather) {
-    //     block_size = 1024;
-    // }
-    
     TLLM_CHECK(block_size <= 1024 && cluster_size > 0);
 
     int grid_size = (std::min(sm_count, cluster_num * cluster_size) / cluster_size) * cluster_size;

cpp/tensorrt_llm/nanobind/runtime/bindings.cpp

@@ -345,6 +345,7 @@ void initBindings(nb::module_& m)
         .def("get_mc_buffer", &tensorrt_llm::runtime::McastGPUBuffer::getMCBuffer);
 
     nb::enum_<tensorrt_llm::kernels::AllReduceFusionOp>(m, "AllReduceFusionOp")
+	.value("ALLGATHER", tensorrt_llm::kernels::AllReduceFusionOp::ALLGATHER)
         .value("NONE", tensorrt_llm::kernels::AllReduceFusionOp::NONE)
         .value("RESIDUAL_RMS_NORM", tensorrt_llm::kernels::AllReduceFusionOp::RESIDUAL_RMS_NORM)
         .value("LAST_PROCESS_FOR_UB", tensorrt_llm::kernels::AllReduceFusionOp::LAST_PROCESS_FOR_UB)

tensorrt_llm/_torch/speculative/mtp.py

@@ -1069,6 +1069,7 @@ def get_local_max_and_combined(self, logits):
         original_last_dim = combined.shape[-1]
 
         # Ensure the combined tensor has at least 4 elements by padding with zeros
+        # This is required by the Lamport ALLGATHER kernel implementation
         if combined.numel() < 4:
             padding_size = 4 - combined.numel()
             # Create padding tensor with same shape as combined except for the last dimension

tests/unittest/_torch/multi_gpu/test_allreduce.py

@@ -106,7 +106,6 @@ def e2m1_and_ufp8sf_scale_to_float_v2(e2m1_tensor,
             e2m1_tensor, ufp8_scale_tensor, global_scale_tensor, sf_vec_size,
             ufp8_type, is_sf_swizzled_layout)
 
-    print(f"DBG AMEY: run_allreduce_op: i am here")
     x = x.cuda()
     residual = residual.cuda()
     norm_weight = torch.randn((hidden_size, ), dtype=dtype, device="cuda")
@@ -383,14 +382,12 @@ def ref_allgather(x, res):
 def test_allreduce_fusion_patterns(seq_len, hidden_size, fusion_op,
                                    mpi_pool_executor):
     torch.manual_seed(0)
-    print("DBG AMEY: test_allreduce_fusion_patterns: i am here")
     # dtype = torch.bfloat16
     dtype = torch.float32
     tensor_parallel_size = mpi_pool_executor.num_workers
     x = torch.randn((seq_len, hidden_size), dtype=dtype)
     residual = torch.randn_like(x)
     linear_weight = torch.randn((hidden_size, hidden_size), dtype=dtype)
-    print(f"DBG AMEY: test_allreduce_fusion_patterns: seq_len={seq_len}, hidden_size={hidden_size}, fusion_op={fusion_op}")
     results = mpi_pool_executor.map(
         run_single_rank,
         *zip(*[(tensor_parallel_size, run_allreduce_op, x, residual,