DeepSeek V3 FP8 Support

cpp/micro_benchmarks/mixtureOfExpertsBackendBenchmarkFixture.h

@@ -425,6 +425,8 @@ class MixtureOfExpertsBenchmark : public ::benchmark::Fixture
     QuantParams mQuantParams{};
     bool mUseLora = false;
     LoraParams mLoraParams{};
+    bool mUseDeepSeek = false;
+    BlockScaleParams mDeepseekParams{};
 
     std::optional<tensorrt_llm::cutlass_extensions::CutlassGemmConfig> mSelectedConfig = std::nullopt;
 
@@ -678,7 +680,7 @@ class MixtureOfExpertsBenchmark : public ::benchmark::Fixture
         mMoERunner.runMoe(mInputTensor, mInputProbabilities, mExpertWeight1, mExpertBias1, mActType, mExpertWeight2,
             mExpertBias2, mQuantParams, mTotalTokens, mHiddenSize, mInterSize, mNumExperts, mK, mWorkspace,
             mFinalOutput, nullptr, mTotalTokens, mScaleProbs, mSourceToExpandedMap, mSelectedExpert, 0.01,
-            parallelism_config, mNormMode, mUseLora, mLoraParams, stream);
+            parallelism_config, mNormMode, mUseLora, mLoraParams, mUseDeepSeek, mDeepseekParams, stream);
     }
 
     void runBenchmark(benchmark::State& state);

cpp/tensorrt_llm/CMakeLists.txt

@@ -369,6 +369,7 @@ set(TRTLLM_LINK_LIBS
     selective_scan_src
     fpA_intB_gemm_src
     moe_gemm_src
+    fp8_blockscale_gemm_lib
     fb_gemm_src
     gemm_swiglu_sm90_src
     cutlass_src
@@ -377,6 +378,13 @@ set(TRTLLM_LINK_LIBS
     userbuffers_src
     ${DECODER_SHARED_TARGET})
 
+set(BLOCKSCALEGEMM_LIB_LOC
+    "${CMAKE_CURRENT_SOURCE_DIR}/kernels/cutlass_kernels/fp8_blockscale_gemm/libfp8_blockscale_gemm.a"
+)
+add_library(fp8_blockscale_gemm_lib STATIC IMPORTED)
+set_property(TARGET fp8_blockscale_gemm_lib PROPERTY IMPORTED_LOCATION
+                                                     ${BLOCKSCALEGEMM_LIB_LOC})
+
 if(ENABLE_MULTI_DEVICE)
   set(TRTLLM_LINK_LIBS ${TRTLLM_LINK_LIBS} ${MPI_C_LIBRARIES} ${NCCL_LIB})
 endif()

cpp/tensorrt_llm/kernels/cutlass_kernels/CMakeLists.txt

@@ -78,6 +78,7 @@ file(GLOB_RECURSE SRC_CU *.cu)
 set(ALL_SRCS ${SRC_CPP};${SRC_CU})
 list(FILTER ALL_SRCS EXCLUDE REGEX "fpA_intB_gemm/.*")
 list(FILTER ALL_SRCS EXCLUDE REGEX "moe_gemm/.*")
+list(FILTER ALL_SRCS EXCLUDE REGEX "fp8_blockscale_gemm/.*")
 list(FILTER ALL_SRCS EXCLUDE REGEX "fp8_rowwise_gemm/.*")
 list(REMOVE_ITEM ALL_SRCS
      "${CMAKE_CURRENT_SOURCE_DIR}/fused_gated_gemm/gemm_swiglu_e4m3.cu")
@@ -90,6 +91,7 @@ message(
   "Group srcs ${GROUPED_SRC_CU} ${GROUPED_SRC_CPP} ${GROUPED_CU_INSTANTIATIONS}"
 )
 message(VERBOSE "Fbgemm srcs ${FBGEMM_SRC_CU} ${FBGEMM_CU_INSTANTIATIONS}")
+message(VERBOSE "Blockscale srcs ${BLOCKSCALEGEMM_SRC_CU} ")
 message(VERBOSE "All srcs ${ALL_SRCS}")
 
 add_library(cutlass_src STATIC ${ALL_SRCS})

cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_gemm.h

@@ -0,0 +1,102 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 1993-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * 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.
+ */
+
+#pragma once
+#include <cuda_runtime_api.h>
+#include <vector>
+#include <cstdint>
+#include <cuda_fp8.h>
+
+// non-persistent-cooperative GEMM
+namespace tensorrt_llm::kernels
+{
+namespace small_m_gemm
+{
+
+class CutlassFp8BlockScaleGemmRunnerInterface
+{
+public:
+    CutlassFp8BlockScaleGemmRunnerInterface() {}
+
+    virtual ~CutlassFp8BlockScaleGemmRunnerInterface() {}
+
+    virtual void gemm(void* mat_d, void const* mat_a, void const* mat_b, int shape_m, int shape_n, int shape_k,
+        char* workspace_ptr, cudaStream_t stream, float const* scales_a = nullptr, float const* scales_b = nullptr)
+        = 0;
+
+    virtual void gemm(__nv_fp8_e4m3 const* mat_a, int ld_a, __nv_fp8_e4m3 const* mat_b, int ld_b, __nv_bfloat16* mat_d, int ld_d,
+                int shape_m, int shape_n, int shape_k, float const* scales_a, float const* scales_b, cudaStream_t stream) = 0;
+
+    virtual void moeGemm(void *mat_d, void const *mat_a, void const *mat_b, const int64_t *problem_m_offsets, size_t num_problems, size_t shape_n,
+                         size_t shape_k, char *workspace_ptr, cudaStream_t stream, float const *scales_a = nullptr,
+                         float const *scales_b = nullptr) = 0;
+
+    virtual void strideBatchGemm(__nv_bfloat16* mat_d, int ld_d, int stride_d, __nv_fp8_e4m3* mat_a, int ld_a, int stride_a,
+        __nv_fp8_e4m3* mat_b, int ld_b, int stride_b, int num_problems, int shape_m, int shape_n, int shape_k, cudaStream_t stream, float* scales_a, int stride_scales_a, float* scales_b) = 0;
+
+    virtual void fp8CS1x128(__nv_fp8_e4m3* mat_quant, float* scales, __nv_bfloat16 const* mat, int shape_x, int shape_y, cudaStream_t stream) = 0;
+    virtual void fp8CS1x128Reshape(__nv_fp8_e4m3* mat_quant, float* scales, __nv_bfloat16 const* mat, int shape_x, int shape_h, int shape_y, int stride_x, cudaStream_t stream) = 0;
+    virtual void fp8CS128x128(__nv_fp8_e4m3* mat_quant, float* scales, __nv_bfloat16 const* mat, int shape_x, int shape_y, cudaStream_t stream) = 0;
+    // Returns desired workspace size in bytes.
+    virtual size_t getWorkspaceSize(size_t max_shape_m, size_t shape_n, size_t shape_k, size_t num_problems = 1) = 0;
+
+    virtual size_t getFP8DataSize(int shape_m, int shape_n, bool is_act) = 0;
+    virtual size_t getActScaleSize(int shape_m, int shape_k) = 0;
+    virtual size_t getWeightScaleSize(int shape_n, int shape_k) = 0;
+    virtual size_t getActWorkspaceSize(int shape_m, int shape_k) = 0;
+    virtual size_t getWeightWorkspaceSize(int shape_n, int shape_k) = 0;
+};
+
+template <typename ElementA, typename ElementB, typename ElementD>
+class CutlassFp8BlockScaleGemmRunner : public virtual CutlassFp8BlockScaleGemmRunnerInterface
+{
+public:
+    CutlassFp8BlockScaleGemmRunner();
+    ~CutlassFp8BlockScaleGemmRunner();
+
+    void gemm(void* mat_d, void const* mat_a, void const* mat_b, int shape_m, int shape_n, int shape_k,
+        char* workspace_ptr, cudaStream_t stream, float const* scales_a = nullptr,
+        float const* scales_b = nullptr) override;
+
+    void gemm(__nv_fp8_e4m3 const* mat_a, int ld_a, __nv_fp8_e4m3 const* mat_b, int ld_b, __nv_bfloat16* mat_d, int ld_d,
+                int shape_m, int shape_n, int shape_k, float const* scales_a, float const* scales_b, cudaStream_t stream) override;
+
+    void moeGemm(void *mat_d, void const *mat_a, void const *mat_b, const int64_t *problem_m_offsets, size_t num_problems, size_t shape_n,
+                 size_t shape_k, char *workspace_ptr, cudaStream_t stream, float const *scales_a = nullptr,
+                 float const *scales_b = nullptr) override;
+
+    void strideBatchGemm(__nv_bfloat16* mat_d, int ld_d, int stride_d, __nv_fp8_e4m3* mat_a, int ld_a, int stride_a,
+        __nv_fp8_e4m3* mat_b, int ld_b, int stride_b, int num_problems, int shape_m, int shape_n, int shape_k, cudaStream_t stream, float* scales_a, int stride_scales_a, float* scales_b) override;
+
+    void fp8CS1x128(__nv_fp8_e4m3* mat_quant, float* scales, __nv_bfloat16 const* mat, int shape_x, int shape_y, cudaStream_t stream) override;
+    void fp8CS1x128Reshape(__nv_fp8_e4m3* mat_quant, float* scales, __nv_bfloat16 const* mat, int shape_x, int shape_h, int shape_y, int stride_x, cudaStream_t stream) override;
+    void fp8CS128x128(__nv_fp8_e4m3* mat_quant, float* scales, __nv_bfloat16 const* mat, int shape_x, int shape_y, cudaStream_t stream) override;
+
+    // Returns desired workspace size in bytes.
+    size_t getWorkspaceSize(size_t max_shape_m, size_t shape_n, size_t shape_k, size_t num_problems = 1) override;
+
+    size_t getFP8DataSize(int shape_m, int shape_n, bool is_act) override;
+    size_t getActScaleSize(int shape_m, int shape_k) override;
+    size_t getWeightScaleSize(int shape_n, int shape_k) override;
+    size_t getActWorkspaceSize(int shape_m, int shape_k) override;
+    size_t getWeightWorkspaceSize(int shape_n, int shape_k) override;
+private:
+    int64_t max_shape_m_4_align_ = 0;
+};
+
+} // namespace small_m_gemm
+} // namespace tensorrt_llm::kernels

cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_kernels_bf16_fp8.cu

@@ -0,0 +1,25 @@
+/*
+ * 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 "tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_kernels_template.h"
+
+namespace tensorrt_llm
+{
+
+#ifdef ENABLE_BF16
+template class MoeGemmRunner<__nv_bfloat16, __nv_fp8_e4m3, __nv_bfloat16>;
+#endif
+} // namespace tensorrt_llm