add fp4 gemm + allreduce

Signed-off-by: benzh <[email protected]>

Author: benzh-2025

cpp/tensorrt_llm/thop/CMakeLists.txt

@@ -104,7 +104,8 @@ add_library(
   loraOp.cpp
   finegrained_mixed_dtype_gemm_thop.cpp
   tinygemm2.cpp
-  dsv3RopeOp.cpp)
+  dsv3RopeOp.cpp
+  fusedGemmAllreduceOp.cpp)
 set_property(TARGET th_common PROPERTY POSITION_INDEPENDENT_CODE ON)
 target_link_libraries(
   th_common PRIVATE ${TORCH_LIBRARIES} th_utils ${Python3_LIBRARIES}

cpp/tensorrt_llm/thop/fusedGemmAllreduceOp.cpp

@@ -0,0 +1,265 @@
+/*
+ * Copyright (c) 2020-2023, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "cutlass_extensions/gemm_configs.h"
+
+#include "tensorrt_llm/common/cudaUtils.h"
+#include "tensorrt_llm/kernels/cutlass_kernels/include/allreduce_gemm_runner.h"
+#include "tensorrt_llm/runtime/ipcNvlsMemory.h"
+#include "tensorrt_llm/thop/thUtils.h"
+
+#include <cstddef>
+#include <cuda_fp16.h>
+
+#include <cstdint>
+#include <functional>
+#include <type_traits>
+#include <vector>
+
+using tensorrt_llm::kernels::opened_cutlass_kernels::GemmAllReduceImplRunner;
+using tensorrt_llm::kernels::opened_cutlass_kernels::GemmAllReduceImplInterface;
+using tensorrt_llm::kernels::opened_cutlass_kernels::GemmTypes;
+using tensorrt_llm::kernels::opened_cutlass_kernels::PersistentWorkspaceInterface;
+
+namespace
+{
+struct AllocationKey
+{
+    c10::cuda::CUDAStream stream;
+    std::set<int> group;
+
+    bool operator==(AllocationKey const& other) const
+    {
+        return stream.id() == other.stream.id() && stream.device_index() == other.stream.device_index()
+            && group == other.group;
+    }
+
+    std::string toString() const
+    {
+        std::stringstream ss;
+        ss << "AllocationKey(stream: " << stream.id() << ", device: " << (int) stream.device_index() << ", group: [";
+        for (int rank : group)
+        {
+            ss << rank << ", ";
+        }
+        ss << "])";
+        return ss.str();
+    }
+};
+
+struct AllocationKeyHash
+{
+    size_t operator()(AllocationKey const& key) const
+    {
+        size_t seed = 0;
+
+        // Hash the stream (using stream ID and device index)
+        hash_combine(seed, key.stream.id());
+        hash_combine(seed, key.stream.device_index());
+
+        // Hash the set elements
+        for (auto const& elem : key.group)
+        {
+            hash_combine(seed, elem);
+        }
+
+        return seed;
+    }
+
+private:
+    template <typename T>
+    static void hash_combine(size_t& seed, T const& val)
+    {
+        seed ^= std::hash<T>()(val) + 0x9e3779b9 + (seed << 6) + (seed >> 2);
+    }
+};
+
+class IpcNvlsHandleWrapper
+{
+public:
+    IpcNvlsHandleWrapper(size_t size, std::set<int> groups)
+        : mSize(size)
+    {
+        mHandle = tensorrt_llm::runtime::ipcNvlsAllocate(size, groups);
+    }
+
+    tensorrt_llm::runtime::IpcNvlsHandle* getHandle() const
+    {
+        return mHandle;
+    }
+
+    size_t getSize() const
+    {
+        return mSize;
+    }
+
+    ~IpcNvlsHandleWrapper()
+    {
+        tensorrt_llm::runtime::ipcNvlsFree(mHandle);
+    }
+
+private:
+    size_t mSize;
+    tensorrt_llm::runtime::IpcNvlsHandle* mHandle;
+};
+
+class NvlsMemoryManager
+{
+public:
+    PersistentWorkspaceInterface* getWorkspace(GemmAllReduceImplInterface* runner,
+        GemmAllReduceImplInterface::ProblemArgs const& problem, AllocationKey const& key)
+    {
+        auto newWorkspace = runner->getPersistentWorkspace(problem);
+        auto curWorkspace = mWorkspaces[key];
+
+        if (curWorkspace == nullptr || curWorkspace->size() < newWorkspace->size())
+        {
+            TLLM_LOG_WARNING("NvlsMemoryManager allocate workspace: key=%s, size=%zu.", key.toString().c_str(),
+                newWorkspace->size());
+            newWorkspace->allocate();
+            mWorkspaces[key] = newWorkspace;
+        }
+        return mWorkspaces[key].get();
+    }
+
+    tensorrt_llm::runtime::IpcNvlsHandle* getD(size_t size, AllocationKey const& key)
+    {
+        constexpr size_t PREFER_SIZE = 1024 * 16384 * 2;
+        auto handle = mHandles[key];
+        if (handle == nullptr || handle->getSize() < size)
+        {
+            if (size < PREFER_SIZE)
+            {
+                handle = std::make_shared<IpcNvlsHandleWrapper>(PREFER_SIZE, key.group);
+                TLLM_LOG_WARNING(
+                    "NvlsMemoryManager allocate D: key=%s, size=%zu bytes.", key.toString().c_str(), PREFER_SIZE);
+            }
+            else
+            {
+                handle = std::make_shared<IpcNvlsHandleWrapper>(size, key.group);
+                TLLM_LOG_WARNING("NvlsMemoryManager allocate D: key=%s, size=%zu bytes.", key.toString().c_str(), size);
+            }
+            mHandles[key] = handle;
+        }
+
+        return mHandles[key]->getHandle();
+    }
+
+private:
+    std::unordered_map<AllocationKey, std::shared_ptr<PersistentWorkspaceInterface>, AllocationKeyHash> mWorkspaces;
+    std::unordered_map<AllocationKey, std::shared_ptr<IpcNvlsHandleWrapper>, AllocationKeyHash> mHandles;
+};
+
+NvlsMemoryManager gNvlsMemoryManager;
+} // namespace
+
+namespace torch_ext
+{
+
+class Fp4GemmAllreduceRunner : public torch::CustomClassHolder
+{
+public:
+    explicit Fp4GemmAllreduceRunner(at::ScalarType outputDtype, int64_t rank, torch::List<int64_t> group)
+        : mOutputDtype(outputDtype)
+        , mRank(rank)
+    {
+        for (int64_t rank : group)
+        {
+            mGroup.insert(static_cast<int>(rank));
+        }
+
+        if (outputDtype == at::ScalarType::Half)
+        {
+            using Traits = GemmTypes<cutlass::float_e2m1_t, cutlass::float_e2m1_t, cutlass::half_t, cutlass::half_t,
+                cutlass::float_ue4m3_t, cutlass::float_ue4m3_t, cutlass::layout::RowMajor, cutlass::layout::ColumnMajor,
+                cutlass::layout::RowMajor, cutlass::layout::RowMajor>;
+            mRunner = std::make_shared<GemmAllReduceImplRunner<Traits>>();
+        }
+        else if (outputDtype == at::ScalarType::BFloat16)
+        {
+            using Traits = GemmTypes<cutlass::float_e2m1_t, cutlass::float_e2m1_t, cutlass::bfloat16_t,
+                cutlass::bfloat16_t, cutlass::float_ue4m3_t, cutlass::float_ue4m3_t, cutlass::layout::RowMajor,
+                cutlass::layout::ColumnMajor, cutlass::layout::RowMajor, cutlass::layout::RowMajor>;
+            mRunner = std::make_shared<GemmAllReduceImplRunner<Traits>>();
+        }
+        else
+        {
+            C10_THROW_ERROR(NotImplementedError, "Unsupported input or output dtype");
+        }
+
+        mConfigs = mRunner->getSupportedLaunchConfigs();
+    }
+
+    at::Tensor runGemm(at::Tensor const& mat1, at::Tensor const& mat2, at::Tensor const& mat1Scale,
+        at::Tensor const& mat2Scale, at::Tensor const& alpha, int64_t configIdx) const
+    {
+        if (configIdx < 0)
+            configIdx = 0;
+
+        TORCH_CHECK(configIdx < int64_t(mConfigs.size()), "configIdx out of bounds");
+        const int64_t M = mat1.size(0);
+        const int64_t N = mat2.size(0);
+        const int64_t K = mat1.size(1) * 2;
+
+        GemmAllReduceImplInterface::ProblemArgs problemArgs;
+        problemArgs.argProblemShape(M, N, K, 1);
+        problemArgs.argA(mat1.data_ptr());
+        problemArgs.argB(mat2.data_ptr());
+        problemArgs.argAScale(mat1Scale.data_ptr());
+        problemArgs.argBScale(mat2Scale.data_ptr());
+        problemArgs.argC(nullptr);
+        problemArgs.argAlphaPtr(reinterpret_cast<float const*>(alpha.const_data_ptr()));
+        problemArgs.argBeta(0.f);
+        problemArgs.argRanks(mRank, mGroup);
+        problemArgs.argLaunchConfig(mConfigs[configIdx]);
+
+        size_t dSize = M * N * c10::elementSize(mOutputDtype);
+        auto stream = at::cuda::getCurrentCUDAStream(mat1.get_device());
+
+        auto handle = gNvlsMemoryManager.getD(dSize, AllocationKey{stream, mGroup});
+        problemArgs.argD((void*) handle->uc_ptr, (void*) handle->mc_ptr, (void**) handle->ipc_uc_ptrs.data());
+
+        auto workspace = gNvlsMemoryManager.getWorkspace(mRunner.get(), problemArgs, AllocationKey{stream, mGroup});
+        problemArgs.argWorkspace(workspace);
+        mRunner->run(problemArgs, stream);
+        auto options = mat1.options().dtype(mOutputDtype);
+        auto D = at::from_blob((void*) handle->uc_ptr, {M, N}, {N, 1}, options);
+
+        return D;
+    }
+
+    int64_t getNumConfigs() const
+    {
+        return static_cast<int64_t>(mConfigs.size());
+    }
+
+private:
+    at::ScalarType mOutputDtype;
+    int mRank;
+    std::set<int> mGroup;
+    std::shared_ptr<GemmAllReduceImplInterface> mRunner{nullptr};
+    std::vector<GemmAllReduceImplInterface::LaunchConfig> mConfigs;
+};
+
+} // namespace torch_ext
+
+TORCH_LIBRARY_FRAGMENT(trtllm, m)
+{
+    m.class_<torch_ext::Fp4GemmAllreduceRunner>("Fp4GemmAllreduceRunner")
+        .def(torch::init<at::ScalarType, int64_t, torch::List<int64_t>>())
+        .def("run_gemm", &torch_ext::Fp4GemmAllreduceRunner::runGemm)
+        .def("get_num_configs", &torch_ext::Fp4GemmAllreduceRunner::getNumConfigs);
+}

tensorrt_llm/_torch/custom_ops/torch_custom_ops.py

@@ -1621,3 +1621,99 @@ def record_stream(tensor: torch.Tensor, stream_id: int) -> None:
     stream = get_stream(stream_id)
     assert stream is not None
     tensor.record_stream(stream)
+
+
+class Fp4GemmAllreduceRunner(TunableRunner):
+    runner_dict = dict()
+    tuning_config = TuningConfig(dynamic_tensor_specs=(DynamicTensorSpec(
+        0, 0, get_last_power_of_2_num_tokens_buckets,
+        last_positive_power_of_2), ),
+                                 constraint_specs=(ConstraintSpec(
+                                     2, 0, fp4_scale_infer_shape), ))
+
+    def __init__(
+        self,
+        output_dtype: torch.dtype,
+        tp_rank: int,
+        tp_group: List[int],
+    ):
+        self.output_dtype = output_dtype
+        self.tp_rank = tp_rank
+        self.tp_group_str = '-'.join(str(g) for g in tp_group)
+        instance_key = (output_dtype, self.tp_group_str)
+        if instance_key not in Fp4GemmAllreduceRunner.runner_dict:
+            Fp4GemmAllreduceRunner.runner_dict[
+                instance_key] = torch.classes.trtllm.Fp4GemmAllreduceRunner(
+                    output_dtype, tp_rank, tp_group)
+        self.fp4_gemm_all_reduce_runner = Fp4GemmAllreduceRunner.runner_dict[
+            instance_key]
+
+    def unique_id(self):
+        return (self.output_dtype, self.tp_group_str)
+
+    def get_valid_tactics(self, inputs: List[torch.Tensor],
+                          profile: OptimizationProfile, **kwargs) -> List[int]:
+        return list(range(self.fp4_gemm_all_reduce_runner.get_num_configs()))
+
+    def forward(
+        self,
+        inputs: List[torch.Tensor],
+        tactic: int = 0,
+    ) -> torch.Tensor:
+        mat1, mat2, mat1_scale, mat2_scale, global_scale = inputs
+        return self.fp4_gemm_all_reduce_runner.run_gemm(
+            mat1,
+            mat2,
+            mat1_scale,
+            mat2_scale,
+            global_scale,
+            tactic,
+        )
+
+
+@torch.library.custom_op("trtllm::nvfp4_gemm_allreduce", mutates_args=())
+def nvfp4_gemm_allreduce(
+    act_fp4: torch.Tensor,
+    weight: torch.Tensor,
+    act_sf: torch.Tensor,
+    weight_scale: torch.Tensor,
+    alpha: torch.Tensor,
+    output_dtype: torch.dtype,
+    tp_rank: int,
+    tp_group: List[int],
+) -> torch.Tensor:
+    AutoTuner.get()
+
+    # Use Cutlass runner with predefined configs
+    nvfp4_gemm_allreduce_runner = Fp4GemmAllreduceRunner(
+        output_dtype, tp_rank, tp_group)
+
+    # TODO: Enable auto-tuning
+    # runner_type = type(nvfp4_gemm_allreduce_runner).__name__
+    # _, best_tactic = tuner.choose_one(
+    #     f"trtllm::nvfp4_gemm_allreduce::{runner_type}",
+    #     [nvfp4_gemm_allreduce_runner],
+    #     nvfp4_gemm_allreduce_runner.tuning_config,
+    #     [act_fp4, weight, act_sf, weight_scale, alpha],
+    # )
+
+    best_tactic = -1
+
+    return nvfp4_gemm_allreduce_runner(
+        inputs=[act_fp4, weight, act_sf, weight_scale, alpha],
+        tactic=best_tactic)
+
+
+@nvfp4_gemm_allreduce.register_fake
+def _(
+    act_fp4: torch.Tensor,
+    weight: torch.Tensor,
+    act_sf: torch.Tensor,
+    weight_scale: torch.Tensor,
+    alpha: torch.Tensor,
+    output_dtype: torch.dtype,
+    tp_rank: int,
+    tp_group: List[int],
+) -> torch.Tensor:
+    return act_fp4.new_empty((act_fp4.size(0), weight.size(0)),
+                             dtype=output_dtype)