perf: Eliminate the need for attention DP padding when possible

cpp/tensorrt_llm/thop/allgatherOp.cpp

@@ -40,9 +40,8 @@ namespace
 class AllgatherOp
 {
 public:
-    AllgatherOp(std::set<int> group, nvinfer1::DataType type)
+    AllgatherOp(std::set<int> group)
         : mGroup(std::move(group))
-        , mType(type)
     {
     }
 
@@ -56,55 +55,113 @@ class AllgatherOp
         return 0;
     }
 
-    torch::Tensor run(torch::Tensor input) noexcept
+    torch::Tensor run(torch::Tensor input, torch::optional<torch::List<int64_t>> sizes) noexcept
     {
+        TLLM_CHECK_WITH_INFO(mNcclComm.get() != nullptr, "mNcclComm should be initialized before used");
         auto stream = at::cuda::getCurrentCUDAStream(input.get_device());
+        auto type = tensorrt_llm::runtime::TorchUtils::dataType(input.scalar_type());
         std::vector<int64_t> outputShape = input.sizes().vec();
-        outputShape.insert(outputShape.begin(), mGroup.size());
+        if (sizes.has_value())
+        {
+            outputShape[0] = std::accumulate(sizes.value().begin(), sizes.value().end(), 0, std::plus<>{});
+        }
+        else
+        {
+            outputShape[0] *= mGroup.size();
+        }
         auto output = torch::empty(outputShape, input.options());
-        size_t size = input.numel();
-        TLLM_CHECK_WITH_INFO(mNcclComm.get() != nullptr, "mNcclComm should be initialized before used");
-        NCCLCHECK(ncclAllGather(
-            input.data_ptr(), output.mutable_data_ptr(), size, (*getDtypeMap())[mType], *mNcclComm, stream));
+        if (sizes.has_value())
+        {
+            size_t numel_base = std::accumulate(outputShape.cbegin() + 1, outputShape.cend(), 1, std::multiplies<>{});
+            int64_t split_offset = 0;
+            ncclGroupStart();
+            for (int root = 0; root < static_cast<int>(mGroup.size()); ++root)
+            {
+                auto split_size = sizes.value()[root];
+                NCCLCHECK(ncclBroadcast(input.data_ptr(),
+                    output.index({torch::indexing::Slice(split_offset, torch::indexing::None)}).mutable_data_ptr(),
+                    numel_base * split_size, (*getDtypeMap())[type], root, *mNcclComm, stream));
+                split_offset += split_size;
+            }
+            ncclGroupEnd();
+        }
+        else
+        {
+            NCCLCHECK(ncclAllGather(input.data_ptr(), output.mutable_data_ptr(), input.numel(), (*getDtypeMap())[type],
+                *mNcclComm, stream));
+        }
         return output;
     }
 
+    std::vector<torch::Tensor> run_list(
+        torch::TensorList input_list, torch::optional<torch::List<int64_t>> sizes) noexcept
+    {
+        std::vector<torch::Tensor> output_list;
+        output_list.reserve(input_list.size());
+        ncclGroupStart();
+        for (auto const& input : input_list)
+        {
+            auto output = run(input, sizes);
+            output_list.push_back(output);
+        }
+        ncclGroupEnd();
+        return output_list;
+    }
+
 private:
     std::set<int> mGroup;
-    nvinfer1::DataType mType;
     std::shared_ptr<ncclComm_t> mNcclComm;
 };
 
 } // namespace
 
 #endif // ENABLE_MULTI_DEVICE
 
-torch::Tensor allgather(torch::Tensor input, torch::List<int64_t> group_)
+torch::Tensor allgather(torch::Tensor input, torch::optional<torch::List<int64_t>> sizes, torch::List<int64_t> group_)
 {
 #if ENABLE_MULTI_DEVICE
-    auto const type = tensorrt_llm::runtime::TorchUtils::dataType(input.scalar_type());
     std::set<int> group;
     for (int64_t rank : group_)
     {
         group.insert(static_cast<int>(rank));
     }
-    AllgatherOp op(group, type);
+    AllgatherOp op(group);
     op.initialize();
-    auto output = op.run(input);
+    auto output = op.run(input, sizes);
     return output;
 #else
     return input;
 #endif // ENABLE_MULTI_DEVICE
 }
 
+std::vector<torch::Tensor> allgather_list(
+    torch::TensorList input_list, torch::optional<torch::List<int64_t>> sizes, torch::List<int64_t> group_)
+{
+#if ENABLE_MULTI_DEVICE
+    std::set<int> group;
+    for (int64_t rank : group_)
+    {
+        group.insert(static_cast<int>(rank));
+    }
+    AllgatherOp op(group);
+    op.initialize();
+    auto output_list = op.run_list(input_list, sizes);
+    return output_list;
+#else
+    return input_list.vec();
+#endif // ENABLE_MULTI_DEVICE
+}
+
 } // namespace torch_ext
 
 TORCH_LIBRARY_FRAGMENT(trtllm, m)
 {
-    m.def("allgather(Tensor input, int[] group) -> Tensor");
+    m.def("allgather(Tensor input, int[]? sizes, int[] group) -> Tensor");
+    m.def("allgather_list(Tensor[] input_list, int[]? sizes, int[] group) -> Tensor[]");
 }
 
 TORCH_LIBRARY_IMPL(trtllm, CUDA, m)
 {
     m.impl("allgather", &torch_ext::allgather);
+    m.impl("allgather_list", &torch_ext::allgather_list);
 }

cpp/tensorrt_llm/thop/reducescatterOp.cpp

@@ -40,9 +40,8 @@ namespace
 class ReducescatterOp
 {
 public:
-    ReducescatterOp(std::set<int> group, nvinfer1::DataType type)
+    ReducescatterOp(std::set<int> group)
         : mGroup(std::move(group))
-        , mType(type)
     {
     }
 
@@ -56,55 +55,124 @@ class ReducescatterOp
         return 0;
     }
 
-    torch::Tensor run(torch::Tensor const& input) noexcept
+    torch::Tensor run(torch::Tensor const& input, torch::optional<torch::List<int64_t>> sizes) noexcept
     {
+        TLLM_CHECK_WITH_INFO(mNcclComm.get() != nullptr, "mNcclComm should be initialized before used");
         auto stream = at::cuda::getCurrentCUDAStream(input.get_device());
+        auto type = tensorrt_llm::runtime::TorchUtils::dataType(input.scalar_type());
         std::vector<int64_t> outputShape = input.sizes().vec();
-        outputShape[0] = outputShape[0] / mGroup.size();
+        if (sizes.has_value())
+        {
+            auto rank = COMM_SESSION.getRank();
+            int groupRank = 0;
+            for (auto const& currentRank : mGroup)
+            {
+                if (rank == currentRank)
+                    break;
+                ++groupRank;
+            }
+            TLLM_CHECK(static_cast<size_t>(groupRank) < mGroup.size());
+            outputShape[0] = sizes.value()[groupRank];
+        }
+        else
+        {
+            outputShape[0] = outputShape[0] / mGroup.size();
+        }
         auto output = torch::empty(outputShape, input.options());
-        size_t const size = output.numel();
-        TLLM_CHECK_WITH_INFO(mNcclComm.get() != nullptr, "mNcclComm should be initialized before used");
-        NCCLCHECK(ncclReduceScatter(
-            input.data_ptr(), output.mutable_data_ptr(), size, (*getDtypeMap())[mType], ncclSum, *mNcclComm, stream));
+        if (sizes.has_value())
+        {
+            size_t numel_base = std::accumulate(outputShape.cbegin() + 1, outputShape.cend(), 1, std::multiplies<>{});
+            int64_t split_offset = 0;
+            ncclGroupStart();
+            for (int root = 0; root < static_cast<int>(mGroup.size()); ++root)
+            {
+                auto split_size = sizes.value()[root];
+                NCCLCHECK(
+                    ncclReduce(input.index({torch::indexing::Slice(split_offset, torch::indexing::None)}).data_ptr(),
+                        output.mutable_data_ptr(), numel_base * split_size, (*getDtypeMap())[type], ncclSum, root,
+                        *mNcclComm, stream));
+                split_offset += split_size;
+            }
+            ncclGroupEnd();
+        }
+        else
+        {
+            NCCLCHECK(ncclReduceScatter(input.data_ptr(), output.mutable_data_ptr(), output.numel(),
+                (*getDtypeMap())[type], ncclSum, *mNcclComm, stream));
+        }
         return output;
     }
 
+    std::vector<torch::Tensor> run_list(
+        torch::TensorList input_list, torch::optional<torch::List<int64_t>> sizes) noexcept
+    {
+        std::vector<torch::Tensor> output_list;
+        output_list.reserve(input_list.size());
+        ncclGroupStart();
+        for (auto const& input : input_list)
+        {
+            auto output = run(input, sizes);
+            output_list.push_back(output);
+        }
+        ncclGroupEnd();
+        return output_list;
+    }
+
 private:
     std::set<int> mGroup;
-    nvinfer1::DataType mType;
     std::shared_ptr<ncclComm_t> mNcclComm;
 };
 
 } // namespace
 
 #endif // ENABLE_MULTI_DEVICE
 
-extern torch::Tensor reducescatter(torch::Tensor input, torch::List<int64_t> group_)
+extern torch::Tensor reducescatter(
+    torch::Tensor input, torch::optional<torch::List<int64_t>> sizes, torch::List<int64_t> group_)
 {
 #if ENABLE_MULTI_DEVICE
-    auto const type = tensorrt_llm::runtime::TorchUtils::dataType(input.scalar_type());
     std::set<int> group;
     for (int64_t rank : group_)
     {
         group.insert(static_cast<int>(rank));
     }
-    ReducescatterOp op(group, type);
+    ReducescatterOp op(group);
     op.initialize();
-    auto output = op.run(input);
+    auto output = op.run(input, sizes);
     return output;
 #else
     return input;
 #endif // ENABLE_MULTI_DEVICE
 }
 
+extern std::vector<torch::Tensor> reducescatter_list(
+    torch::TensorList input_list, torch::optional<torch::List<int64_t>> sizes, torch::List<int64_t> group_)
+{
+#if ENABLE_MULTI_DEVICE
+    std::set<int> group;
+    for (int64_t rank : group_)
+    {
+        group.insert(static_cast<int>(rank));
+    }
+    ReducescatterOp op(group);
+    op.initialize();
+    auto output_list = op.run_list(input_list, sizes);
+    return output_list;
+#else
+    return input_list.vec();
+#endif // ENABLE_MULTI_DEVICE
+}
+
 } // namespace torch_ext
 
 TORCH_LIBRARY_FRAGMENT(trtllm, m)
 {
-    m.def("reducescatter(Tensor input, int[] group) -> Tensor");
+    m.def("reducescatter(Tensor input, int[]? sizes, int[] group) -> Tensor");
+    m.def("reducescatter_list(Tensor[] input_list, int[]? sizes, int[] group) -> Tensor[]");
 }
 
 TORCH_LIBRARY_IMPL(trtllm, CUDA, m)
 {
     m.impl("reducescatter", &torch_ext::reducescatter);
+    m.impl("reducescatter_list", &torch_ext::reducescatter_list);
 }

tensorrt_llm/_torch/auto_deploy/custom_ops/dist.py

@@ -1,16 +1,20 @@
 """Custom ops required for implementing tensor parallelism."""
 
+from typing import List, Optional
+
 import torch
 
 from ..distributed import common as dist
 from ..distributed import trtllm as trtllm_dist
 
 
 @torch.library.custom_op("dist::all_gather", mutates_args=(), device_types="cuda")
-def all_gather(tensor: torch.Tensor, dim: int = 0) -> torch.Tensor:
+def all_gather(
+    tensor: torch.Tensor, dim: int = 0, sizes: Optional[List[int]] = None
+) -> torch.Tensor:
     """All gather followed by concat in dim = 0. This is the default nccl behavior."""
     if trtllm_dist.is_trtllm_op_available():
-        return trtllm_dist.trtllm_allgather(tensor, dim=dim)
+        return trtllm_dist.trtllm_allgather(tensor, dim=dim, sizes=sizes)
     tl = [torch.zeros_like(tensor) for _ in range(dist.get_world_size())]
     dist.all_gather(tl, tensor)
     return torch.cat(tl, dim=dim)

tensorrt_llm/_torch/auto_deploy/distributed/trtllm.py

@@ -8,10 +8,10 @@
     from ...distributed import AllReduce, allgather
     from ...modules.linear import AllReduceFusionOp, AllReduceParams
 
-    def trtllm_allgather(tensor, dim):
+    def trtllm_allgather(tensor, dim, sizes=None):
         rank, world_size = get_rank_world_size()
         p_config = Mapping(world_size=world_size, tp_size=world_size, rank=rank)
-        return allgather(tensor, p_config, gather_dim=dim)
+        return allgather(tensor, p_config, dim=dim, sizes=sizes)
 
     def trtllm_allreduce(tensor, op, all_reduce_params=None):
         rank, world_size = get_rank_world_size()
@@ -45,7 +45,7 @@ def fused_allreduce_residual_rmsnorm_fake(
     TRTLLM_OP_AVAILABLE = True
 except ImportError:
 
-    def trtllm_allgather(tensor, dim):
+    def trtllm_allgather(tensor, dim, sizes=None):
         raise ImportError("TRT-LLM is not available.")
 
     def trtllm_allreduce(tensor, op):

tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py

@@ -61,8 +61,11 @@ def _(residual, norm_weight, device_num_experts, scale_input,
         return [norm_out, residual_out]
 
     @torch.library.register_fake("trtllm::allgather")
-    def _(input, group):
-        output_shape = (len(group), *input.shape)
+    def _(input, sizes, group):
+        if sizes is None:
+            output_shape = (len(group) * input.shape[0], *input.shape[1:])
+        else:
+            output_shape = (sum(sizes), *input.shape[1:])
         return input.new_empty(output_shape)
 
     @torch.library.register_fake("trtllm::cublas_scaled_mm")