diff --git a/cpp/micro_benchmarks/mixtureOfExpertsBackendBenchmarkFixture.h b/cpp/micro_benchmarks/mixtureOfExpertsBackendBenchmarkFixture.h index ec3da705f6d9..12819f48845a 100644 --- a/cpp/micro_benchmarks/mixtureOfExpertsBackendBenchmarkFixture.h +++ b/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 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); diff --git a/cpp/tensorrt_llm/CMakeLists.txt b/cpp/tensorrt_llm/CMakeLists.txt index ad979fe60e0b..c165eed6a9ad 100644 --- a/cpp/tensorrt_llm/CMakeLists.txt +++ b/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() diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/CMakeLists.txt b/cpp/tensorrt_llm/kernels/cutlass_kernels/CMakeLists.txt index a457b77aa5a9..a805bed007ea 100644 --- a/cpp/tensorrt_llm/kernels/cutlass_kernels/CMakeLists.txt +++ b/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}) diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_gemm.h b/cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_gemm.h new file mode 100644 index 000000000000..810c59d613e9 --- /dev/null +++ b/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 +#include +#include +#include + +// 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 +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 diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/libfp8_blockscale_gemm.a b/cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/libfp8_blockscale_gemm.a new file mode 100644 index 000000000000..2e99d3ca4d36 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/libfp8_blockscale_gemm.a @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:3130be05e0b4f136a40b5ac2726ff90994f7abba0206ca205fc9345a273c34b7 +size 2123196 diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_kernels_bf16_fp8.cu b/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_kernels_bf16_fp8.cu new file mode 100644 index 000000000000..cc7b0d0fb244 --- /dev/null +++ b/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 diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_kernels_template.h b/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_kernels_template.h index 2a337e6ca4ec..3dbd4c50a5e5 100644 --- a/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_kernels_template.h +++ b/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_kernels_template.h @@ -45,6 +45,7 @@ #include "cutlass_extensions/gemm/kernel/default_fpA_intB_traits.h" #include "cutlass_extensions/gemm/kernel/moe_cutlass_kernel.h" #include "cutlass_extensions/gemm/threadblock/default_mma.h" +#include #ifdef __GNUC__ // Restore GCC-specific diagnostics #pragma GCC diagnostic pop @@ -74,126 +75,155 @@ namespace kernels::cutlass_kernels { // ============================= Variable batched Gemm things =========================== + +template +struct isGeneralMoe +{ + + static constexpr bool value + = (std::is_same_v && std::is_same_v) == false; +}; + template -void genericMoeGemmKernelLauncher(T const* A, WeightType const* B, GemmOutputType const* weight_scales, - GemmOutputType const* biases, bool bias_is_broadcast, GemmOutputType* C, - int64_t const* total_tokens_including_expert, int64_t num_rows, int64_t gemm_n, int64_t gemm_k, int num_experts, - cutlass_extensions::CutlassGemmConfig gemm_config, int const multi_processor_count, bool use_fused_moe, - float const** alpha_scale_ptr_array, cudaStream_t stream, int* kernel_occupancy = nullptr) +struct genericMoeGemmKernelLauncher { + static void call(T const* A, WeightType const* B, GemmOutputType const* weight_scales, GemmOutputType const* biases, + bool bias_is_broadcast, GemmOutputType* C, int64_t const* total_tokens_including_expert, int64_t num_rows, + int64_t gemm_n, int64_t gemm_k, int num_experts, cutlass_extensions::CutlassGemmConfig gemm_config, + int const multi_processor_count, bool use_fused_moe, float const** alpha_scale_ptr_array, cudaStream_t stream, + int* kernel_occupancy = nullptr) + { #if defined(ENABLE_FP8) - static_assert(cutlass::platform::is_same::value || cutlass::platform::is_same::value - || cutlass::platform::is_same::value - || cutlass::platform::is_same::value || cutlass::platform::is_same::value, - "Specialized for fp8, bfloat16, half, float"); + static_assert(cutlass::platform::is_same::value || cutlass::platform::is_same::value + || cutlass::platform::is_same::value + || cutlass::platform::is_same::value || cutlass::platform::is_same::value, + "Specialized for fp8, bfloat16, half, float"); #elif defined(ENABLE_BF16) - static_assert(cutlass::platform::is_same::value || cutlass::platform::is_same::value - || cutlass::platform::is_same::value, - "Specialized for bfloat16, half, float"); + static_assert(cutlass::platform::is_same::value || cutlass::platform::is_same::value + || cutlass::platform::is_same::value, + "Specialized for bfloat16, half, float"); #else - static_assert(cutlass::platform::is_same::value || cutlass::platform::is_same::value, - "Specialized for half, float"); + static_assert(cutlass::platform::is_same::value || cutlass::platform::is_same::value, + "Specialized for half, float"); #endif - static_assert(cutlass::platform::is_same::value - || cutlass::platform::is_same::value - || cutlass::platform::is_same::value, - ""); + static_assert(cutlass::platform::is_same::value + || cutlass::platform::is_same::value + || cutlass::platform::is_same::value, + ""); - static_assert(!cutlass::platform::is_same::value, - "Sm90 architecture should use specialised kernels"); + static_assert(!cutlass::platform::is_same::value, + "Sm90 architecture should use specialised kernels"); - // The cutlass type for the input elements. This is needed to convert to cutlass::half_t if necessary. - using ElementType = typename TllmToCutlassTypeAdapter::type; - using CutlassGemmOutputType = typename TllmToCutlassTypeAdapter::type; - using CutlassWeightType = typename TllmToCutlassTypeAdapter::type; - if (!use_fused_moe) - { - // We need separate config for each architecture since we will target different tensorcore instructions. For - // float, we do not target TCs. - using MixedGemmArchTraits = cutlass::gemm::kernel::MixedGemmArchTraits; - using ElementAccumulator = typename MixedGemmArchTraits::AccType; + // The cutlass type for the input elements. This is needed to convert to cutlass::half_t if necessary. + using ElementType = typename TllmToCutlassTypeAdapter::type; + using CutlassGemmOutputType = typename TllmToCutlassTypeAdapter::type; + using CutlassWeightType = typename TllmToCutlassTypeAdapter::type; + if (!use_fused_moe) + { + // We need separate config for each architecture since we will target different tensorcore instructions. For + // float, we do not target TCs. + using MixedGemmArchTraits + = cutlass::gemm::kernel::MixedGemmArchTraits; + using ElementAccumulator = typename MixedGemmArchTraits::AccType; - using EpilogueOp = typename tensorrt_llm::cutlass_extensions::Epilogue::Op; + using EpilogueOp = typename tensorrt_llm::cutlass_extensions::Epilogue::Op; - typename EpilogueOp::Params epilogue_op( - ElementAccumulator(1.f), biases ? ElementAccumulator(1.f) : ElementAccumulator(0.f)); + typename EpilogueOp::Params epilogue_op( + ElementAccumulator(1.f), biases ? ElementAccumulator(1.f) : ElementAccumulator(0.f)); #if defined(ENABLE_FP8) - if constexpr ((std::is_same_v - || std::is_same_v) &&std::is_same_v) - { - TLLM_CHECK_WITH_INFO(weight_scales == nullptr && biases == nullptr && alpha_scale_ptr_array, - "weight_scales and biases should be nullptr and alpha_scale_ptr_array shouldn't be nullptr for FP8 " - "Ada"); - epilogue_op.alpha_ptr_array = alpha_scale_ptr_array; - } + if constexpr ((std::is_same_v + || std::is_same_v) &&std::is_same_v) + { + TLLM_CHECK_WITH_INFO(weight_scales == nullptr && biases == nullptr && alpha_scale_ptr_array, + "weight_scales and biases should be nullptr and alpha_scale_ptr_array shouldn't be nullptr for FP8 " + "Ada"); + epilogue_op.alpha_ptr_array = alpha_scale_ptr_array; + } #endif - // Finally, set up the kernel. - using GemmKernel_ = typename cutlass::gemm::kernel::DefaultGemmGrouped::GemmKernel; - - using GemmKernel = cutlass::gemm::kernel::MoeFCGemm; - - using GemmGrouped = cutlass::gemm::device::GemmGrouped; - - if (kernel_occupancy != nullptr) + // Finally, set up the kernel. + using GemmKernel_ = typename cutlass::gemm::kernel::DefaultGemmGrouped::GemmKernel; + + using GemmKernel = cutlass::gemm::kernel::MoeFCGemm; + + using GemmGrouped = cutlass::gemm::device::GemmGrouped; + + if (kernel_occupancy != nullptr) + { + *kernel_occupancy = tensorrt_llm::cutlass_extensions::compute_occupancy_for_kernel(); + return; + } + int occupancy = std::min(2, GemmGrouped::maximum_active_blocks()); + TLLM_CHECK_WITH_INFO(occupancy > 0, "GPU lacks the shared memory resources to run GroupedGEMM kernel"); + int const threadblock_count = multi_processor_count * occupancy; + + int const group_size = gemm_k; + typename GemmGrouped::Arguments args(num_experts, threadblock_count, group_size, epilogue_op, + reinterpret_cast(A), reinterpret_cast(B), + reinterpret_cast(weight_scales), + reinterpret_cast(biases), bias_is_broadcast, + reinterpret_cast(C), total_tokens_including_expert, gemm_n, gemm_k); + + GemmGrouped gemm; + + auto can_implement = gemm.can_implement(args); + TLLM_CHECK_WITH_INFO(can_implement == cutlass::Status::kSuccess, + "MoE FC kernel will fail for params. Error: " + std::string(cutlassGetStatusString(can_implement))); + + auto init_status = gemm.initialize(args); + TLLM_CHECK_WITH_INFO(init_status == cutlass::Status::kSuccess, + "Failed to initialize cutlass grouped gemm. Error: " + + std::string(cutlassGetStatusString(init_status))); + + auto run_status = gemm.run(stream); + TLLM_CHECK_WITH_INFO(run_status == cutlass::Status::kSuccess, + "Failed to run cutlass grouped gemm. Error: " + std::string(cutlassGetStatusString(run_status))); + } + else if constexpr (sizeof(ElementType) == 2 && sizeof(CutlassWeightType) == 2 + && (std::is_same_v + || std::is_same_v) ) // use fused moe gemm + // kernel.. (only support + // fp16 or bf16) { - *kernel_occupancy = tensorrt_llm::cutlass_extensions::compute_occupancy_for_kernel(); - return; + sm80_generic_fused_moe_gemm_kernelLauncher( + reinterpret_cast(A), reinterpret_cast(B), + reinterpret_cast(biases), bias_is_broadcast, reinterpret_cast(C), + total_tokens_including_expert, num_rows, gemm_n, gemm_k, num_experts, multi_processor_count, stream, + kernel_occupancy); } - int occupancy = std::min(2, GemmGrouped::maximum_active_blocks()); - TLLM_CHECK_WITH_INFO(occupancy > 0, "GPU lacks the shared memory resources to run GroupedGEMM kernel"); - int const threadblock_count = multi_processor_count * occupancy; - - int const group_size = gemm_k; - typename GemmGrouped::Arguments args(num_experts, threadblock_count, group_size, epilogue_op, - reinterpret_cast(A), reinterpret_cast(B), - reinterpret_cast(weight_scales), - reinterpret_cast(biases), bias_is_broadcast, - reinterpret_cast(C), total_tokens_including_expert, gemm_n, gemm_k); - - GemmGrouped gemm; - - auto can_implement = gemm.can_implement(args); - TLLM_CHECK_WITH_INFO(can_implement == cutlass::Status::kSuccess, - "MoE FC kernel will fail for params. Error: " + std::string(cutlassGetStatusString(can_implement))); - - auto init_status = gemm.initialize(args); - TLLM_CHECK_WITH_INFO(init_status == cutlass::Status::kSuccess, - "Failed to initialize cutlass grouped gemm. Error: " + std::string(cutlassGetStatusString(init_status))); - - auto run_status = gemm.run(stream); - TLLM_CHECK_WITH_INFO(run_status == cutlass::Status::kSuccess, - "Failed to run cutlass grouped gemm. Error: " + std::string(cutlassGetStatusString(run_status))); } - else if constexpr (sizeof(ElementType) == 2 && sizeof(CutlassWeightType) == 2 - && (std::is_same_v - || std::is_same_v) ) // use fused moe gemm - // kernel.. (only support - // fp16 or bf16) +}; + +template +struct genericMoeGemmKernelLauncher<__nv_bfloat16, __nv_fp8_e4m3, GemmOutputType, arch, EpilogueTag, ThreadblockShape, + WarpShape, Stages> +{ + static void call(__nv_bfloat16 const* A, __nv_fp8_e4m3 const* B, GemmOutputType const* weight_scales, + GemmOutputType const* biases, bool bias_is_broadcast, GemmOutputType* C, + int64_t const* total_tokens_including_expert, int64_t num_rows, int64_t gemm_n, int64_t gemm_k, int num_experts, + cutlass_extensions::CutlassGemmConfig gemm_config, int const multi_processor_count, bool use_fused_moe, + float const** alpha_scale_ptr_array, cudaStream_t stream, int* kernel_occupancy = nullptr) { - sm80_generic_fused_moe_gemm_kernelLauncher(reinterpret_cast(A), - reinterpret_cast(B), reinterpret_cast(biases), - bias_is_broadcast, reinterpret_cast(C), total_tokens_including_expert, num_rows, gemm_n, - gemm_k, num_experts, multi_processor_count, stream, kernel_occupancy); } -} - +}; } // namespace kernels::cutlass_kernels template ) ) { kernels::cutlass_kernels::genericMoeGemmKernelLauncher(A, B, weight_scales, biases, bias_is_broadcast, C, + ThreadblockShape, WarpShape, Stages>::call(A, B, weight_scales, biases, bias_is_broadcast, C, total_tokens_including_expert, num_rows, gemm_n, gemm_k, num_experts, gemm_config, multi_processor_count, use_fused_moe, alpha_scale_ptr_array, stream, occupancy); } diff --git a/cpp/tensorrt_llm/kernels/mixtureOfExperts/moe_kernels.cu b/cpp/tensorrt_llm/kernels/mixtureOfExperts/moe_kernels.cu index 992ac406826e..1858f4aa6758 100644 --- a/cpp/tensorrt_llm/kernels/mixtureOfExperts/moe_kernels.cu +++ b/cpp/tensorrt_llm/kernels/mixtureOfExperts/moe_kernels.cu @@ -1429,7 +1429,7 @@ template CutlassMoeFCRunner::getWorkspaceDeviceBufferSizes( int64_t const num_rows, int64_t const hidden_size, int64_t const inter_size, int const num_experts, int const num_experts_per_node, int const k, ActivationType activation_type, - MOEExpertScaleNormalizationMode norm_mode, bool use_lora) const + MOEExpertScaleNormalizationMode norm_mode, bool use_lora, bool use_deepseek) const { size_t const num_moe_inputs = k * num_rows; size_t const permuted_elems = num_moe_inputs * hidden_size; @@ -1512,14 +1512,24 @@ std::vector CutlassMoeFCRunner workspace{source_rows_size, // - permuted_rows_size, // - permuted_experts_size, // - expert_first_token_offset_size, // - sparse_mixer_out_size, // - softmax_out_size, // - permuted_scales_size, // - sorter_size, // + if (use_deepseek) + { + int factor = is_gated_activation ? 2 : 1; + int blockscale_fc1_output_size = factor * interbuf_elems * gemm_output_dtype; + int blockscale_fc2_output_size = permuted_elems * gemm_output_dtype; + overlapped_gemm1_gemm2_inputs = std::max(permuted_data_size, fc1_result_size); + overlapped_gemm1_gemm2_outputs = std::max(blockscale_fc1_output_size, blockscale_fc2_output_size); + } + + std::vector workspace{ + source_rows_size, // + permuted_rows_size, // + permuted_experts_size, // + expert_first_token_offset_size, // + sparse_mixer_out_size, // + softmax_out_size, // + permuted_scales_size, // + sorter_size, // // These pointers reuse the same memory overlapped_gemm1_gemm2_inputs, // overlapped_gemm1_gemm2_outputs, // @@ -1529,7 +1539,8 @@ std::vector CutlassMoeFCRunner::getWorkspaceSize(int64_t const num_rows, int64_t const hidden_size, int64_t const inter_size, int const num_experts, int const k, ActivationType activation_type, MOEExpertScaleNormalizationMode norm_mode, MOEParallelismConfig parallelism_config, - bool use_lora) const + bool use_lora, bool use_deepseek) const { int const ep_size = parallelism_config.ep_size; TLLM_CHECK_WITH_INFO(num_experts % ep_size == 0, "Number of experts must be a multiple of ep size"); - auto workspace = getWorkspaceDeviceBufferSizes( - num_rows, hidden_size, inter_size, num_experts, num_experts / ep_size, k, activation_type, norm_mode, use_lora); + auto workspace = getWorkspaceDeviceBufferSizes(num_rows, hidden_size, inter_size, num_experts, + num_experts / ep_size, k, activation_type, norm_mode, use_lora, use_deepseek); auto ws_size = tensorrt_llm::common::calculateTotalWorkspaceSize(workspace.data(), workspace.size()); TLLM_LOG_DEBUG("Mixture Of Experts Plugin requires workspace of %2f MiB", ws_size / 1024.f / 1024.f); return ws_size; @@ -1552,11 +1563,10 @@ template ::configureWsPtrs(char* ws_ptr, int64_t const num_rows, int64_t const hidden_size, int64_t const inter_size, int const num_experts, int const num_experts_per_node, int const k, ActivationType activation_type, - MOEExpertScaleNormalizationMode norm_mode, bool use_lora) - + MOEExpertScaleNormalizationMode norm_mode, bool use_lora, bool use_deepseek) { - auto ws_sizes = getWorkspaceDeviceBufferSizes( - num_rows, hidden_size, inter_size, num_experts, num_experts_per_node, k, activation_type, norm_mode, use_lora); + auto ws_sizes = getWorkspaceDeviceBufferSizes(num_rows, hidden_size, inter_size, num_experts, num_experts_per_node, + k, activation_type, norm_mode, use_lora, use_deepseek); std::vector ws_sliced{(int8_t*) ws_ptr}; for (auto size : ws_sizes) @@ -1628,6 +1638,14 @@ void CutlassMoeFCRunner::confi lora_add_bias_ = (ScaleBiasType*) ws_sliced[15]; lora_fc2_result_ = (ScaleBiasType*) ws_sliced[16]; } + + if (use_deepseek) + { + permuted_data_ = (T*) ws_sliced[8]; + fc1_result_ = (T*) ws_sliced[8]; + glu_inter_result_ = (T*) ws_sliced[9]; + fc2_result_ = (T*) ws_sliced[9]; + } } void sortAndScanSoftmaxOutput(int* expert_for_source_row, int* source_rows, int* permuted_experts, int* permuted_rows, @@ -1682,7 +1700,6 @@ void CutlassMoeFCRunner::gemm1 num_experts_per_node, input, fc1_expert_weights, fc1_fp8_dequant, nullptr, static_cast(gemm_output), stream); sync_check_cuda_error(); - gemm_runner.moeGemm(input, nullptr, nullptr, nullptr, total_tokens_including_expert, hopper_input, expanded_num_rows, fc1_out_size, hidden_size, num_experts_per_node, false, alpha_scale_ptr_array, stream, config); @@ -1835,6 +1852,58 @@ void CutlassMoeFCRunner::gemm2 sync_check_cuda_error(); } +template +void CutlassMoeFCRunner::BlockScaleFC1(T const* const input, + T* const output, void* const gemm_output, int64_t const* const expert_first_token_offset, + WeightType const* const fc1_expert_weights, ScaleBiasType const* const fc1_expert_biases, + float const* const fc2_fp8_quant, int64_t const expanded_num_rows, int64_t const hidden_size, + int64_t const inter_size, int const num_experts_per_node, ActivationType fc1_activation_type, + BlockScaleParams& deepseek_params, cudaStream_t stream) +{ + auto gemm_runner = deepseek_params.blockscale_gemm_iml; + TLLM_CHECK_WITH_INFO(gemm_runner, "blockscale gemm runner must be instantiated"); + bool const is_gated_activation = isGatedActivation(fc1_activation_type); + + int shape_n = is_gated_activation ? inter_size * 2 : inter_size; + int shape_k = hidden_size; + + gemm_runner->moeGemm(gemm_output, input, fc1_expert_weights, expert_first_token_offset, num_experts_per_node, + shape_n, shape_k, deepseek_params.workspace, stream, nullptr, deepseek_params.fc1_scales_ptrs); + + sync_check_cuda_error(); + constexpr bool bias_is_broadcast = true; + doActivation(output, static_cast(gemm_output), + fc2_fp8_quant, fc1_expert_biases, bias_is_broadcast, expert_first_token_offset, num_experts_per_node, + inter_size, expanded_num_rows, fc1_activation_type, stream); + + sync_check_cuda_error(); +} + +template +void CutlassMoeFCRunner::BlockScaleFC2(T const* const input, + void* const gemm_output, OutputType* const final_output, int64_t const* const expert_first_token_offset, + WeightType const* const fc2_expert_weights, ScaleBiasType const* const fc2_expert_biases, + float const* const token_topk_unpermuted_scales, int const* const expanded_source_row_to_expanded_dest_row, + int const* const expert_for_source_row, int64_t const* const num_valid_tokens_ptr, int64_t const num_rows, + int64_t const expanded_num_rows, int64_t const hidden_size, int64_t const inter_size, + int const num_experts_per_node, int64_t const k, BlockScaleParams& deepseek_params, cudaStream_t stream, + MOEParallelismConfig parallelism_config) +{ + int shape_n = hidden_size; + int shape_k = inter_size; + auto gemm_runner = deepseek_params.blockscale_gemm_iml; + TLLM_CHECK_WITH_INFO(gemm_runner, "blockscale gemm runner must be instantiated"); + gemm_runner->moeGemm(gemm_output, input, fc2_expert_weights, expert_first_token_offset, num_experts_per_node, + shape_n, shape_k, deepseek_params.workspace, stream, nullptr, deepseek_params.fc2_scales_ptrs); + + sync_check_cuda_error(); + + finalizeMoeRoutingKernelLauncher( + static_cast(gemm_output), final_output, fc2_expert_biases, + token_topk_unpermuted_scales, expanded_source_row_to_expanded_dest_row, expert_for_source_row, num_rows, + hidden_size, k, num_valid_tokens_ptr, parallelism_config, MOEExpertScaleNormalizationMode::NONE, stream); +} + template bool CutlassMoeFCRunner::setupLoraWorkspace(int64_t expanded_num_rows, int64_t num_rows, int64_t inter_size, int64_t hidden_size, int start_expert, bool is_gated_activation, @@ -2037,7 +2106,8 @@ void CutlassMoeFCRunner::runMo int const num_experts, int const k, char* workspace_ptr, void* final_output_void, bool const* finished, int64_t const active_rows, void* token_topk_final_scales_void, int* expanded_source_row_to_expanded_dest_row, int* expert_for_source_row, float sparse_mixer_epsilon, MOEParallelismConfig parallelism_config, - MOEExpertScaleNormalizationMode normalization_mode, bool use_lora, LoraParams& lora_params, cudaStream_t stream) + MOEExpertScaleNormalizationMode normalization_mode, bool use_lora, LoraParams& lora_params, bool use_deepseek, + BlockScaleParams& deepseek_params, cudaStream_t stream) { static constexpr bool int_scales_required = std::is_same::value || std::is_same::value; @@ -2097,7 +2167,7 @@ void CutlassMoeFCRunner::runMo TLLM_CHECK_WITH_INFO(fc1_fp8_dequant == nullptr && fc2_fp8_quant == nullptr && fc2_fp8_dequant == nullptr, "FP8 scales are provided for integer quantization"); } - else if (fp8_scales_required) + else if (fp8_scales_required && (!use_deepseek)) { TLLM_CHECK_WITH_INFO(fc1_expert_biases == nullptr, "Bias is not supported with FP8"); TLLM_CHECK_WITH_INFO(fc2_expert_biases == nullptr, "Bias is not supported with FP8"); @@ -2133,7 +2203,7 @@ void CutlassMoeFCRunner::runMo int const num_experts_per_node = num_experts / parallelism_config.ep_size; configureWsPtrs(workspace_ptr, num_rows, hidden_size, inter_size, num_experts, num_experts_per_node, k, - fc1_activation_type, normalization_mode, use_lora); + fc1_activation_type, normalization_mode, use_lora, use_deepseek); int const start_expert = num_experts_per_node * parallelism_config.ep_rank; int const end_expert = start_expert + num_experts_per_node; @@ -2192,10 +2262,19 @@ void CutlassMoeFCRunner::runMo } } - Self::gemm1(moe_gemm_runner_, permuted_data_, fc1_result_, glu_inter_result_, expert_first_token_offset_, - hopper_grouped_gemm_input_, fc1_expert_weights, fc1_expert_biases, num_valid_tokens_ptr, fc1_int_scales, - fc1_fp8_dequant, fc2_fp8_quant, expanded_num_rows, hidden_size, inter_size, num_experts_per_node, - fc1_activation_type, alpha_scale_ptr_array_, !use_lora, stream, *gemm1_config_); + if (use_deepseek) + { + BlockScaleFC1(permuted_data_, fc1_result_, glu_inter_result_, expert_first_token_offset_, fc1_expert_weights, + fc1_expert_biases, fc2_fp8_quant, expanded_num_rows, hidden_size, inter_size, num_experts_per_node, + fc1_activation_type, deepseek_params, stream); + } + else + { + Self::gemm1(moe_gemm_runner_, permuted_data_, fc1_result_, glu_inter_result_, expert_first_token_offset_, + hopper_grouped_gemm_input_, fc1_expert_weights, fc1_expert_biases, num_valid_tokens_ptr, fc1_int_scales, + fc1_fp8_dequant, fc2_fp8_quant, expanded_num_rows, hidden_size, inter_size, num_experts_per_node, + fc1_activation_type, alpha_scale_ptr_array_, !use_lora, stream, *gemm1_config_); + } sync_check_cuda_error(); @@ -2206,13 +2285,22 @@ void CutlassMoeFCRunner::runMo sync_check_cuda_error(); } - Self::gemm2(moe_gemm_runner_, fc1_result_, fc2_result_, final_output, expert_first_token_offset_, - hopper_grouped_gemm_input_, fc2_expert_weights, fc2_expert_biases, fc2_int_scales, fc2_fp8_dequant, - token_topk_unpermuted_scales, permuted_scales_, expanded_source_row_to_expanded_dest_row, permuted_rows_, - expert_for_source_row, num_valid_tokens_ptr, num_rows, expanded_num_rows, hidden_size, inter_size, - num_experts_per_node, k, !use_deterministic_hopper_reduce_, alpha_scale_ptr_array_, use_lora, lora_fc2_result_, - stream, parallelism_config, *gemm2_config_); - + if (use_deepseek) + { + BlockScaleFC2(fc1_result_, fc2_result_, final_output, expert_first_token_offset_, fc2_expert_weights, + fc2_expert_biases, token_topk_unpermuted_scales, expanded_source_row_to_expanded_dest_row, + expert_for_source_row, num_valid_tokens_ptr, num_rows, expanded_num_rows, hidden_size, inter_size, + num_experts_per_node, k, deepseek_params, stream, parallelism_config); + } + else + { + Self::gemm2(moe_gemm_runner_, fc1_result_, fc2_result_, final_output, expert_first_token_offset_, + hopper_grouped_gemm_input_, fc2_expert_weights, fc2_expert_biases, fc2_int_scales, fc2_fp8_dequant, + token_topk_unpermuted_scales, permuted_scales_, expanded_source_row_to_expanded_dest_row, permuted_rows_, + expert_for_source_row, num_valid_tokens_ptr, num_rows, expanded_num_rows, hidden_size, inter_size, + num_experts_per_node, k, use_deterministic_hopper_reduce_, alpha_scale_ptr_array_, use_lora, + lora_fc2_result_, stream, parallelism_config, *gemm2_config_); + } sync_check_cuda_error(); } @@ -2631,6 +2719,7 @@ template class CutlassMoeFCRunner; template class CutlassMoeFCRunner<__nv_fp8_e4m3, __nv_fp8_e4m3, half>; #ifdef ENABLE_BF16 template class CutlassMoeFCRunner<__nv_fp8_e4m3, __nv_fp8_e4m3, __nv_bfloat16>; +template class CutlassMoeFCRunner<__nv_bfloat16, __nv_fp8_e4m3, __nv_bfloat16>; #endif #endif diff --git a/cpp/tensorrt_llm/kernels/mixtureOfExperts/moe_kernels.h b/cpp/tensorrt_llm/kernels/mixtureOfExperts/moe_kernels.h index fdf5fe9ba5e0..8a003fb9727a 100644 --- a/cpp/tensorrt_llm/kernels/mixtureOfExperts/moe_kernels.h +++ b/cpp/tensorrt_llm/kernels/mixtureOfExperts/moe_kernels.h @@ -19,6 +19,7 @@ #include "cutlass/gemm/gemm.h" #include "tensorrt_llm/common/assert.h" #include "tensorrt_llm/common/quantization.h" +#include "tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_gemm.h" #include "tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_kernels.h" #include "tensorrt_llm/kernels/lora/lora.h" #include @@ -208,13 +209,38 @@ struct LoraParams } }; +struct BlockScaleParams +{ + using BlockScaleGemmImplPtr = std::shared_ptr; + BlockScaleGemmImplPtr blockscale_gemm_iml; + + float const* fc1_scales_ptrs = nullptr; + float const* fc2_scales_ptrs = nullptr; + + char* workspace; + + cudaEvent_t* memcpy_event_ptr; + + BlockScaleParams() = default; + + BlockScaleParams(float const* fc1_scales_ptrs, float const* fc2_scales_ptrs, + BlockScaleGemmImplPtr blockscale_gemm_iml, char* workspace, cudaEvent_t* memcpy_event_ptr) + : fc1_scales_ptrs(fc1_scales_ptrs) + , fc2_scales_ptrs(fc2_scales_ptrs) + , blockscale_gemm_iml(blockscale_gemm_iml) + , workspace(workspace) + , memcpy_event_ptr(memcpy_event_ptr) + { + } +}; + class CutlassMoeFCRunnerInterface { public: virtual ~CutlassMoeFCRunnerInterface() = default; virtual size_t getWorkspaceSize(int64_t const num_rows, int64_t const hidden_size, int64_t const inter_size, int const num_experts, int const k, ActivationType activation_type, MOEExpertScaleNormalizationMode norm_mode, - MOEParallelismConfig parallelism_config, bool use_lora) const + MOEParallelismConfig parallelism_config, bool use_lora, bool use_deepseek) const = 0; virtual void setTactic(std::optional gemm1_config, std::optional gemm2_config) @@ -228,7 +254,7 @@ class CutlassMoeFCRunnerInterface bool const* finished, int64_t const active_rows, void* token_topk_unpermuted_scales, int* expanded_source_row_to_expanded_dest_row, int* expert_for_source_row, float sparse_mixer_epsilon, MOEParallelismConfig parallelism_config, MOEExpertScaleNormalizationMode normalization_mode, bool use_lora, - LoraParams& lora_params, cudaStream_t stream) + LoraParams& lora_params, bool use_deepseek, BlockScaleParams& deepseek_params, cudaStream_t stream) = 0; // Aliases for profiling the gemms @@ -255,7 +281,6 @@ class CutlassMoeFCRunnerInterface = 0; virtual size_t getGemmWorkspaceSize(int num_experts) const = 0; - bool is_profiler = false; bool use_deterministic_hopper_reduce_ = false; }; @@ -295,7 +320,7 @@ class CutlassMoeFCRunner : public CutlassMoeFCRunnerInterface size_t getWorkspaceSize(int64_t const num_rows, int64_t const hidden_size, int64_t const fc1_output_size, int const num_experts, int const k, ActivationType activation_type, MOEExpertScaleNormalizationMode norm_mode, - MOEParallelismConfig parallelism_config, bool use_lora) const override; + MOEParallelismConfig parallelism_config, bool use_lora, bool use_deepseek) const override; void setTactic(std::optional gemm1_config, std::optional gemm2_config) override @@ -322,7 +347,7 @@ class CutlassMoeFCRunner : public CutlassMoeFCRunnerInterface bool const* finished, int64_t const active_rows, void* token_topk_unpermuted_scales, int* expanded_source_row_to_expanded_dest_row, int* expert_for_source_row, float sparse_mixer_epsilon, MOEParallelismConfig parallelism_config, MOEExpertScaleNormalizationMode normalization_mode, bool use_lora, - LoraParams& lora_params, cudaStream_t stream) override; + LoraParams& lora_params, bool use_deepseek, BlockScaleParams& deepseek_params, cudaStream_t stream) override; // We make these GEMM1 & GEMM2 static because they need to be stateless for the profiler to work static void gemm1(MoeGemmRunner& gemm_runner, T const* const input, @@ -398,10 +423,11 @@ class CutlassMoeFCRunner : public CutlassMoeFCRunnerInterface cudaStream_t stream); std::vector getWorkspaceDeviceBufferSizes(int64_t const num_rows, int64_t const hidden_size, int64_t const inter_size, int const num_experts, int const num_experts_per_node, int const k, - ActivationType activation_type, MOEExpertScaleNormalizationMode norm_mode, bool use_lora) const; + ActivationType activation_type, MOEExpertScaleNormalizationMode norm_mode, bool use_lora, + bool use_deepseek) const; void configureWsPtrs(char* ws_ptr, int64_t const num_rows, int64_t const hidden_size, int64_t const inter_size, int const num_experts, int const num_experts_per_node, int const k, ActivationType activation_type, - MOEExpertScaleNormalizationMode norm_mode, bool use_lora); + MOEExpertScaleNormalizationMode norm_mode, bool use_lora, bool use_deepseek); private: bool mayHaveDifferentGEMMOutputType() const @@ -429,6 +455,21 @@ class CutlassMoeFCRunner : public CutlassMoeFCRunnerInterface int64_t const* num_valid_tokens_ptr, int64_t num_tokens, LoraParams& lora_params, float const* fc2_fp8_quant, cudaStream_t stream); + void BlockScaleFC1(T const* const input, T* const output, void* const intermediate_result, + int64_t const* const expert_first_token_offset, WeightType const* const fc1_expert_weights, + ScaleBiasType const* const fc1_expert_biases, float const* const fc2_fp8_quant, int64_t const expanded_num_rows, + int64_t const hidden_size, int64_t const inter_size, int const num_experts_per_node, + ActivationType fc1_activation_type, BlockScaleParams& deepseek_params, cudaStream_t stream); + + void BlockScaleFC2(T const* const input, void* const gemm_output, OutputType* const final_output, + int64_t const* const expert_first_token_offset, WeightType const* const fc2_expert_weights, + ScaleBiasType const* const fc2_expert_biases, float const* const token_topk_unpermuted_scales, + int const* const expanded_source_row_to_expanded_dest_row, int const* const expert_for_source_row, + int64_t const* const num_valid_tokens_ptr, int64_t const num_rows, + int64_t const expanded_num_rows, int64_t const hidden_size, int64_t const inter_size, + int const num_experts_per_node, int64_t const k, BlockScaleParams& deepseek_params, cudaStream_t stream, + MOEParallelismConfig parallelism_config); + CubKeyValueSorter sorter_; MoeGemmRunner moe_gemm_runner_; @@ -471,6 +512,14 @@ class CutlassMoeFCRunner : public CutlassMoeFCRunnerInterface }; HostLoraWorkspace host_lora_workspace_; + + struct HostDeepSeekWorkspace + { + std::vector host_expert_first_token_offset; + std::vector host_num_tokens_per_expert; + }; + + HostDeepSeekWorkspace host_deepseek_workspace_; }; void makeLoadBalancedRoutingConfiguration( diff --git a/cpp/tensorrt_llm/kernels/mlaKernels.cu b/cpp/tensorrt_llm/kernels/mlaKernels.cu index 6394ded30c83..929bb95e6f0b 100644 --- a/cpp/tensorrt_llm/kernels/mlaKernels.cu +++ b/cpp/tensorrt_llm/kernels/mlaKernels.cu @@ -255,9 +255,9 @@ __global__ void applyMLARopeAndAssignQKVKernelOptContext(T* qkv_output, T const* } template -__global__ void applyMLARopeAndAssignQKVKernelGeneration(T* qkv_output, T const* fuse_buf, KVCacheBuffer kv_cache, - float2 const* cos_sin_cache, size_t head_num, int head_size, int c_q, int c_k, int total_s_len, int* seqQOffset, - uint32_t* fmha_tile_counter, int32_t const* kv_cache_lengths, int* seqKVOffsets) +__global__ void applyMLARopeAndAssignQKVKernelGeneration(T* qkv_output, T* q_buf, T const* fuse_buf, + KVCacheBuffer kv_cache, float2 const* cos_sin_cache, size_t head_num, int head_size, int c_q, int c_k, + int total_s_len, int* seqQOffset, uint32_t* fmha_tile_counter, int32_t const* kv_cache_lengths, int* seqKVOffsets) { // Constants. @@ -322,12 +322,12 @@ __global__ void applyMLARopeAndAssignQKVKernelGeneration(T* qkv_output, T const* } else { - auto const src_q_global_offset = static_cast(global_token_idx) * head_num * (c_k + ROPE_DIM) - + (c_k + ROPE_DIM) * head_idx + c_k; + auto const src_q_global_offset + = static_cast(global_token_idx) * head_num * (head_size + ROPE_DIM) + + (head_size + ROPE_DIM) * head_idx + head_size; for (int i = 0; i < 2; ++i) { - ref[i] = *reinterpret_cast( - &qkv_output[src_q_global_offset + src_bias + i * ELTS_PER_VEC]); + ref[i] = *reinterpret_cast(&q_buf[src_q_global_offset + src_bias + i * ELTS_PER_VEC]); } } @@ -451,7 +451,7 @@ void invokeMLARopeGeneration(mlaParams& params, KVCacheBuffer kv_cache_buffer dim3 grid(int(tensorrt_llm::common::divUp(params.acc_q_len, 32)), params.head_num + 1 + 8); auto head_size = params.meta.qk_nope_head_dim; applyMLARopeAndAssignQKVKernelGeneration - <<>>(params.attention_input_buf, params.fused_a_input, kv_cache_buffer, + <<>>(params.attention_input_buf, params.q_buf, params.fused_a_input, kv_cache_buffer, params.cos_sin_cache, params.head_num, head_size, params.meta.q_lora_rank, params.meta.kv_lora_rank, params.acc_q_len, params.seqQOffset, params.fmha_tile_counter, params.cache_seq_lens, params.cu_kv_seqlens); } diff --git a/cpp/tensorrt_llm/kernels/mlaKernels.h b/cpp/tensorrt_llm/kernels/mlaKernels.h index 74b683886dd9..94099a214b85 100644 --- a/cpp/tensorrt_llm/kernels/mlaKernels.h +++ b/cpp/tensorrt_llm/kernels/mlaKernels.h @@ -39,12 +39,16 @@ struct mlaMetaParams template struct mlaParams { - T const* fused_a_input; // [b, s, c_q + c_k + r] - T* attention_input_buf; // [b, s, 3, h, d_h + r] + T const* fused_a_input; // [b, s, c_q + c_k + r] + T* attention_input_buf; // [b, s, 3, h, d_h + r] T* context_buf; - T const* fused_q_proj; // [c_k + r, d] - T const* q_b_proj; // [(d_h + r) * h, c_q] - T const* kv_b_proj; // [h * d_h * 2, c_k] + T* q_buf; // [b, h, d_h + r] + T const* q_b_proj; // [(d_h + r) * h, c_q] + T const* kv_b_proj; // [h * d_h * 2, c_k] + T const* k_b_proj_trans; // [h * c_k, d_h] + float const* q_b_scale; + float const* kv_b_scale; + float const* k_b_trans_scale; float2 const* cos_sin_cache; // [s, rope] int32_t batch_size; int32_t acc_q_len; diff --git a/cpp/tensorrt_llm/plugins/CMakeLists.txt b/cpp/tensorrt_llm/plugins/CMakeLists.txt index 40fdff9f5cd4..c38aa095c193 100755 --- a/cpp/tensorrt_llm/plugins/CMakeLists.txt +++ b/cpp/tensorrt_llm/plugins/CMakeLists.txt @@ -40,6 +40,7 @@ set(PLUGIN_LISTS identityPlugin gemmPlugin gemmSwigluPlugin + fp8CurrentScalingGemmPlugin fp8RowwiseGemmPlugin smoothQuantGemmPlugin quantizePerTokenPlugin diff --git a/cpp/tensorrt_llm/plugins/api/tllmPlugin.cpp b/cpp/tensorrt_llm/plugins/api/tllmPlugin.cpp index 8c54ba94c054..f05580ee4a2f 100644 --- a/cpp/tensorrt_llm/plugins/api/tllmPlugin.cpp +++ b/cpp/tensorrt_llm/plugins/api/tllmPlugin.cpp @@ -20,6 +20,7 @@ #include "tensorrt_llm/runtime/tllmLogger.h" #include "tensorrt_llm/plugins/bertAttentionPlugin/bertAttentionPlugin.h" +#include "tensorrt_llm/plugins/fp8CurrentScalingGemmPlugin/fp8CurrentScalingGemmPlugin.h" #include "tensorrt_llm/plugins/fp8RowwiseGemmPlugin/fp8RowwiseGemmPlugin.h" #include "tensorrt_llm/plugins/gemmPlugin/gemmPlugin.h" #include "tensorrt_llm/plugins/gemmSwigluPlugin/gemmSwigluPlugin.h" @@ -207,6 +208,7 @@ extern "C" static tensorrt_llm::plugins::GemmPluginCreator gemmPluginCreator; static tensorrt_llm::plugins::GemmSwigluPluginCreator gemmSwigluPluginCreator; static tensorrt_llm::plugins::Fp8RowwiseGemmPluginCreator fp8RowwiseGemmPluginCreator; + static tensorrt_llm::plugins::Fp8CurrentScalingGemmPluginCreator fp8CurrentScalingGemmPluginCreator; static tensorrt_llm::plugins::MixtureOfExpertsPluginCreator moePluginCreator; #if ENABLE_MULTI_DEVICE static tensorrt_llm::plugins::SendPluginCreator sendPluginCreator; @@ -245,6 +247,7 @@ extern "C" creatorPtr(gemmPluginCreator), creatorPtr(gemmSwigluPluginCreator), creatorPtr(fp8RowwiseGemmPluginCreator), + creatorPtr(fp8CurrentScalingGemmPluginCreator), creatorPtr(moePluginCreator), #if ENABLE_MULTI_DEVICE creatorPtr(sendPluginCreator), diff --git a/cpp/tensorrt_llm/plugins/fp8CurrentScalingGemmPlugin/CMakeLists.txt b/cpp/tensorrt_llm/plugins/fp8CurrentScalingGemmPlugin/CMakeLists.txt new file mode 100644 index 000000000000..3b714a3928fb --- /dev/null +++ b/cpp/tensorrt_llm/plugins/fp8CurrentScalingGemmPlugin/CMakeLists.txt @@ -0,0 +1,21 @@ +# +# SPDX-FileCopyrightText: Copyright (c) 1993-2022 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. +# +file(GLOB SRCS *.cpp *.cu) +set(PLUGIN_SOURCES ${PLUGIN_SOURCES} ${SRCS}) +set(PLUGIN_SOURCES + ${PLUGIN_SOURCES} + PARENT_SCOPE) diff --git a/cpp/tensorrt_llm/plugins/fp8CurrentScalingGemmPlugin/fp8CurrentScalingGemmPlugin.cpp b/cpp/tensorrt_llm/plugins/fp8CurrentScalingGemmPlugin/fp8CurrentScalingGemmPlugin.cpp new file mode 100644 index 000000000000..75b80bdc0553 --- /dev/null +++ b/cpp/tensorrt_llm/plugins/fp8CurrentScalingGemmPlugin/fp8CurrentScalingGemmPlugin.cpp @@ -0,0 +1,382 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 1993-2022 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. + */ + +#include "fp8CurrentScalingGemmPlugin.h" +#include "cutlass_extensions/gemm_configs.h" + +#include +#include + +using namespace nvinfer1; +using namespace tensorrt_llm::common; +using namespace tensorrt_llm::kernels::small_m_gemm; +using tensorrt_llm::plugins::Fp8CurrentScalingGemmPluginCreator; +using tensorrt_llm::plugins::Fp8CurrentScalingGemmPlugin; + +static char const* FP8_CURRENT_SCALING_GEMM_PLUGIN_VERSION{"1"}; +static char const* FP8_CURRENT_SCALING_GEMM_PLUGIN_NAME{"Fp8CurrentScalingGemm"}; +PluginFieldCollection Fp8CurrentScalingGemmPluginCreator::mFC{}; +std::vector Fp8CurrentScalingGemmPluginCreator::mPluginAttributes; + +Fp8CurrentScalingGemmPlugin::Fp8CurrentScalingGemmPlugin( + int need_quantize_acts_on_demand, int need_quantize_weights_on_demand, nvinfer1::DataType type) +{ + init(type, need_quantize_acts_on_demand, need_quantize_weights_on_demand); +} + +// Parameterized constructor +Fp8CurrentScalingGemmPlugin::Fp8CurrentScalingGemmPlugin(void const* data, size_t length) +{ + char const *d = reinterpret_cast(data), *a = d; + nvinfer1::DataType type; + int need_quantize_acts_on_demand = 0; + int need_quantize_weights_on_demand = 0; + + read(d, need_quantize_acts_on_demand); + read(d, need_quantize_weights_on_demand); + read(d, type); + + init(type, need_quantize_acts_on_demand, need_quantize_weights_on_demand); + + TLLM_CHECK_WITH_INFO(d == a + length, + "Expected length (%d) != real length (%d). This is often " + "caused by using different TensorRT-LLM version to build " + "engine and run engine.", + (int) length, (int) (d - a)); +} + +void Fp8CurrentScalingGemmPlugin::init( + nvinfer1::DataType type, int need_quantize_acts_on_demand, int need_quantize_weights_on_demand) +{ + mArch = tensorrt_llm::common::getSMVersion(); + mType = type; + mNeedQuantizeActsOnDemand = need_quantize_acts_on_demand; + mNeedQuantizeWeightsOnDemand = need_quantize_weights_on_demand; + mInputIdx = 0; + mWeightInputIdx = mInputIdx + 1; + mWeightScalesIdx = mNeedQuantizeWeightsOnDemand ? mWeightInputIdx : mWeightInputIdx + 1; + mInputScalesIdx = mNeedQuantizeActsOnDemand ? mWeightScalesIdx : mWeightScalesIdx + 1; + + if (mType == nvinfer1::DataType::kBF16) + { + if (mNeedQuantizeActsOnDemand && !mNeedQuantizeWeightsOnDemand) + { + mGemmRunner + = std::make_shared>(); + } + else if (mNeedQuantizeActsOnDemand && mNeedQuantizeWeightsOnDemand) + { + mGemmRunner + = std::make_shared>(); + } + else + { + mGemmRunner + = std::make_shared>(); + } + } + else + { + TLLM_THROW("Fp8 current scaling Gemm plugin doesn't support this type now"); + } +} + +// IPluginV2DynamicExt Methods +nvinfer1::IPluginV2DynamicExt* Fp8CurrentScalingGemmPlugin::clone() const noexcept +{ + auto* plugin = new Fp8CurrentScalingGemmPlugin(*this); + return plugin; +} + +nvinfer1::DimsExprs Fp8CurrentScalingGemmPlugin::getOutputDimensions( + int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, nvinfer1::IExprBuilder& exprBuilder) noexcept +{ + // inputs + // 0 activations [M, K] + // 1 weights [K, N] + // 2 weight scales [K // 128, N // 128] + // 3 activation scales [M, K // 128] (optional) + + try + { + TLLM_CHECK(nbInputs == mInputScalesIdx + 1); + TLLM_CHECK(outputIndex == 0); + int const nbDimsA = inputs[0].nbDims; + int const nbDimsB = inputs[mWeightInputIdx].nbDims; + TLLM_CHECK(nbDimsA >= 2); + TLLM_CHECK(nbDimsB == 2); + DimsExprs ret; + ret.nbDims = nbDimsA; + for (int ii = 0; ii < nbDimsA - 1; ++ii) + { + ret.d[ii] = inputs[0].d[ii]; + } + ret.d[nbDimsA - 1] = inputs[1].d[0]; + return ret; + } + catch (std::exception const& e) + { + caughtError(e); + } + return DimsExprs{}; +} + +bool Fp8CurrentScalingGemmPlugin::supportsFormatCombination( + int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept +{ + if (pos == mInputIdx) + { + // activations + if (mNeedQuantizeActsOnDemand) + { + return inOut[pos].type == mType && inOut[pos].format == TensorFormat::kLINEAR; + } + else + { + return inOut[pos].type == nvinfer1::DataType::kFP8 && inOut[pos].format == TensorFormat::kLINEAR; + } + } + else if (pos == mWeightInputIdx) + { + // weights + if (mNeedQuantizeWeightsOnDemand) + { + return inOut[pos].type == mType && inOut[pos].format == TensorFormat::kLINEAR; + } + else + { + return inOut[pos].type == nvinfer1::DataType::kFP8 && inOut[pos].format == TensorFormat::kLINEAR; + } + } + else if (!mNeedQuantizeWeightsOnDemand && pos == mWeightScalesIdx) + { + // weight scales + return inOut[pos].type == nvinfer1::DataType::kFLOAT && inOut[pos].format == TensorFormat::kLINEAR; + } + else if (!mNeedQuantizeActsOnDemand && pos == mInputScalesIdx) + { + // input scales + return inOut[pos].type == nvinfer1::DataType::kFLOAT && inOut[pos].format == TensorFormat::kLINEAR; + } + else if (pos == mInputScalesIdx + 1) + { + // outputs + return inOut[pos].type == mType && inOut[pos].format == TensorFormat::kLINEAR; + } + else + { + return false; + } +} + +void Fp8CurrentScalingGemmPlugin::configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, + nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept +{ + auto const minM = std::accumulate(in[0].min.d, in[0].min.d + in[0].min.nbDims - 1, 1, std::multiplies()); + auto const maxM = std::accumulate(in[0].max.d, in[0].max.d + in[0].max.nbDims - 1, 1, std::multiplies()); + + int const maxK = in[0].max.d[in[0].max.nbDims - 1]; + int const maxN = in[1].max.d[0]; + int const minK = in[0].min.d[in[0].min.nbDims - 1]; + int const minN = in[1].min.d[0]; + + TLLM_CHECK_WITH_INFO(minN == maxN, "Variable out channels is not allowed"); + TLLM_CHECK_WITH_INFO(minK == maxK, "Variable in channels is not allowed"); + + mWorkspaceMaxSize = mGemmRunner->getWorkspaceSize(maxM, maxN, maxK); +} + +size_t Fp8CurrentScalingGemmPlugin::getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, + nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept +{ + return mWorkspaceMaxSize; +} + +int Fp8CurrentScalingGemmPlugin::enqueue(nvinfer1::PluginTensorDesc const* inputDesc, + nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace, + cudaStream_t stream) noexcept +{ + // inputs + // 0 activations [M, K] + // 1 weights [K, N] + // 2 weight scales[K // 128, N // 128] + // 3 scales [M, K // 128] (optional) + // + // outputs + // mat [M, N] + int m = 1; + for (int ii = 0; ii < inputDesc[0].dims.nbDims - 1; ++ii) + { + m *= inputDesc[0].dims.d[ii]; + } + int const n = inputDesc[1].dims.d[0]; + int const k = inputDesc[0].dims.d[inputDesc[0].dims.nbDims - 1]; + // size_t const wsSize = mGemmRunner->getWorkspaceSize(m, n, k); + + mGemmRunner->gemm(outputs[0], inputs[0], inputs[1], m, n, k, reinterpret_cast(workspace), stream, + reinterpret_cast(inputs[3]), reinterpret_cast(inputs[2])); + sync_check_cuda_error(); + + return 0; +} + +// IPluginV2Ext Methods +nvinfer1::DataType Fp8CurrentScalingGemmPlugin::getOutputDataType( + int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept +{ + TLLM_CHECK(index == 0); + return mType; +} + +// IPluginV2 Methods + +char const* Fp8CurrentScalingGemmPlugin::getPluginType() const noexcept +{ + return FP8_CURRENT_SCALING_GEMM_PLUGIN_NAME; +} + +char const* Fp8CurrentScalingGemmPlugin::getPluginVersion() const noexcept +{ + return FP8_CURRENT_SCALING_GEMM_PLUGIN_VERSION; +} + +int Fp8CurrentScalingGemmPlugin::getNbOutputs() const noexcept +{ + return 1; +} + +int Fp8CurrentScalingGemmPlugin::initialize() noexcept +{ + // Modify here, maybe do nothing + configGemm(); // gemm profiler in action + return 0; +} + +void Fp8CurrentScalingGemmPlugin::terminate() noexcept {} + +size_t Fp8CurrentScalingGemmPlugin::getSerializationSize() const noexcept +{ + return sizeof(int) + // need_quantize_acts_on_demand + sizeof(int) + // need_quantize_weights_on_demand + sizeof(nvinfer1::DataType); // dtype +} + +void Fp8CurrentScalingGemmPlugin::serialize(void* buffer) const noexcept +{ + char *d = static_cast(buffer), *a = d; + write(d, mNeedQuantizeActsOnDemand); + write(d, mNeedQuantizeWeightsOnDemand); + write(d, mType); + + TLLM_CHECK(d == a + getSerializationSize()); +} + +void Fp8CurrentScalingGemmPlugin::destroy() noexcept +{ + // This gets called when the network containing plugin is destroyed + delete this; +} + +void Fp8CurrentScalingGemmPlugin::configGemm() {} + +Fp8CurrentScalingGemmPluginCreator::Fp8CurrentScalingGemmPluginCreator() +{ + // Fill PluginFieldCollection with PluginField arguments metadata + mPluginAttributes.clear(); + mPluginAttributes.emplace_back(PluginField("need_quantize_acts_on_demand", nullptr, PluginFieldType::kINT32, 1)); + mPluginAttributes.emplace_back(PluginField("need_quantize_weights_on_demand", nullptr, PluginFieldType::kINT32, 1)); + mPluginAttributes.emplace_back(PluginField("type_id", nullptr, PluginFieldType::kINT32, 1)); + mFC.nbFields = mPluginAttributes.size(); + mFC.fields = mPluginAttributes.data(); +} + +char const* Fp8CurrentScalingGemmPluginCreator::getPluginName() const noexcept +{ + return FP8_CURRENT_SCALING_GEMM_PLUGIN_NAME; +} + +char const* Fp8CurrentScalingGemmPluginCreator::getPluginVersion() const noexcept +{ + return FP8_CURRENT_SCALING_GEMM_PLUGIN_VERSION; +} + +PluginFieldCollection const* Fp8CurrentScalingGemmPluginCreator::getFieldNames() noexcept +{ + return &mFC; +} + +IPluginV2* Fp8CurrentScalingGemmPluginCreator::createPlugin(char const* name, PluginFieldCollection const* fc) noexcept +{ + PluginField const* fields = fc->fields; + TLLM_CHECK(fc->nbFields == 3); + int needQuantizeActsOnDemand; + int needQuantizeWeightsOnDemand; + + nvinfer1::DataType type; + // Read configurations from each fields + for (int i = 0; i < fc->nbFields; ++i) + { + char const* attrName = fields[i].name; + if (!strcmp(attrName, "need_quantize_acts_on_demand")) + { + TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); + needQuantizeActsOnDemand = static_cast(*(static_cast(fields[i].data))); + } + else if (!strcmp(attrName, "need_quantize_weights_on_demand")) + { + TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); + needQuantizeWeightsOnDemand = static_cast(*(static_cast(fields[i].data))); + } + else if (!strcmp(attrName, "type_id")) + { + TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); + type = static_cast(*(static_cast(fields[i].data))); + } + } + try + { + // Fp8CurrentScalingGemmPluginCreator is unique and shared for an engine generation + // Create plugin profiler with shared tactics map + auto* obj = new Fp8CurrentScalingGemmPlugin(needQuantizeActsOnDemand, needQuantizeWeightsOnDemand, type); + obj->setPluginNamespace(mNamespace.c_str()); + return obj; + } + catch (std::exception const& e) + { + caughtError(e); + } + return nullptr; +} + +IPluginV2* Fp8CurrentScalingGemmPluginCreator::deserializePlugin( + char const* name, void const* serialData, size_t serialLength) noexcept +{ + // This object will be deleted when the network is destroyed, which will + // call Fp8CurrentScalingGemmPlugin::destroy() + try + { + // Create plugin profiler with private tactics map which is read from the serialized engine + auto* obj = new Fp8CurrentScalingGemmPlugin(serialData, serialLength); + obj->setPluginNamespace(mNamespace.c_str()); + return obj; + } + catch (std::exception const& e) + { + caughtError(e); + } + return nullptr; +} diff --git a/cpp/tensorrt_llm/plugins/fp8CurrentScalingGemmPlugin/fp8CurrentScalingGemmPlugin.h b/cpp/tensorrt_llm/plugins/fp8CurrentScalingGemmPlugin/fp8CurrentScalingGemmPlugin.h new file mode 100644 index 000000000000..f38de4795f39 --- /dev/null +++ b/cpp/tensorrt_llm/plugins/fp8CurrentScalingGemmPlugin/fp8CurrentScalingGemmPlugin.h @@ -0,0 +1,116 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 1993-2022 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 + +// Add our real GEMM kernel +#include "tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_gemm.h" +#include "tensorrt_llm/plugins/common/gemmPluginProfiler.h" +#include "tensorrt_llm/plugins/common/plugin.h" +#include +#include +#include +#include + +namespace tensorrt_llm::plugins +{ + +using Fp8BlockScaleGemmRunnerPtr + = std::shared_ptr; + +class Fp8CurrentScalingGemmPlugin : public BasePlugin +{ +public: + Fp8CurrentScalingGemmPlugin() = delete; + + Fp8CurrentScalingGemmPlugin( + int need_quantize_acts_on_demand, int need_quantize_weights_on_demand, nvinfer1::DataType type); + + Fp8CurrentScalingGemmPlugin(void const* data, size_t length); + + ~Fp8CurrentScalingGemmPlugin() override = default; + + // IPluginV2DynamicExt Methods + nvinfer1::IPluginV2DynamicExt* clone() const noexcept override; + nvinfer1::DimsExprs getOutputDimensions(int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, + nvinfer1::IExprBuilder& exprBuilder) noexcept override; + bool supportsFormatCombination( + int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept override; + void configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, + nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept override; + size_t getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, + nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept override; + int enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, + void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept override; + + // IPluginV2Ext Methods + nvinfer1::DataType getOutputDataType( + int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept override; + + // IPluginV2 Methods + char const* getPluginType() const noexcept override; + char const* getPluginVersion() const noexcept override; + int getNbOutputs() const noexcept override; + int initialize() noexcept override; + void terminate() noexcept override; + size_t getSerializationSize() const noexcept override; + void serialize(void* buffer) const noexcept override; + void destroy() noexcept override; + +private: + void init(nvinfer1::DataType type, int need_quantize_acts_on_demand, int need_quantize_weights_on_demand); + + void configGemm(); + +private: + const std::string mLayerName; + + Fp8BlockScaleGemmRunnerPtr mGemmRunner; + size_t mWorkspaceMaxSize; + nvinfer1::DataType mType; + + int mArch; + int mNeedQuantizeActsOnDemand; + int mNeedQuantizeWeightsOnDemand; + + int mInputIdx; + int mWeightInputIdx; + int mWeightScalesIdx; + int mInputScalesIdx; +}; + +class Fp8CurrentScalingGemmPluginCreator : public BaseCreator +{ +public: + Fp8CurrentScalingGemmPluginCreator(); + + char const* getPluginName() const noexcept override; + + char const* getPluginVersion() const noexcept override; + + nvinfer1::PluginFieldCollection const* getFieldNames() noexcept override; + + nvinfer1::IPluginV2* createPlugin(char const* name, nvinfer1::PluginFieldCollection const* fc) noexcept override; + + nvinfer1::IPluginV2* deserializePlugin( + char const* name, void const* serialData, size_t serialLength) noexcept override; + +private: + static nvinfer1::PluginFieldCollection mFC; + static std::vector mPluginAttributes; +}; + +} // namespace tensorrt_llm::plugins diff --git a/cpp/tensorrt_llm/plugins/gptAttentionCommon/gptAttentionCommon.cpp b/cpp/tensorrt_llm/plugins/gptAttentionCommon/gptAttentionCommon.cpp index 36298d53539b..0e947bd2902a 100644 --- a/cpp/tensorrt_llm/plugins/gptAttentionCommon/gptAttentionCommon.cpp +++ b/cpp/tensorrt_llm/plugins/gptAttentionCommon/gptAttentionCommon.cpp @@ -426,8 +426,8 @@ GPTAttentionPluginCommon::GPTAttentionPluginCommon(int layer_idx, int num_heads, bool use_paged_context_fmha, bool use_fp8_context_fmha, bool has_full_attention_mask, bool use_cache, bool is_spec_decoding_enabled, bool spec_decoding_is_generation_length_variable, int32_t spec_decoding_max_generation_length, bool is_mla_enabled, int q_lora_rank, int kv_lora_rank, - int qk_nope_head_dim, int qk_rope_head_dim, int v_head_dim, bool skip_attn, int cp_size, int cp_rank, - std::set cp_group) + int qk_nope_head_dim, int qk_rope_head_dim, int v_head_dim, bool is_ptp128c_enabled, bool is_fp8_model, + bool skip_attn, int cp_size, int cp_rank, std::set cp_group) : mLayerIdx(layer_idx) , mNumHeads(num_heads) , mVisionStart(vision_start) @@ -478,6 +478,8 @@ GPTAttentionPluginCommon::GPTAttentionPluginCommon(int layer_idx, int num_heads, , mSpecDecodingMaxGenerationLength(spec_decoding_max_generation_length) , mIsMLAEnabled(is_mla_enabled) , mMLAParams({q_lora_rank, kv_lora_rank, qk_nope_head_dim, qk_rope_head_dim, v_head_dim}) + , mIsPTP128CEnabled(is_ptp128c_enabled) + , mIsFP8Model(is_fp8_model) , mCpSize(cp_size) , mCpRank(cp_rank) , mCpGroup(move(cp_group)) @@ -616,6 +618,8 @@ GPTAttentionPluginCommon::GPTAttentionPluginCommon(void const* data, size_t leng read(d, mSpecDecodingMaxGenerationLength); read(d, mIsMLAEnabled); read(d, mMLAParams); + read(d, mIsPTP128CEnabled); + read(d, mIsFP8Model); read(d, mNbMultiBlockSemaphores); read(d, mSkipAttn); read(d, mCpSize); @@ -725,16 +729,38 @@ size_t GPTAttentionPluginCommon::getWorkspaceSizeForGeneration( // MLA use fmha instead of XQA in generation if (mIsMLAEnabled) { - size_t cu_seqlens_size = sizeof(int) * (max_num_tokens + 1); + size_t cu_seqlens_size = sizeof(int) * (max_num_seq + 1); size_t fmha_scheduler_counter = sizeof(uint32_t); - size_t o_buffer_size = size * max_num_tokens * mNumHeads * mMLAParams.kv_lora_rank; - int const NUM_BUFFERS = 5; + size_t q_buffer_size = size * max_num_seq * mNumHeads * (mMLAParams.q_lora_rank + mMLAParams.qk_rope_head_dim); + size_t o_buffer_size = size * max_num_seq * mNumHeads * (mMLAParams.kv_lora_rank); + size_t act_buffer_size = 0; + size_t weight_buffer_size = 0; + if (mIsPTP128CEnabled) + { + act_buffer_size + = std::max(mGemmRunner->getActWorkspaceSize(max_num_seq, + mMLAParams.q_lora_rank + mMLAParams.kv_lora_rank + mMLAParams.qk_rope_head_dim), + mGemmRunner->getActWorkspaceSize( + max_num_seq, mNumHeads * (mMLAParams.kv_lora_rank + mMLAParams.qk_rope_head_dim))); + if (!mIsFP8Model) + { + weight_buffer_size = std::max(mGemmRunner->getWeightWorkspaceSize(mNumHeads + * (mMLAParams.qk_nope_head_dim + mMLAParams.qk_rope_head_dim), + mMLAParams.q_lora_rank), + mGemmRunner->getWeightWorkspaceSize( + mNumHeads * mMLAParams.qk_nope_head_dim * 2, mMLAParams.kv_lora_rank)); + } + } + int const NUM_BUFFERS = 8; size_t workspaces[NUM_BUFFERS]; workspaces[0] = CUBLAS_WORKSPACE_SIZE; workspaces[1] = cu_seqlens_size; // cu_q_len workspaces[2] = cu_seqlens_size; // cu_kv_len workspaces[3] = fmha_scheduler_counter; - workspaces[4] = o_buffer_size; + workspaces[4] = q_buffer_size; + workspaces[5] = o_buffer_size; + workspaces[6] = act_buffer_size; + workspaces[7] = weight_buffer_size; generation_workspace_size = tc::calculateTotalWorkspaceSize(workspaces, NUM_BUFFERS); return generation_workspace_size; } @@ -818,45 +844,164 @@ int GPTAttentionPluginCommon::mlaPreContext( auto c_k = mMLAParams.kv_lora_rank; auto context_head_size = mMLAParams.qk_nope_head_dim + mMLAParams.qk_rope_head_dim; auto v_head_dim = mMLAParams.v_head_dim; - - // kv = self.kv_b_proj(compressed_kv) [b*s, c_k] * [c_k, h * (d_h * 2)] -> [b*s, h * (d_h * 2)] + if (!mIsPTP128CEnabled) { - auto transa = CUBLAS_OP_T; - auto transb = CUBLAS_OP_N; - int m = params.head_num * context_head_size; - int n = params.acc_q_len; - int k = c_q; - int lda = k, ldb = c_q + c_k + rope_dim; - int ldc - = params.head_num * (2 * context_head_size + v_head_dim); // output shape: [(b * s), (3 * h * (d_h + rope)] - mCublasWrapper->createDescriptors(transa, transb, m, n, k, lda, ldb, ldc); - mCublasWrapper->Gemm( - transa, transb, m, n, k, params.q_b_proj, lda, params.fused_a_input, ldb, params.attention_input_buf, ldc); - mCublasWrapper->destroyDescriptors(); - } + { + auto transa = CUBLAS_OP_T; + auto transb = CUBLAS_OP_N; + int m = params.head_num * context_head_size; + int n = params.acc_q_len; + int k = c_q; + int lda = k, ldb = c_q + c_k + rope_dim; + int ldc = params.head_num + * (2 * context_head_size + v_head_dim); // output shape: [(b * s), (3 * h * (d_h + rope)] + mCublasWrapper->createDescriptors(transa, transb, m, n, k, lda, ldb, ldc); + mCublasWrapper->Gemm(transa, transb, m, n, k, params.q_b_proj, lda, params.fused_a_input, ldb, + params.attention_input_buf, ldc); + mCublasWrapper->destroyDescriptors(); + } + { + auto transa = CUBLAS_OP_T; + auto transb = CUBLAS_OP_N; + int m = mMLAParams.qk_nope_head_dim; + int n = params.acc_q_len; + int k = c_k; + // int lda = k, ldb = k + params.rope_dim + params.c_q; + // int ldc = (params.head_size + params.rope_dim) * params.head_num; //params.head_size * params.c_k; + int lda = k, ldb = c_q + c_k + rope_dim; + int ldc = params.head_num * (2 * context_head_size + v_head_dim); + mCublasWrapper->createDescriptors(transa, transb, m, n, k, lda, ldb, ldc); + mCublasWrapper->stridedBatchedGemm(transa, transb, m, n, k, params.kv_b_proj, lda, + mMLAParams.qk_nope_head_dim * c_k, params.fused_a_input + c_q, ldb, 0, + params.attention_input_buf + static_cast(params.head_num) * context_head_size, + params.head_num * (context_head_size * 2 + v_head_dim), context_head_size, params.head_num, 1.0f, 0.0); + + mCublasWrapper->Gemm(transa, transb, params.head_num * m, n, k, + params.kv_b_proj + mMLAParams.qk_nope_head_dim * c_k * params.head_num, lda, params.fused_a_input + c_q, + ldb, params.attention_input_buf + 2 * static_cast(params.head_num) * context_head_size, ldc); + mCublasWrapper->destroyDescriptors(); + } + } + else { - auto transa = CUBLAS_OP_T; - auto transb = CUBLAS_OP_N; - int m = mMLAParams.qk_nope_head_dim; - int n = params.acc_q_len; - int k = c_k; - // int lda = k, ldb = k + params.rope_dim + params.c_q; - // int ldc = (params.head_size + params.rope_dim) * params.head_num; //params.head_size * params.c_k; - int lda = k, ldb = c_q + c_k + rope_dim; - int ldc = params.head_num * (context_head_size * 2 + v_head_dim); - mCublasWrapper->createDescriptors(transa, transb, m, n, k, lda, ldb, ldc); - - mCublasWrapper->stridedBatchedGemm(transa, transb, m, n, k, params.kv_b_proj, lda, - mMLAParams.qk_nope_head_dim * c_k, params.fused_a_input + c_q, ldb, 0, - params.attention_input_buf + static_cast(params.head_num) * context_head_size, - params.head_num * (context_head_size * 2 + v_head_dim), context_head_size, params.head_num, 1.0f, 0.0); - mCublasWrapper->stridedBatchedGemm(transa, transb, m, n, k, - params.kv_b_proj + mMLAParams.qk_nope_head_dim * c_k * params.head_num, lda, v_head_dim * c_k, - params.fused_a_input + c_q, ldb, 0, - params.attention_input_buf + 2 * static_cast(params.head_num) * context_head_size, - params.head_num * (context_head_size * 2 + v_head_dim), v_head_dim, params.head_num, 1.0f, 0.0); - mCublasWrapper->destroyDescriptors(); + int8_t* workspace_byte_ptr = reinterpret_cast(params.workspace); + size_t offset = CUBLAS_WORKSPACE_SIZE; + + size_t act_size = mGemmRunner->getFP8DataSize(params.acc_q_len, c_q + c_k + rope_dim, true); + size_t act_scale_size = mGemmRunner->getActScaleSize(params.acc_q_len, c_q + c_k + rope_dim); + __nv_fp8_e4m3* act_buffer + = reinterpret_cast<__nv_fp8_e4m3*>(nextWorkspacePtr(workspace_byte_ptr, offset, act_size)); + float* act_scale_buffer + = reinterpret_cast(nextWorkspacePtr(workspace_byte_ptr, offset, act_scale_size)); + size_t act_offset = offset; + + mGemmRunner->fp8CS1x128(act_buffer, act_scale_buffer, + reinterpret_cast<__nv_bfloat16 const*>(params.fused_a_input), c_q + c_k + rope_dim, params.acc_q_len, + stream); + sync_check_cuda_error(); + + // mat_a: [params.acc_q_len, c_q], mat_b: [c_q, params.head_num * context_head_size] + // + { + int m = params.head_num * context_head_size; + int n = params.acc_q_len; + int k = c_q; + int lda = k, ldb = c_q + c_k + rope_dim; + int ldc = params.head_num + * (2 * context_head_size + v_head_dim); // output shape: [(b * s), (3 * h * (d_h + rope)] + + __nv_fp8_e4m3 const* q_gemm_weight_ptr; + float const* q_gemm_scale_ptr; + if (mIsFP8Model) + { + q_gemm_weight_ptr = reinterpret_cast<__nv_fp8_e4m3 const*>(params.q_b_proj); + q_gemm_scale_ptr = reinterpret_cast(params.q_b_scale); + } + else + { + size_t q_gemm_size = mGemmRunner->getFP8DataSize(params.head_num * context_head_size, c_q, false); + size_t q_gemm_scale_size = mGemmRunner->getWeightScaleSize(params.head_num * context_head_size, c_q); + __nv_fp8_e4m3* q_weight_buffer + = reinterpret_cast<__nv_fp8_e4m3*>(nextWorkspacePtr(workspace_byte_ptr, offset, q_gemm_size)); + float* q_scale_buffer + = reinterpret_cast(nextWorkspacePtr(workspace_byte_ptr, offset, q_gemm_scale_size)); + mGemmRunner->fp8CS128x128(q_weight_buffer, q_scale_buffer, + reinterpret_cast<__nv_bfloat16 const*>(params.q_b_proj), c_q, params.head_num * context_head_size, + stream); + + q_gemm_weight_ptr = reinterpret_cast<__nv_fp8_e4m3 const*>(q_weight_buffer); + q_gemm_scale_ptr = reinterpret_cast(q_scale_buffer); + } + + // npcgemm2::fp8_gemm_run(act_buffer, ldb, const_cast<__nv_fp8_e4m3*>(q_gemm_weight_ptr), lda, + // reinterpret_cast<__nv_bfloat16*>(params.attention_input_buf), ldc, n, m, k, act_scale_buffer, + // const_cast(q_gemm_scale_ptr), stream); + mGemmRunner->gemm(act_buffer, ldb, const_cast<__nv_fp8_e4m3*>(q_gemm_weight_ptr), lda, + reinterpret_cast<__nv_bfloat16*>(params.attention_input_buf), ldc, n, m, k, act_scale_buffer, + const_cast(q_gemm_scale_ptr), stream); + sync_check_cuda_error(); + } + + { + offset = act_offset; + __nv_fp8_e4m3 const* kv_gemm_weight_ptr; + float const* kv_gemm_scale_ptr; + if (mIsFP8Model) + { + kv_gemm_weight_ptr = reinterpret_cast<__nv_fp8_e4m3 const*>(params.kv_b_proj); + kv_gemm_scale_ptr = reinterpret_cast(params.kv_b_scale); + } + else + { + size_t kv_gemm_size + = mGemmRunner->getFP8DataSize(params.head_num * mMLAParams.qk_nope_head_dim * 2, c_k, false); + size_t kv_gemm_scale_size + = mGemmRunner->getWeightScaleSize(params.head_num * mMLAParams.qk_nope_head_dim * 2, c_k); + __nv_fp8_e4m3* kv_weight_buffer + = reinterpret_cast<__nv_fp8_e4m3*>(nextWorkspacePtr(workspace_byte_ptr, offset, kv_gemm_size)); + float* kv_scale_buffer + = reinterpret_cast(nextWorkspacePtr(workspace_byte_ptr, offset, kv_gemm_scale_size)); + mGemmRunner->fp8CS128x128(kv_weight_buffer, kv_scale_buffer, + reinterpret_cast<__nv_bfloat16 const*>(params.kv_b_proj), c_k, + params.head_num * mMLAParams.qk_nope_head_dim * 2, stream); + + kv_gemm_weight_ptr = reinterpret_cast<__nv_fp8_e4m3 const*>(kv_weight_buffer); + kv_gemm_scale_ptr = reinterpret_cast(kv_scale_buffer); + } + + auto transa = CUBLAS_OP_T; + auto transb = CUBLAS_OP_N; + int m = mMLAParams.qk_nope_head_dim; + int n = params.acc_q_len; + int k = c_k; + + int shape_n_4_align = ((params.acc_q_len - 1) / 4 + 1) * 4; + int lda = k, ldb = c_q + c_k + rope_dim; + int ldc = params.head_num * (2 * context_head_size + v_head_dim); + + mGemmRunner->strideBatchGemm(reinterpret_cast<__nv_bfloat16*>(params.attention_input_buf) + + static_cast(params.head_num) * context_head_size, + ldc, context_head_size, act_buffer + c_q, ldb, 0, const_cast<__nv_fp8_e4m3*>(kv_gemm_weight_ptr), lda, + m * k, params.head_num, shape_n_4_align, m, k, stream, act_scale_buffer + shape_n_4_align * (c_q / 128), + 0, const_cast(kv_gemm_scale_ptr)); + sync_check_cuda_error(); + // npcgemm2::fp8_gemm_run(act_buffer + c_q, ldb, + // const_cast<__nv_fp8_e4m3*>(kv_gemm_weight_ptr) + params.head_num * m * k, lda, + // reinterpret_cast<__nv_bfloat16*>( + // params.attention_input_buf + 2 * static_cast(params.head_num) * context_head_size), + // ldc, n, params.head_num * m, k, act_scale_buffer + shape_n_4_align * (c_q / 128), + // const_cast(kv_gemm_scale_ptr) + params.head_num * (k / 128) * (m / 128), stream); + + mGemmRunner->gemm(act_buffer + c_q, ldb, + const_cast<__nv_fp8_e4m3*>(kv_gemm_weight_ptr) + params.head_num * m * k, lda, + reinterpret_cast<__nv_bfloat16*>( + params.attention_input_buf + 2 * static_cast(params.head_num) * context_head_size), + ldc, n, params.head_num * m, k, act_scale_buffer + shape_n_4_align * (c_q / 128), + const_cast(kv_gemm_scale_ptr) + params.head_num * (k / 128) * (m / 128), stream); + sync_check_cuda_error(); + mCublasWrapper->destroyDescriptors(); + } } return 0; @@ -917,34 +1062,143 @@ int GPTAttentionPluginCommon::mlaGeneration( int8_t* workspace_byte_ptr = reinterpret_cast(params.workspace); size_t offset = CUBLAS_WORKSPACE_SIZE; - // output[b, s, :1, h * (d_h + rope)] = self.q_b_proj(q_buf) [b*s, c_q] * [c_q, h * (d_h + rope)] -> [b*s, h * (d_h - // + rope)] - { - auto transa = CUBLAS_OP_T; - auto transb = CUBLAS_OP_N; - int m = params.head_num * (mMLAParams.kv_lora_rank + mMLAParams.qk_rope_head_dim); - int n = params.acc_q_len; - int k = mMLAParams.q_lora_rank; - int lda = k, ldb = mMLAParams.q_lora_rank + mMLAParams.kv_lora_rank + mMLAParams.qk_rope_head_dim; - int ldc = m; - mCublasWrapper->createDescriptors(transa, transb, m, n, k, lda, ldb, ldc); - mCublasWrapper->Gemm(transa, transb, m, n, k, params.fused_q_proj, lda, params.fused_a_input, ldb, - params.attention_input_buf, ldc); - mCublasWrapper->destroyDescriptors(); - } - size_t const cu_seqlens_size = sizeof(int) * (params.batch_size + 1); size_t const fmha_scheduler_counter = sizeof(uint32_t); + size_t q_buffer_size = size * batch_beam * mNumHeads * (mMLAParams.qk_nope_head_dim + mMLAParams.qk_rope_head_dim); size_t o_buffer_size = size * batch_beam * mNumHeads * mMLAParams.kv_lora_rank; int* cu_q_seqlens = reinterpret_cast(nextWorkspacePtr(workspace_byte_ptr, offset, cu_seqlens_size)); int* cu_kv_seqlens = reinterpret_cast(nextWorkspacePtr(workspace_byte_ptr, offset, cu_seqlens_size)); uint32_t* fmha_tile_counter_ptr = reinterpret_cast(nextWorkspacePtr(workspace_byte_ptr, offset, fmha_scheduler_counter)); + T* q_buffer = reinterpret_cast(nextWorkspacePtr(workspace_byte_ptr, offset, q_buffer_size)); T* o_buffer = reinterpret_cast(nextWorkspacePtr(workspace_byte_ptr, offset, o_buffer_size)); + size_t quant_offset = offset; + + if (!mIsPTP128CEnabled) + { + { + auto transa = CUBLAS_OP_T; + auto transb = CUBLAS_OP_N; + int m = params.head_num * (mMLAParams.qk_nope_head_dim + mMLAParams.qk_rope_head_dim); + int n = params.acc_q_len; + int k = mMLAParams.q_lora_rank; + int lda = k, ldb = mMLAParams.q_lora_rank + mMLAParams.kv_lora_rank + mMLAParams.qk_rope_head_dim; + int ldc = m; + mCublasWrapper->createDescriptors(transa, transb, m, n, k, lda, ldb, ldc); + mCublasWrapper->Gemm( + transa, transb, m, n, k, params.q_b_proj, lda, params.fused_a_input, ldb, q_buffer, ldc); + mCublasWrapper->destroyDescriptors(); + } + + { + auto transa = CUBLAS_OP_T; + auto transb = CUBLAS_OP_N; + int m = mMLAParams.kv_lora_rank; + int n = params.acc_q_len; + int k = mMLAParams.qk_nope_head_dim; + int lda = k, ldb = params.head_num * (mMLAParams.qk_nope_head_dim + mMLAParams.qk_rope_head_dim); + int ldc = params.head_num * (mMLAParams.kv_lora_rank + mMLAParams.qk_rope_head_dim); + mCublasWrapper->createDescriptors(transa, transb, m, n, k, lda, ldb, ldc); + mCublasWrapper->stridedBatchedGemm(transa, transb, m, n, k, params.k_b_proj_trans, lda, + mMLAParams.qk_nope_head_dim * mMLAParams.kv_lora_rank, q_buffer, ldb, + mMLAParams.qk_nope_head_dim + mMLAParams.qk_rope_head_dim, params.attention_input_buf, ldc, + mMLAParams.kv_lora_rank + mMLAParams.qk_rope_head_dim, params.head_num, 1.0f, 0.0); + mCublasWrapper->destroyDescriptors(); + } + } + else + { + { + int m = params.head_num * (mMLAParams.qk_nope_head_dim + mMLAParams.qk_rope_head_dim); + int n = params.acc_q_len; + int k = mMLAParams.q_lora_rank; + int lda = k, ldb = mMLAParams.q_lora_rank + mMLAParams.kv_lora_rank + mMLAParams.qk_rope_head_dim; + int ldc = m; + size_t act_size = mGemmRunner->getFP8DataSize(n, ldb, true); + size_t act_scale_size = mGemmRunner->getActScaleSize(n, ldb); + __nv_fp8_e4m3* act_buffer + = reinterpret_cast<__nv_fp8_e4m3*>(nextWorkspacePtr(workspace_byte_ptr, offset, act_size)); + float* act_scale_buffer + = reinterpret_cast(nextWorkspacePtr(workspace_byte_ptr, offset, act_scale_size)); + __nv_fp8_e4m3 const* q_gemm_weight_ptr; + float const* q_gemm_scale_ptr; + if (mIsFP8Model) + { + q_gemm_weight_ptr = reinterpret_cast<__nv_fp8_e4m3 const*>(params.q_b_proj); + q_gemm_scale_ptr = reinterpret_cast(params.q_b_scale); + } + else + { + size_t q_gemm_size = mGemmRunner->getFP8DataSize(m, lda, false); + size_t q_gemm_scale_size = mGemmRunner->getWeightScaleSize(m, lda); + __nv_fp8_e4m3* q_gemm_weight_buffer + = reinterpret_cast<__nv_fp8_e4m3*>(nextWorkspacePtr(workspace_byte_ptr, offset, q_gemm_size)); + float* q_gemm_scale_buffer + = reinterpret_cast(nextWorkspacePtr(workspace_byte_ptr, offset, q_gemm_scale_size)); + mGemmRunner->fp8CS128x128(q_gemm_weight_buffer, q_gemm_scale_buffer, + reinterpret_cast<__nv_bfloat16 const*>(params.q_b_proj), lda, m, stream); + q_gemm_weight_ptr = reinterpret_cast<__nv_fp8_e4m3 const*>(q_gemm_weight_buffer); + q_gemm_scale_ptr = reinterpret_cast(q_gemm_scale_buffer); + } + mGemmRunner->fp8CS1x128(act_buffer, act_scale_buffer, + reinterpret_cast<__nv_bfloat16 const*>(params.fused_a_input), ldb, n, stream); + mGemmRunner->gemm(act_buffer, ldb, const_cast<__nv_fp8_e4m3*>(q_gemm_weight_ptr), lda, + reinterpret_cast<__nv_bfloat16*>(q_buffer), ldc, n, m, k, act_scale_buffer, + const_cast(q_gemm_scale_ptr), stream); + } + { + offset = quant_offset; + int m = mMLAParams.kv_lora_rank; + int n = params.acc_q_len; + int k = mMLAParams.qk_nope_head_dim; + int lda = k, ldb = mMLAParams.qk_nope_head_dim; + int ldc = params.head_num * (mMLAParams.kv_lora_rank + mMLAParams.qk_rope_head_dim); + size_t act_size = mGemmRunner->getFP8DataSize(n, params.head_num * mMLAParams.qk_nope_head_dim, true); + size_t act_scale_size = mGemmRunner->getActScaleSize(n, params.head_num * mMLAParams.qk_nope_head_dim); + __nv_fp8_e4m3* act_buffer + = reinterpret_cast<__nv_fp8_e4m3*>(nextWorkspacePtr(workspace_byte_ptr, offset, act_size)); + float* act_scale_buffer + = reinterpret_cast(nextWorkspacePtr(workspace_byte_ptr, offset, act_scale_size)); + __nv_fp8_e4m3 const* k_gemm_weight_ptr; + float const* k_gemm_scale_ptr; + if (mIsFP8Model) + { + k_gemm_weight_ptr = reinterpret_cast<__nv_fp8_e4m3 const*>(params.k_b_proj_trans); + k_gemm_scale_ptr = reinterpret_cast(params.k_b_trans_scale); + } + else + { + size_t k_gemm_size + = mGemmRunner->getFP8DataSize(params.head_num * m, mMLAParams.qk_nope_head_dim, false); + size_t k_gemm_scale_size + = mGemmRunner->getWeightScaleSize(params.head_num * m, mMLAParams.qk_nope_head_dim); + __nv_fp8_e4m3* k_gemm_weight_buffer + = reinterpret_cast<__nv_fp8_e4m3*>(nextWorkspacePtr(workspace_byte_ptr, offset, k_gemm_size)); + float* k_gemm_scale_buffer + = reinterpret_cast(nextWorkspacePtr(workspace_byte_ptr, offset, k_gemm_scale_size)); + mGemmRunner->fp8CS128x128(k_gemm_weight_buffer, k_gemm_scale_buffer, + reinterpret_cast<__nv_bfloat16 const*>(params.k_b_proj_trans), mMLAParams.qk_nope_head_dim, + params.head_num * m, stream); + k_gemm_weight_ptr = reinterpret_cast<__nv_fp8_e4m3 const*>(k_gemm_weight_buffer); + k_gemm_scale_ptr = reinterpret_cast(k_gemm_scale_buffer); + } + int shape_n_4_align = ((n - 1) / 4 + 1) * 4; + mGemmRunner->fp8CS1x128Reshape(act_buffer, act_scale_buffer, + reinterpret_cast<__nv_bfloat16 const*>(q_buffer), mMLAParams.qk_nope_head_dim, params.head_num, n, + mMLAParams.qk_nope_head_dim + mMLAParams.qk_rope_head_dim, stream); + + mGemmRunner->strideBatchGemm(reinterpret_cast<__nv_bfloat16*>(params.attention_input_buf), ldc, + mMLAParams.kv_lora_rank + mMLAParams.qk_rope_head_dim, act_buffer, ldb, n * k, + const_cast<__nv_fp8_e4m3*>(k_gemm_weight_ptr), lda, m * k, params.head_num, n, m, k, stream, + act_scale_buffer, shape_n_4_align * (mMLAParams.qk_nope_head_dim / 128), + const_cast(k_gemm_scale_ptr)); + } + } params.seqQOffset = cu_q_seqlens; params.cu_kv_seqlens = cu_kv_seqlens; params.fmha_tile_counter = fmha_tile_counter_ptr; + params.q_buf = q_buffer; invokeMLARopeGeneration(params, kv_cache_buffer, stream); @@ -980,22 +1234,68 @@ int GPTAttentionPluginCommon::mlaGeneration( // Run the fmha kernel mDecoderFMHARunner->run(fmhaParams); - + if (!mIsPTP128CEnabled) + { + { + auto transa = CUBLAS_OP_T; + auto transb = CUBLAS_OP_N; + int m = mMLAParams.v_head_dim; + int n = params.batch_size; + int k = mMLAParams.kv_lora_rank; + int lda = k, ldb = k; + int ldc = m; // params.head_size * params.c_k; + mCublasWrapper->createDescriptors(transa, transb, m, n, k, lda, ldb, ldc); + + mCublasWrapper->stridedBatchedGemm(transa, transb, m, n, k, + params.kv_b_proj + params.head_num * mMLAParams.kv_lora_rank * mMLAParams.qk_nope_head_dim, lda, + mMLAParams.v_head_dim * mMLAParams.kv_lora_rank, o_buffer, ldb * params.head_num, ldb, + params.context_buf, ldc * params.head_num, ldc, params.head_num, 1.0f, 0.0); + mCublasWrapper->destroyDescriptors(); + } + } + else { - auto transa = CUBLAS_OP_T; - auto transb = CUBLAS_OP_N; + // ??? 16 128 + offset = quant_offset; int m = mMLAParams.v_head_dim; int n = params.batch_size; int k = mMLAParams.kv_lora_rank; + int h = params.head_num; int lda = k, ldb = k; - int ldc = m; // params.head_size * params.c_k; - mCublasWrapper->createDescriptors(transa, transb, m, n, k, lda, ldb, ldc); - - mCublasWrapper->stridedBatchedGemm(transa, transb, m, n, k, - params.kv_b_proj + params.head_num * mMLAParams.kv_lora_rank * mMLAParams.qk_nope_head_dim, lda, - mMLAParams.v_head_dim * mMLAParams.kv_lora_rank, o_buffer, ldb * params.head_num, ldb, params.context_buf, - ldc * params.head_num, ldc, params.head_num, 1.0f, 0.0); - mCublasWrapper->destroyDescriptors(); + int ldc = h * m; + size_t act_size = mGemmRunner->getFP8DataSize(n, h * k, true); + size_t act_scale_size = mGemmRunner->getActScaleSize(n, h * k); + __nv_fp8_e4m3* act_buffer + = reinterpret_cast<__nv_fp8_e4m3*>(nextWorkspacePtr(workspace_byte_ptr, offset, act_size)); + float* act_scale_buffer + = reinterpret_cast(nextWorkspacePtr(workspace_byte_ptr, offset, act_scale_size)); + __nv_fp8_e4m3 const* v_gemm_weight_ptr; + float const* v_gemm_scale_ptr; + if (mIsFP8Model) + { + v_gemm_weight_ptr = reinterpret_cast<__nv_fp8_e4m3 const*>(params.kv_b_proj) + h * k * m; + v_gemm_scale_ptr = reinterpret_cast(params.kv_b_scale) + h * (k / 128) * (m / 128); + } + else + { + size_t v_gemm_size = mGemmRunner->getFP8DataSize(h * m, k, false); + size_t v_gemm_scale_size = mGemmRunner->getWeightScaleSize(h * m, k); + __nv_fp8_e4m3* v_gemm_weight_buffer + = reinterpret_cast<__nv_fp8_e4m3*>(nextWorkspacePtr(workspace_byte_ptr, offset, v_gemm_size)); + float* v_gemm_scale_buffer + = reinterpret_cast(nextWorkspacePtr(workspace_byte_ptr, offset, v_gemm_scale_size)); + mGemmRunner->fp8CS128x128(v_gemm_weight_buffer, v_gemm_scale_buffer, + reinterpret_cast<__nv_bfloat16 const*>(params.kv_b_proj) + h * m * k, k, h * m, stream); + + v_gemm_weight_ptr = reinterpret_cast<__nv_fp8_e4m3 const*>(v_gemm_weight_buffer); + v_gemm_scale_ptr = reinterpret_cast(v_gemm_scale_buffer); + } + int shape_n_4_align = ((n - 1) / 4 + 1) * 4; + mGemmRunner->fp8CS1x128Reshape( + act_buffer, act_scale_buffer, reinterpret_cast<__nv_bfloat16 const*>(o_buffer), k, h, n, k, stream); + mGemmRunner->strideBatchGemm(reinterpret_cast<__nv_bfloat16*>(params.context_buf), ldc, m, act_buffer, ldb, + n * k, const_cast<__nv_fp8_e4m3*>(v_gemm_weight_ptr), lda, m * k, params.head_num, n, m, k, stream, + act_scale_buffer, shape_n_4_align * (k / 128), const_cast(v_gemm_scale_ptr)); } sync_check_cuda_error(); @@ -2368,6 +2668,17 @@ int GPTAttentionPluginCommon::initialize() noexcept // Only deepseek must using fmha. TLLM_CHECK_WITH_INFO(mFMHARunner->isFmhaSupported() && mDecoderFMHARunner->isFmhaSupported(), "Deepseek should be supported by fmha in context and generation part."); + + if (mIsPTP128CEnabled && mIsFP8Model) + { + mGemmRunner = std::make_shared>(); + } + else if (mIsPTP128CEnabled && !mIsFP8Model) + { + mGemmRunner = std::make_shared>(); + } } // Fall back to unfused MHA kernels if not supported. @@ -2448,8 +2759,8 @@ size_t GPTAttentionPluginCommon::getCommonSerializationSize() const noexcept + sizeof(mPagedContextFMHA) + sizeof(mFP8ContextFMHA) + sizeof(mHasFullAttentionMask) + sizeof(mUseKVCache) + sizeof(mUnfuseQkvGemm) + sizeof(mUseLognScaling) + sizeof(mIsSpecDecodingEnabled) + sizeof(mUseSpecDecoding) + sizeof(mSpecDecodingIsGenerationLengthVariable) + sizeof(mSpecDecodingMaxGenerationLength) - + sizeof(mNbMultiBlockSemaphores) + sizeof(mIsMLAEnabled) + sizeof(mMLAParams) + sizeof(mSkipAttn) - + sizeof(uint32_t) // size of DecoderXQARunnerResource buffer. + + sizeof(mNbMultiBlockSemaphores) + sizeof(mIsMLAEnabled) + sizeof(mMLAParams) + sizeof(mIsPTP128CEnabled) + + sizeof(mIsFP8Model) + sizeof(mSkipAttn) + sizeof(uint32_t) // size of DecoderXQARunnerResource buffer. + sizeof(mCpSize) + sizeof(mCpRank) + sizeof(int32_t) * mCpGroup.size() + DecoderXQARunner::getResourceGlobal()->getSerializationSize(); } @@ -2507,6 +2818,8 @@ void GPTAttentionPluginCommon::serializeCommon(void* buffer) const noexcept write(d, mSpecDecodingMaxGenerationLength); write(d, mIsMLAEnabled); write(d, mMLAParams); + write(d, mIsPTP128CEnabled); + write(d, mIsFP8Model); write(d, mNbMultiBlockSemaphores); write(d, mSkipAttn); write(d, mCpSize); @@ -2617,6 +2930,8 @@ GPTAttentionPluginCreatorCommon::GPTAttentionPluginCreatorCommon() mPluginAttributes.emplace_back(PluginField("qk_nope_head_dim", nullptr, PluginFieldType::kINT32, 0)); mPluginAttributes.emplace_back(PluginField("qk_rope_head_dim", nullptr, PluginFieldType::kINT32, 0)); mPluginAttributes.emplace_back(PluginField("v_head_dim", nullptr, PluginFieldType::kINT32, 0)); + mPluginAttributes.emplace_back(PluginField("is_ptp128c_enabled", nullptr, PluginFieldType::kINT8, 0)); + mPluginAttributes.emplace_back(PluginField("is_fp8_model", nullptr, PluginFieldType::kINT8, 0)); mPluginAttributes.emplace_back(PluginField("skip_attn", nullptr, PluginFieldType::kINT8, 0)); mPluginAttributes.emplace_back(PluginField("cp_size", nullptr, PluginFieldType::kINT32, 0)); mPluginAttributes.emplace_back(PluginField("cp_rank", nullptr, PluginFieldType::kINT32, 0)); diff --git a/cpp/tensorrt_llm/plugins/gptAttentionCommon/gptAttentionCommon.h b/cpp/tensorrt_llm/plugins/gptAttentionCommon/gptAttentionCommon.h index 0ae6cba4b055..da1a05e2e89f 100644 --- a/cpp/tensorrt_llm/plugins/gptAttentionCommon/gptAttentionCommon.h +++ b/cpp/tensorrt_llm/plugins/gptAttentionCommon/gptAttentionCommon.h @@ -21,6 +21,7 @@ #include "tensorrt_llm/kernels/contextFusedMultiHeadAttention/fmhaRunner.h" #include "tensorrt_llm/kernels/contextFusedMultiHeadAttention/fused_multihead_attention_common.h" #include "tensorrt_llm/kernels/decoderMaskedMultiheadAttention/decoderXQARunner.h" +#include "tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_gemm.h" #include "tensorrt_llm/kernels/gptKernels.h" #include "tensorrt_llm/kernels/kvCacheUtils.h" #include "tensorrt_llm/kernels/mlaKernels.h" @@ -33,6 +34,8 @@ namespace tensorrt_llm::plugins { +using Fp8BlockScaleGemmRunnerPtr + = std::shared_ptr; class GPTAttentionPluginCommon : public BasePlugin { public: @@ -57,8 +60,8 @@ class GPTAttentionPluginCommon : public BasePlugin bool use_cache = true, bool is_spec_decoding_enabled = false, bool spec_decoding_is_generation_length_variable = false, int32_t spec_decoding_max_generation_length = 1, bool is_mla_enabled = false, int q_lora_rank = 0, int kv_lora_rank = 0, int qk_nope_head_dim = 0, - int qk_rope_head_dim = 0, int v_head_dim = 0, bool skip_attn = false, int cp_size = 1, int cp_rank = 0, - std::set cp_group = {}); + int qk_rope_head_dim = 0, int v_head_dim = 0, bool is_ptp128c_enabled = false, bool is_fp8_model = false, + bool skip_attn = false, int cp_size = 1, int cp_rank = 0, std::set cp_group = {}); GPTAttentionPluginCommon(void const* data, size_t length); @@ -415,6 +418,8 @@ class GPTAttentionPluginCommon : public BasePlugin int32_t mSpecDecodingMaxGenerationLength = 1; bool mIsMLAEnabled = false; tensorrt_llm::kernels::mlaMetaParams mMLAParams; + bool mIsPTP128CEnabled = false; + bool mIsFP8Model = false; int mCpSize = 1; int mCpRank = 0; std::set mCpGroup = {}; @@ -438,6 +443,7 @@ class GPTAttentionPluginCommon : public BasePlugin UniqPtrWNullCopy mFMHARunner; UniqPtrWNullCopy mDecoderFMHARunner; UniqPtrWNullCopy mDecoderXQARunner; + Fp8BlockScaleGemmRunnerPtr mGemmRunner; bool mMultiBlockMode; bool mEnableXQA; diff --git a/cpp/tensorrt_llm/plugins/gptAttentionPlugin/gptAttentionPlugin.cpp b/cpp/tensorrt_llm/plugins/gptAttentionPlugin/gptAttentionPlugin.cpp index 3f94474a7883..5bcdecf76da7 100644 --- a/cpp/tensorrt_llm/plugins/gptAttentionPlugin/gptAttentionPlugin.cpp +++ b/cpp/tensorrt_llm/plugins/gptAttentionPlugin/gptAttentionPlugin.cpp @@ -19,6 +19,7 @@ #include "tensorrt_llm/batch_manager/contextProgress.h" #include "tensorrt_llm/common/logger.h" +#include "tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_gemm.h" #include "tensorrt_llm/kernels/decoderMaskedMultiheadAttention.h" #include "tensorrt_llm/kernels/gptKernels.h" #include "tensorrt_llm/kernels/unfusedAttentionKernels.h" @@ -63,8 +64,8 @@ GPTAttentionPlugin::GPTAttentionPlugin(int layer_idx, int num_heads, int vision_ bool use_paged_context_fmha, bool use_fp8_context_fmha, bool has_full_attention_mask, bool use_cache, bool is_spec_decoding_enabled, bool spec_decoding_is_generation_length_variable, int spec_decoding_max_generation_length, bool is_mla_enabled, int q_lora_rank, int kv_lora_rank, - int qk_nope_head_dim, int qk_rope_head_dim, int v_head_dim, bool skip_attn, int cp_size, int cp_rank, - std::set cp_group) + int qk_nope_head_dim, int qk_rope_head_dim, int v_head_dim, bool is_ptp128c_enabled, bool is_fp8_model, + bool skip_attn, int cp_size, int cp_rank, std::set cp_group) : GPTAttentionPluginCommon(layer_idx, num_heads, vision_start, vision_length, num_kv_heads, layer_idx_in_cache_pool, head_size, unidirectional, q_scaling, attn_logit_softcapping_scale, position_embedding_type, rotary_embedding_dim, rotary_embedding_base, rotary_embedding_scale_type, rotary_embedding_scale, @@ -74,9 +75,8 @@ GPTAttentionPlugin::GPTAttentionPlugin(int layer_idx, int num_heads, int vision_ type, max_context_length, qkv_bias_enabled, cross_attention, max_distance, pos_shift_enabled, dense_context_fmha, use_paged_context_fmha, use_fp8_context_fmha, has_full_attention_mask, use_cache, is_spec_decoding_enabled, spec_decoding_is_generation_length_variable, spec_decoding_max_generation_length, - is_mla_enabled, q_lora_rank, kv_lora_rank, qk_nope_head_dim, qk_rope_head_dim, v_head_dim, skip_attn, cp_size, - - cp_rank, cp_group) + is_mla_enabled, q_lora_rank, kv_lora_rank, qk_nope_head_dim, qk_rope_head_dim, v_head_dim, is_ptp128c_enabled, + is_fp8_model, skip_attn, cp_size, cp_rank, cp_group) { initEntryIdx(); } @@ -128,6 +128,12 @@ std::string GPTAttentionPlugin::toString(IdxEntry const& entry) const case IdxEntry::LONG_ROPE_ROTARY_COS_SIN: return "LONG_ROPE_ROTARY_COS_SIN"; case IdxEntry::HOST_RUNTIME_PERF_KNOBS: return "HOST_RUNTIME_PERF_KNOBS"; case IdxEntry::HOST_CONTEXT_PROGRESS: return "HOST_CONTEXT_PROGRESS"; + case IdxEntry::MLA_Q_B_PROJ_TENSOR: return "MLA_Q_B_PROJ_TENSOR"; + case IdxEntry::MLA_KV_B_PROJ_TENSOR: return "MLA_KV_B_PROJ_TENSOR"; + case IdxEntry::MLA_K_B_PROJ_TRANS_TENSOR: return "MLA_K_B_PROJ_TRANS_TENSOR"; + case IdxEntry::MLA_Q_B_SCALE_TENSOR: return "MLA_Q_B_SCALE_TENSOR"; + case IdxEntry::MLA_KV_B_SCALE_TENSOR: return "MLA_KV_B_SCALE_TENSOR"; + case IdxEntry::MLA_K_B_TRANS_SCALE_TENSOR: return "MLA_K_B_TRANS_SCALE_TENSOR"; case IdxEntry::SKIP_ATTN: return "SKIP_ATTN"; case IdxEntry::ENUM_SIZE: return "ENUM_SIZE"; } @@ -179,9 +185,12 @@ bool GPTAttentionPlugin::isEntryUsed(IdxEntry const& entry) const case IdxEntry::MROPE_POSITION_DELTAS: return isMRoPE(); case IdxEntry::HOST_RUNTIME_PERF_KNOBS: return true; case IdxEntry::HOST_CONTEXT_PROGRESS: return true; - case IdxEntry::MLA_FUSED_Q_PROJ_TENSOR: return mIsMLAEnabled; case IdxEntry::MLA_Q_B_PROJ_TENSOR: return mIsMLAEnabled; case IdxEntry::MLA_KV_B_PROJ_TENSOR: return mIsMLAEnabled; + case IdxEntry::MLA_K_B_PROJ_TRANS_TENSOR: return mIsMLAEnabled; + case IdxEntry::MLA_Q_B_SCALE_TENSOR: return mIsMLAEnabled && mIsPTP128CEnabled && mIsFP8Model; + case IdxEntry::MLA_KV_B_SCALE_TENSOR: return mIsMLAEnabled && mIsPTP128CEnabled && mIsFP8Model; + case IdxEntry::MLA_K_B_TRANS_SCALE_TENSOR: return mIsMLAEnabled && mIsPTP128CEnabled && mIsFP8Model; case IdxEntry::SKIP_ATTN: return mSkipAttn; default: return false; } @@ -323,6 +332,20 @@ bool GPTAttentionPlugin::supportsFormatCombination( posCaseLine = __LINE__; result = inOut[pos].type == nvinfer1::DataType::kFLOAT; } + else if (mIsMLAEnabled && mIsPTP128CEnabled && mIsFP8Model + && (pos == getIdx(IdxEntry::MLA_Q_B_SCALE_TENSOR) || pos == getIdx(IdxEntry::MLA_KV_B_SCALE_TENSOR) + || pos == getIdx(IdxEntry::MLA_K_B_TRANS_SCALE_TENSOR))) + { + posCaseLine = __LINE__; + result = inOut[pos].type == nvinfer1::DataType::kFLOAT && inOut[pos].format == TensorFormat::kLINEAR; + } + else if (mIsMLAEnabled && mIsPTP128CEnabled && mIsFP8Model + && (pos == getIdx(IdxEntry::MLA_Q_B_PROJ_TENSOR) || pos == getIdx(IdxEntry::MLA_KV_B_PROJ_TENSOR) + || pos == getIdx(IdxEntry::MLA_K_B_PROJ_TRANS_TENSOR))) + { + posCaseLine = __LINE__; + result = inOut[pos].type == nvinfer1::DataType::kFP8 && inOut[pos].format == TensorFormat::kLINEAR; + } else if (isLongRoPE() && (pos == getIdx(IdxEntry::LONG_ROPE_ROTARY_INV_FREQ) || pos == getIdx(IdxEntry::LONG_ROPE_ROTARY_COS_SIN))) { @@ -555,6 +578,7 @@ size_t GPTAttentionPlugin::getWorkspaceSize(nvinfer1::PluginTensorDesc const* in = getWorkspaceSizeForGeneration(type, max_num_seq, max_kv_cache_length, max_num_tokens); size_t attention_input_workspace_size = 0; + size_t context_mla_fp8_quant_size = 0; if (mIsMLAEnabled) { int32_t const size_per_head @@ -565,6 +589,25 @@ size_t GPTAttentionPlugin::getWorkspaceSize(nvinfer1::PluginTensorDesc const* in size_t workspaces[1]; workspaces[0] = attention_input_size; attention_input_workspace_size = tensorrt_llm::common::calculateTotalWorkspaceSize(workspaces, 1); + + if (mIsPTP128CEnabled) + { + size_t act_quant_size = mGemmRunner->getActWorkspaceSize( + max_num_tokens, mMLAParams.qk_nope_head_dim + mMLAParams.kv_lora_rank + mMLAParams.qk_rope_head_dim); + size_t weight_size = 0; + if (!mIsFP8Model) + { + weight_size = std::max(mGemmRunner->getWeightWorkspaceSize( + mNumHeads * (mMLAParams.qk_nope_head_dim + mMLAParams.qk_rope_head_dim), + mMLAParams.q_lora_rank), + mGemmRunner->getWeightWorkspaceSize( + 2 * mNumHeads * mMLAParams.qk_nope_head_dim, mMLAParams.kv_lora_rank)); + } + size_t workspaces[2]; + workspaces[0] = act_quant_size; + workspaces[1] = weight_size; + context_mla_fp8_quant_size = tensorrt_llm::common::calculateTotalWorkspaceSize(workspaces, 2); + } } else if (mUnfuseQkvGemm) { @@ -579,7 +622,8 @@ size_t GPTAttentionPlugin::getWorkspaceSize(nvinfer1::PluginTensorDesc const* in attention_input_workspace_size = tensorrt_llm::common::calculateTotalWorkspaceSize(workspaces, 1); } - return std::max(context_workspace_size, generation_workspace_size) + attention_input_workspace_size; + return std::max(context_mla_fp8_quant_size, std::max(context_workspace_size, generation_workspace_size)) + + attention_input_workspace_size; } static size_t getStride(nvinfer1::Dims const& dims, int n) @@ -668,9 +712,10 @@ mlaParams GPTAttentionPlugin::enqueueMLAPreprocess(int32_t localNbSeq, int32_ { auto const* input = static_cast(inputs[getIdx(IdxEntry::QKV_TENSOR)]); - auto const* fused_q_proj = static_cast(inputs[getIdx(IdxEntry::MLA_FUSED_Q_PROJ_TENSOR)]); auto const* q_b_proj = static_cast(inputs[getIdx(IdxEntry::MLA_Q_B_PROJ_TENSOR)]); auto const* kv_b_proj = static_cast(inputs[getIdx(IdxEntry::MLA_KV_B_PROJ_TENSOR)]); + auto const* k_b_proj_trans = static_cast(inputs[getIdx(IdxEntry::MLA_K_B_PROJ_TRANS_TENSOR)]); + float2 const* cos_sin_cache = static_cast(inputs[getIdx(IdxEntry::ROTARY_COS_SIN)]); AttentionOutT* context_buf_ = static_cast(outputs[0]); @@ -678,14 +723,20 @@ mlaParams GPTAttentionPlugin::enqueueMLAPreprocess(int32_t localNbSeq, int32_ mlaParams mla_params; mla_params.fused_a_input = input; mla_params.context_buf = reinterpret_cast(context_buf_); - mla_params.fused_q_proj = fused_q_proj; mla_params.q_b_proj = q_b_proj; mla_params.kv_b_proj = kv_b_proj; + mla_params.k_b_proj_trans = k_b_proj_trans; mla_params.cos_sin_cache = cos_sin_cache; mla_params.batch_size = localNbSeq; mla_params.acc_q_len = localNbTokens; mla_params.head_num = mNumHeads; mla_params.meta = mMLAParams; + if (mIsPTP128CEnabled && mIsFP8Model) + { + mla_params.q_b_scale = static_cast(inputs[getIdx(IdxEntry::MLA_Q_B_SCALE_TENSOR)]); + mla_params.kv_b_scale = static_cast(inputs[getIdx(IdxEntry::MLA_KV_B_SCALE_TENSOR)]); + mla_params.k_b_trans_scale = static_cast(inputs[getIdx(IdxEntry::MLA_K_B_TRANS_SCALE_TENSOR)]); + } // { // __nv_bfloat16 *h_input; @@ -1363,6 +1414,8 @@ IPluginV2* GPTAttentionPluginCreator::createPlugin(char const* name, PluginField static_cast(p.getScalar("qk_nope_head_dim").value()), static_cast(p.getScalar("qk_rope_head_dim").value()), static_cast(p.getScalar("v_head_dim").value()), + static_cast(p.getScalar("is_ptp128c_enabled").value()), + static_cast(p.getScalar("is_fp8_model").value()), static_cast(p.getScalar("skip_attn").value()), static_cast(p.getScalar("cp_size").value()), static_cast(p.getScalar("cp_rank").value()), diff --git a/cpp/tensorrt_llm/plugins/gptAttentionPlugin/gptAttentionPlugin.h b/cpp/tensorrt_llm/plugins/gptAttentionPlugin/gptAttentionPlugin.h index ebbeb04fb75a..489d5bb448ec 100644 --- a/cpp/tensorrt_llm/plugins/gptAttentionPlugin/gptAttentionPlugin.h +++ b/cpp/tensorrt_llm/plugins/gptAttentionPlugin/gptAttentionPlugin.h @@ -119,8 +119,8 @@ class GPTAttentionPlugin : public GPTAttentionPluginCommon bool use_cache = true, bool is_spec_decoding_enabled = false, bool spec_decoding_is_generation_length_variable = false, int spec_decoding_max_generation_length = 1, bool is_mla_enabled = false, int q_lora_rank = 0, int kv_lora_rank = 0, int qk_nope_head_dim = 0, - int qk_rope_head_dim = 0, int v_head_dim = 0, bool skip_attn = false, int cp_size = 1, int cp_rank = 0, - std::set cp_group = {}); + int qk_rope_head_dim = 0, int v_head_dim = 0, bool is_ptp128c_enabled = false, bool is_fp8_model = false, + bool skip_attn = false, int cp_size = 1, int cp_rank = 0, std::set cp_group = {}); GPTAttentionPlugin(void const* data, size_t length); @@ -226,9 +226,12 @@ class GPTAttentionPlugin : public GPTAttentionPluginCommon MROPE_POSITION_DELTAS, HOST_RUNTIME_PERF_KNOBS, HOST_CONTEXT_PROGRESS, - MLA_FUSED_Q_PROJ_TENSOR, MLA_Q_B_PROJ_TENSOR, MLA_KV_B_PROJ_TENSOR, + MLA_K_B_PROJ_TRANS_TENSOR, + MLA_Q_B_SCALE_TENSOR, + MLA_KV_B_SCALE_TENSOR, + MLA_K_B_TRANS_SCALE_TENSOR, SKIP_ATTN, LOGN_SCALING, ENUM_SIZE, // Used to count the number of IdxEntry, must put in last diff --git a/cpp/tensorrt_llm/plugins/mixtureOfExperts/mixtureOfExpertsPlugin.cpp b/cpp/tensorrt_llm/plugins/mixtureOfExperts/mixtureOfExpertsPlugin.cpp index 61cec4842616..c4ab883c369c 100644 --- a/cpp/tensorrt_llm/plugins/mixtureOfExperts/mixtureOfExpertsPlugin.cpp +++ b/cpp/tensorrt_llm/plugins/mixtureOfExperts/mixtureOfExpertsPlugin.cpp @@ -45,7 +45,8 @@ MixtureOfExpertsPlugin::MixtureOfExpertsPlugin(bool remove_input_padding, int nu bool use_finished, bool use_bias, int tp_size, int tp_rank, int ep_size, int ep_rank, MOEExpertScaleNormalizationMode normalization_mode, float sparse_mixer_epsilon, bool force_determinism, int side_stream_id, MixtureOfExpertsPluginProfilerPtr gemm_profiler_ptr, bool use_lora, - nvinfer1::DataType lora_type, LoraPluginProfilerPtr lora_profiler, int max_low_rank) + nvinfer1::DataType lora_type, LoraPluginProfilerPtr lora_profiler, int max_low_rank, bool use_deepseek = false, + bool use_deepseek_with_native_fp8_weights = false) : mRemoveInputPadding(remove_input_padding) , mNumExperts(number_of_experts) , mK(top_k) @@ -68,6 +69,8 @@ MixtureOfExpertsPlugin::MixtureOfExpertsPlugin(bool remove_input_padding, int nu , mLoraType(lora_type) , mLoraProfiler(std::move(lora_profiler)) , mMaxLowRank(max_low_rank) + , mUseDeepSeek(use_deepseek) + , mUseDeepSeekWithNativeFp8Weights(use_deepseek_with_native_fp8_weights) { init(); } @@ -105,6 +108,8 @@ tensorrt_llm::plugins::MixtureOfExpertsPlugin::MixtureOfExpertsPlugin(MixtureOfE , mLoraImpl2(other.mLoraImpl2) , mLayerName(other.mLayerName) , mNamespace(other.mNamespace) + , mUseDeepSeek(other.mUseDeepSeek) + , mUseDeepSeekWithNativeFp8Weights(other.mUseDeepSeekWithNativeFp8Weights) { init(); } @@ -116,7 +121,8 @@ size_t MixtureOfExpertsPlugin::getSerializationSize() const noexcept + sizeof(QuantMode::BaseType) + sizeof(mUseFinished) + sizeof(mUseBias) + sizeof(mParallelismConfig) + sizeof(mNormalizationMode) + sizeof(mSparseMixerEpsilon) + sizeof(mDims) + sizeof(mUseDeterministicKernels) + sizeof(mSideStreamId) + mGemmProfiler->getSerializationSize(mGemmId1) - + mGemmProfiler->getSerializationSize(mGemmId2) + sizeof(mUseLora) + sizeof(mLoraType) + sizeof(mMaxLowRank); + + mGemmProfiler->getSerializationSize(mGemmId2) + sizeof(mUseLora) + sizeof(mLoraType) + sizeof(mMaxLowRank) + + sizeof(mUseDeepSeek) + sizeof(mUseDeepSeekWithNativeFp8Weights); if (hasLora()) { @@ -157,6 +163,8 @@ MixtureOfExpertsPlugin::MixtureOfExpertsPlugin(void const* data, size_t length, read(d, mUseLora); read(d, mLoraType); read(d, mMaxLowRank); + read(d, mUseDeepSeek); + read(d, mUseDeepSeekWithNativeFp8Weights); // Call init before deserialising the profiler to initialize mGemmId init(); @@ -202,6 +210,8 @@ void MixtureOfExpertsPlugin::serialize(void* buffer) const noexcept write(d, mUseLora); write(d, mLoraType); write(d, mMaxLowRank); + write(d, mUseDeepSeek); + write(d, mUseDeepSeekWithNativeFp8Weights); mGemmProfiler->serialize(d, mGemmId1); mGemmProfiler->serialize(d, mGemmId2); @@ -250,6 +260,10 @@ void MixtureOfExpertsPlugin::init() { mMOERunner = std::make_unique>(); } + else if (mType == DataType::kBF16 && mWeightType == DataType::kFP8) + { + mMOERunner = std::make_unique>(); + } else if (mType == DataType::kBF16 && mWeightType == DataType::kINT8) { mMOERunner = std::make_unique>(); @@ -314,6 +328,35 @@ void MixtureOfExpertsPlugin::init() TLLM_CUDA_CHECK(cudaEventCreate(&mMemcpyEvent)); } + + if (useDeepSeekWithNativeFp8Weights()) + { + TLLM_CHECK_WITH_INFO( + useDeepSeek(), "the mUseDeepSeek should be true when the mUseDeepSeekWithNativeFp8Weights is true"); + } + if (useDeepSeek()) + { + +#ifdef ENABLE_BF16 + if (mType == DataType::kBF16) + { + if (mWeightType == DataType::kBF16) + { + mBlockScaleGemmImplPtr + = std::make_shared>(); + } +#ifdef ENABLE_FP8 + else if (mWeightType == DataType::kFP8) + { + mBlockScaleGemmImplPtr + = std::make_shared>(); + } + } +#endif +#endif + } mSideStreamPtr = nullptr; mDebugStallMain = tensorrt_llm::runtime::utils::stallStream("TLLM_DEBUG_MOE_STALL_MAIN"); mDebugStallSide = tensorrt_llm::runtime::utils::stallStream("TLLM_DEBUG_MOE_STALL_SIDE"); @@ -380,6 +423,11 @@ bool MixtureOfExpertsPlugin::supportsFormatCombination( { return inOut[pos].type == DataType::kFLOAT; } + else if (useDeepSeekWithNativeFp8Weights() && getExpertDeepseekScale1Index() <= pos + && pos <= getExpertDeepseekScale2Index()) + { + return inOut[pos].type == DataType::kFLOAT; + } else if (hasExpertIntQuantScales() && getExpertIntQuantScale1Index() <= pos && pos <= getExpertIntQuantScale2Index()) { @@ -470,7 +518,7 @@ auto MixtureOfExpertsPlugin::setupWorkspace(void* base_ptr, int64_t num_tokens, size_t dtype_size = tensorrt_llm::common::getDTypeSize(mType); size_t moe_workspace_size = mMOERunner->getWorkspaceSize(num_tokens, mExpertHiddenSize, mExpertInterSize, - mNumExperts, mK, mActivationType, mNormalizationMode, mParallelismConfig, hasLora()); + mNumExperts, mK, mActivationType, mNormalizationMode, mParallelismConfig, hasLora(), useDeepSeek()); // Output of post-softmax routing probabilities size_t scale_probabilities_size = num_tokens * mNumExperts * sizeof(float); @@ -489,12 +537,24 @@ auto MixtureOfExpertsPlugin::setupWorkspace(void* base_ptr, int64_t num_tokens, mLoraImpl2->getWorkspaceSize(num_tokens * mK, num_reqs_lora, mLoraType)); } + size_t deepseek_workspace_size = 0; + if (useDeepSeek()) + { + bool is_gated_actiation = isGatedActivation(mActivationType); + int factor = is_gated_actiation ? 2 : 1; + size_t deepseek_fc1_size = mBlockScaleGemmImplPtr->getWorkspaceSize( + num_tokens * mK, factor * mExpertInterSize, mExpertHiddenSize, mNumExperts); + size_t deepseek_fc2_size = mBlockScaleGemmImplPtr->getWorkspaceSize( + num_tokens * mK, mExpertHiddenSize, mExpertInterSize, mNumExperts); + deepseek_workspace_size = std::max(deepseek_fc1_size, deepseek_fc2_size); + } std::vector workspaces{ moe_workspace_size, scale_probabilities_size, src_to_dest_map_size, selected_expert_size, lora_workspace_size, + deepseek_workspace_size, }; WorkspaceInfo info{}; @@ -507,6 +567,7 @@ auto MixtureOfExpertsPlugin::setupWorkspace(void* base_ptr, int64_t num_tokens, info.src_to_dest_map = nextWorkspacePtr((int8_t*) info.scale_probs, scale_probabilities_size); info.selected_experts = nextWorkspacePtr((int8_t*) info.src_to_dest_map, src_to_dest_map_size); info.lora_workspace = nextWorkspacePtr((int8_t*) info.selected_experts, selected_expert_size); + info.deepseek_workspace = nextWorkspacePtr((int8_t*) info.lora_workspace, lora_workspace_size); } return info; @@ -669,6 +730,17 @@ LoraParams MixtureOfExpertsPlugin::getLoraParams( mLoraExpandGatedWeightPtrs.data()); } +BlockScaleParams MixtureOfExpertsPlugin::getBlockScaleParams( + nvinfer1::PluginTensorDesc const* inputDesc, void const* const* inputs, void* workspace) +{ + TLLM_CHECK(useDeepSeek()); + auto fc1_scales_ptr = static_cast(inputs[getExpertDeepseekScale1Index()]); + auto fc2_scales_ptr = static_cast(inputs[getExpertDeepseekScale2Index()]); + + return BlockScaleParams( + fc1_scales_ptr, fc2_scales_ptr, mBlockScaleGemmImplPtr, static_cast(workspace), &mMemcpyEvent); +} + int MixtureOfExpertsPlugin::enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace_ptr, cudaStream_t stream) noexcept @@ -767,6 +839,20 @@ int MixtureOfExpertsPlugin::enqueue(nvinfer1::PluginTensorDesc const* inputDesc, auto gemm1 = mGemmProfiler->getBestConfig(num_tokens, mGemmId1); auto gemm2 = mGemmProfiler->getBestConfig(num_tokens, mGemmId2); + + BlockScaleParams deepseek_params{}; + if (useDeepSeek()) + { + deepseek_params = getBlockScaleParams(inputDesc, inputs, workspace.deepseek_workspace); + auto config + = cutlass_extensions::CutlassGemmConfig(cutlass_extensions::CutlassTileConfigSM90::CtaShape128x16x128B, + cutlass_extensions::MainloopScheduleType::AUTO, cutlass_extensions::EpilogueScheduleType::AUTO, + cutlass_extensions::ClusterShape::ClusterShape_1x1x1); + + gemm1 = std::make_optional(config); + gemm2 = std::make_optional(config); + } + mMOERunner->setTactic(gemm1, gemm2); mMOERunner->runMoe(inputs[getInputTensorIndex()], static_cast(inputs[getRoutingTensorIndex()]), inputs[getExpertWeights1Index()], hasBias() ? inputs[getExpertBias1Index()] : nullptr, mActivationType, @@ -777,7 +863,7 @@ int MixtureOfExpertsPlugin::enqueue(nvinfer1::PluginTensorDesc const* inputDesc, hasFinishedTensor() ? static_cast(inputs[getFinishedTensorIndex()]) : nullptr, num_not_finished, workspace.scale_probs, static_cast(workspace.src_to_dest_map), static_cast(workspace.selected_experts), mSparseMixerEpsilon, mParallelismConfig, mNormalizationMode, - hasLora(), lora_params, stream); + hasLora(), lora_params, useDeepSeek(), deepseek_params, stream); if (useSideStream()) { @@ -815,6 +901,11 @@ char const* MixtureOfExpertsPlugin::getPluginVersion() const noexcept int MixtureOfExpertsPlugin::initialize() noexcept { + if (useDeepSeek()) + { + mLoraProfiler->setSkip(true); + mGemmProfiler->setSkip(true); + } mGemmProfiler->setGemmToProfile(kernels::GemmProfilerBackend::GemmToProfile::GEMM_1); mGemmProfiler->profileTactics(this, mType, mDims, mGemmId1); mGemmProfiler->setGemmToProfile(kernels::GemmProfilerBackend::GemmToProfile::GEMM_2); @@ -910,6 +1001,9 @@ MixtureOfExpertsPluginCreator::MixtureOfExpertsPluginCreator() mPluginAttributes.emplace_back(nvinfer1::PluginField("use_lora", nullptr, PluginFieldType::kINT32, 0)); mPluginAttributes.emplace_back(nvinfer1::PluginField("lora_type_id", nullptr, PluginFieldType::kINT32, 0)); mPluginAttributes.emplace_back(nvinfer1::PluginField("max_low_rank", nullptr, PluginFieldType::kINT32, 0)); + mPluginAttributes.emplace_back(nvinfer1::PluginField("use_deepseek", nullptr, PluginFieldType::kINT32, 0)); + mPluginAttributes.emplace_back( + nvinfer1::PluginField("use_deepseek_with_native_fp8_weights", nullptr, PluginFieldType::kINT32, 0)); mFC.nbFields = mPluginAttributes.size(); mFC.fields = mPluginAttributes.data(); } @@ -940,6 +1034,8 @@ IPluginV2* MixtureOfExpertsPluginCreator::createPlugin( int mUseLora{}; int mLoraType{INT_MAX}; int mMaxLowRank{0}; + int mUseDeepSeek{0}; + int mUseDeepSeekWithNativeFp8Weights{0}; float mSparseMixerEpsilon = -INFINITY; @@ -977,6 +1073,8 @@ IPluginV2* MixtureOfExpertsPluginCreator::createPlugin( MapPair{"side_stream_id", std::ref(mSideStreamId), true}, MapPair{"lora_type_id", std::ref(mLoraType), true}, MapPair{"max_low_rank", std::ref(mMaxLowRank), true}, + MapPair{"use_deepseek", std::ref(mUseDeepSeek), true}, + MapPair{"use_deepseek_with_native_fp8_weights", std::ref(mUseDeepSeekWithNativeFp8Weights), true}, }; for (int i = 0; i < fc->nbFields; ++i) { @@ -1036,7 +1134,8 @@ IPluginV2* MixtureOfExpertsPluginCreator::createPlugin( QuantMode(mQuantMode), mUseFinished != 0, mUseBias != 0, mTPSize, mTPRank, mEPSize, mEPRank, static_cast(mNormalizationMode), mSparseMixerEpsilon, mRequiresDeterminism != 0, mSideStreamId, gemmProfiler, mUseLora != 0, - static_cast(mLoraType), loraProfiler, mMaxLowRank); + static_cast(mLoraType), loraProfiler, mMaxLowRank, mUseDeepSeek, + mUseDeepSeekWithNativeFp8Weights); obj->setPluginNamespace(mNamespace.c_str()); return obj; } diff --git a/cpp/tensorrt_llm/plugins/mixtureOfExperts/mixtureOfExpertsPlugin.h b/cpp/tensorrt_llm/plugins/mixtureOfExperts/mixtureOfExpertsPlugin.h index 6a3608460308..7959e9077117 100644 --- a/cpp/tensorrt_llm/plugins/mixtureOfExperts/mixtureOfExpertsPlugin.h +++ b/cpp/tensorrt_llm/plugins/mixtureOfExperts/mixtureOfExpertsPlugin.h @@ -20,6 +20,7 @@ #include "NvInferPlugin.h" #include "tensorrt_llm/common/cudaUtils.h" #include "tensorrt_llm/common/quantization.h" +#include "tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_gemm.h" #include "tensorrt_llm/kernels/lora/lora.h" #include "tensorrt_llm/kernels/mixtureOfExperts/moe_kernels.h" #include "tensorrt_llm/plugins/common/gemmPluginProfiler.h" @@ -100,6 +101,7 @@ class MixtureOfExpertsPlugin : public nvinfer1::IPluginV2DynamicExt using MOEExpertScaleNormalizationMode = tensorrt_llm::kernels::MOEExpertScaleNormalizationMode; using LoraPluginProfilerPtr = std::shared_ptr; using LoraImplPtr = std::shared_ptr; + using BlockScaleGemmImplPtr = std::shared_ptr; MixtureOfExpertsPlugin() = delete; MixtureOfExpertsPlugin(bool remove_input_padding, int number_of_experts, int top_k, int expert_hidden_size, @@ -108,7 +110,8 @@ class MixtureOfExpertsPlugin : public nvinfer1::IPluginV2DynamicExt bool use_finished, bool use_bias, int tp_size, int tp_rank, int ep_size, int ep_rank, MOEExpertScaleNormalizationMode normalization_mode, float sparse_mixer_epsilon, bool force_determinism, int side_stream_id, MixtureOfExpertsPluginProfilerPtr gemm_profiler_ptr, bool use_lora, - nvinfer1::DataType lora_type, LoraPluginProfilerPtr lora_profiler, int max_low_rank); + nvinfer1::DataType lora_type, LoraPluginProfilerPtr lora_profiler, int max_low_rank, bool use_deepseek, + bool use_deepseek_with_native_fp8_weights); MixtureOfExpertsPlugin(void const* data, size_t length, MixtureOfExpertsPluginProfilerPtr gemm_profiler_ptr, LoraPluginProfilerPtr lora_profiler); MixtureOfExpertsPlugin(MixtureOfExpertsPlugin const&); @@ -181,6 +184,11 @@ class MixtureOfExpertsPlugin : public nvinfer1::IPluginV2DynamicExt MixtureOfExpertsPluginProfilerPtr mGemmProfiler; + // deepseek related + bool mUseDeepSeek{}; + bool mUseDeepSeekWithNativeFp8Weights{}; + BlockScaleGemmImplPtr mBlockScaleGemmImplPtr; + // lora related bool mUseLora{}; nvinfer1::DataType mLoraType{}; @@ -216,6 +224,7 @@ class MixtureOfExpertsPlugin : public nvinfer1::IPluginV2DynamicExt void* src_to_dest_map{}; void* selected_experts{}; void* lora_workspace{}; + void* deepseek_workspace{}; size_t size{}; }; @@ -230,6 +239,9 @@ class MixtureOfExpertsPlugin : public nvinfer1::IPluginV2DynamicExt kernels::LoraParams getLoraParams( nvinfer1::PluginTensorDesc const* inputDesc, void const* const* inputs, void* workspace); + kernels::BlockScaleParams getBlockScaleParams( + nvinfer1::PluginTensorDesc const* inputDesc, void const* const* inputs, void* workspace); + enum class RequestType : int32_t { kCONTEXT = 0, @@ -295,6 +307,16 @@ class MixtureOfExpertsPlugin : public nvinfer1::IPluginV2DynamicExt return mUseLora; } + bool useDeepSeek() const + { + return mUseDeepSeek; + } + + bool useDeepSeekWithNativeFp8Weights() const + { + return mUseDeepSeekWithNativeFp8Weights; + } + bool hasGatedLoraWeightsAndRanks() const { return mUseLora && isGatedActivation(mActivationType); @@ -315,6 +337,16 @@ class MixtureOfExpertsPlugin : public nvinfer1::IPluginV2DynamicExt return getExpertBias2Index() + hasFinishedTensor(); } + IndexType getExpertDeepseekScale1Index() const + { + return getFinishedTensorIndex() + useDeepSeekWithNativeFp8Weights(); + } + + IndexType getExpertDeepseekScale2Index() const + { + return getExpertDeepseekScale1Index() + useDeepSeekWithNativeFp8Weights(); + } + IndexType getExpertIntQuantScale1Index() const { return getFinishedTensorIndex() + hasExpertIntQuantScales(); @@ -327,7 +359,7 @@ class MixtureOfExpertsPlugin : public nvinfer1::IPluginV2DynamicExt IndexType getExpertFP8Dequant1Index() const { - return getExpertIntQuantScale2Index() + hasExpertFp8QuantScales(); + return std::max(getExpertIntQuantScale2Index(), getExpertDeepseekScale2Index()) + hasExpertFp8QuantScales(); } IndexType getExpertFP8Quant2Index() const diff --git a/cpp/tensorrt_llm/thop/moeOp.cpp b/cpp/tensorrt_llm/thop/moeOp.cpp index 0eab70cac409..06fe4fe0045f 100644 --- a/cpp/tensorrt_llm/thop/moeOp.cpp +++ b/cpp/tensorrt_llm/thop/moeOp.cpp @@ -294,6 +294,7 @@ class FusedMoeRunner : public torch::CustomClassHolder kernels::QuantParams quant_params{}; kernels::LoraParams lora_params{}; + kernels::BlockScaleParams deepseek_params{}; mKernelRunner->runMoe(input.const_data_ptr(), gating_output.const_data_ptr(), fc1_expert_weights.const_data_ptr(), nullptr, activation_type, fc2_expert_weights.const_data_ptr(), nullptr, @@ -301,7 +302,7 @@ class FusedMoeRunner : public torch::CustomClassHolder static_cast(workspace_info.workspace), output.data_ptr(), nullptr, output.sizes()[0], workspace_info.scale_probs, static_cast(workspace_info.src_to_dest_map), static_cast(workspace_info.selected_experts), 0, parallelism_config, norm_mode, false, lora_params, - stream); + false, deepseek_params, stream); return output; } @@ -451,7 +452,7 @@ class FusedMoeRunner : public torch::CustomClassHolder kernels::MOEExpertScaleNormalizationMode norm_mode, kernels::MOEParallelismConfig const& parallelismConfig) { size_t moe_workspace_size = mKernelRunner->getWorkspaceSize(num_rows, hidden_size, inter_size, num_experts, - top_k, activation_type, norm_mode, parallelismConfig, /* use_lora */ false); + top_k, activation_type, norm_mode, parallelismConfig, /* use_lora */ false, /* use_deepseek */ false); size_t scale_prob_size = num_rows * num_experts * sizeof(float); size_t src_to_dest_map_size = top_k * num_rows * sizeof(int); size_t selected_expert_size = top_k * num_rows * sizeof(int); diff --git a/cpp/tests/kernels/mixtureOfExpertsTest.cu b/cpp/tests/kernels/mixtureOfExpertsTest.cu index 609e7df23f21..4592a66969bb 100644 --- a/cpp/tests/kernels/mixtureOfExpertsTest.cu +++ b/cpp/tests/kernels/mixtureOfExpertsTest.cu @@ -239,6 +239,7 @@ protected: bool mUseBias = true; bool mUseLora = false; + bool mUseDeepSeek = false; bool mIsGated = false; int64_t mGatedMultiplier = 1; @@ -288,8 +289,8 @@ protected: // Expert weights size_t const weight_size = hidden_size * (hidden_size * 4) * num_experts * sizeof(WeightStorage) * num_gemms; // Workspace size - size_t const workspace_size = this->mMoERunner.getWorkspaceSize( - num_tokens, hidden_size, hidden_size * 4, num_experts, k, this->mActType, mNormMode, {}, mUseLora); + size_t const workspace_size = this->mMoERunner.getWorkspaceSize(num_tokens, hidden_size, hidden_size * 4, + num_experts, k, this->mActType, mNormMode, {}, mUseLora, mUseDeepSeek); // The input/output buffers size_t const in_out_size = 2 * num_tokens * hidden_size * sizeof(DataType); @@ -333,8 +334,8 @@ protected: mTotalTokens += num_tokens; } - size_t workspace_size = mMoERunner.getWorkspaceSize( - mTotalTokens, mHiddenSize, mInterSize, mNumExperts, mK, mActType, mNormMode, parallelism_config, mUseLora); + size_t workspace_size = mMoERunner.getWorkspaceSize(mTotalTokens, mHiddenSize, mInterSize, mNumExperts, mK, + mActType, mNormMode, parallelism_config, mUseLora, mUseDeepSeek); auto const stream = mStream->get(); @@ -776,12 +777,14 @@ protected: } LoraParams lora_params; + BlockScaleParams blockscale_params; mMoERunner.setTactic(tactic1, tactic2); mMoERunner.runMoe(mInputTensor, mInputProbabilities, weight1_ptr, bias1_ptr, mActType, weight2_ptr, bias2_ptr, quant_params, mTotalTokens, mHiddenSize, mInterSize / parallelism_config.tp_size, mNumExperts, mK, mWorkspace, mFinalOutput, mFinished, mActiveRows, mScaleProbs, mSourceToExpandedMap, mSelectedExpert, - mSparseMixerEpsilon, parallelism_config, mNormMode, mUseLora, lora_params, stream); + mSparseMixerEpsilon, parallelism_config, mNormMode, mUseLora, lora_params, mUseDeepSeek, blockscale_params, + stream); check_cuda_error(cudaStreamSynchronize(stream)); } diff --git a/examples/deepseek_v3/README.md b/examples/deepseek_v3/README.md index 7cb200b904cc..c8b64a86a998 100644 --- a/examples/deepseek_v3/README.md +++ b/examples/deepseek_v3/README.md @@ -16,21 +16,17 @@ This document shows how to build and run the `DeepSeek-v3` model in TensorRT-LLM ## Support Matrix -| Model | FP16 | BF16 | FP8 | W8A16 | W4A16 | TP | EP | IB | -| :------------- | :---: | :---: | :---: | :-----: | :-----: | :-----: | :-----: | :-----: | -| DeepSeek-V3 | Y | Y | . | Y | Y | Y | Y | Y | +| Model | FP16 | BF16 | FP8 | TP | EP | IB | +| :------------- | :---: | :---: | :---: | :-----: | :-----: | :-----: | +| DeepSeek-V3 | Y | Y | Y | Y | Y | Y | -- W8A16: INT8 Weight-Only -- W4A16: INT4 Weight-Only - TP: Tensor Parallel - EP: Expert Parallel - IB: Inflight Batching -- FP8: Support for FP8 is currently in progress and will be released soon -***Please Note:*** +***Please Note:*** - Prefer using BF16 over FP16 for DeepSeek-V3 since model original training precision is FP8 and we found direct convert FP8 -> FP16 may cause unknown accuracy issues. -- Although Int8 and Int4 weight-only quantization can reduce the amount of total GPU memory needed, they may affect the accuracy to a certain degree. ## Prerequisite @@ -41,13 +37,23 @@ First, please download DeepSeek-V3 weights from HF https://huggingface.co/deepse git lfs install git clone https://huggingface.co/deepseek-ai/DeepSeek-V3-Base ``` +**Optional**: Convert the FP8 checkpoint to BF16. This is not necessary unless you want to run the model E2E in BF16 precision. +```bash +git clone https://github.com/deepseek-ai/DeepSeek-V3.git +cd DeepSeek-V3/inference/ +python fp8_cast_bf16.py --input-fp8-hf-path /path/to/DeepSeek-V3 --output-bf16-hf-path /path/to/deepseek-v3-bf16 +cp /path/to/DeepSeek-V3/config.json /path/to/DeepSeek-V3/configuration_deepseek.py /path/to/deepseek-v3-bf16/ +``` ## Hardware -The DeepSeek-V3 model requires at least 32x80G GPU memory, model contains 660B parameters, roughly 1.3TB memory (with BF16 precision). +The DeepSeek-V3 model requires at least 16x80G GPU memory, model contains 660B parameters, roughly 1.3TB memory (with BF16 precision). ***Caution: Current TRT-LLM MLA kernel only supports Hopper architecture (SM90). Ampere architecture (SM80 & SM86) will be supported in the future release.*** +Please follow the instructions [here](https://github.com/NVIDIA/TensorRT-LLM/blob/deepseek/docs/source/installation/build-from-source-linux.md#building-a-tensorrt-llm-docker-image +) to achieve a correct docker image. + ## Overview The TensorRT-LLM DeepSeek-V3 implementation can be found in [tensorrt_llm/models/deepseek_v2/model.py](../../tensorrt_llm/models/deepseek_v2/model.py). The TensorRT-LLM Deepseek-V3 example code is located in [`example/deepseek_v3`](./). There is one main file: @@ -58,6 +64,8 @@ In addition, there are three shared files in the parent folder [`examples`](../) * [`../run.py`](../run.py) to run the model inference output by giving an input text. +* [`../mmlu.py`](../mmlu.py) to running score script from https://github.com/declare-lab/instruct-eval to compare HF model and TensorRT-LLM model on the MMLU dataset. + ## Usage @@ -67,50 +75,49 @@ The TensorRT-LLM DeepSeek-V3 example code is located at [examples/deepseek_v3](. Below is the step-by-step to run DeepSeek-V3 with TensorRT-LLM. -Firstly, convert FP8 weights to BF16: -```bash -git clone https://github.com/deepseek-ai/DeepSeek-V3.git -cd DeepSeek-V3/inference/ -python fp8_cast_bf16.py --input-fp8-hf-path /path/to/DeepSeek-V3 --output-bf16-hf-path /path/to/deepseek-v3-bf16 -cp /path/to/DeepSeek-V3/config.json /path/to/DeepSeek-V3/configuration_deepseek.py /path/to/deepseek-v3-bf16/ -``` -Secondly, the BF16 checkpoint will be converted to the TensorRT-LLM checkpoint format by apply [`convert_checkpoint.py`](./convert_checkpoint.py). After that, the TensorRT engine(s) can be built with the TensorRT-LLM checkpoint. +Firstly, convert the checkpoint to the TensorRT-LLM checkpoint format by running [`convert_checkpoint.py`](./convert_checkpoint.py). After that, the TensorRT engine(s) can be built with the TensorRT-LLM checkpoint. +To convert FP8 checkpoint: ```bash -# Convert Deepseek-v3 HF weights to TensorRT-LLM checkpoint in BF16. +# Convert Deepseek-v3 HF Native FP8 weights to TensorRT-LLM checkpoint. python convert_checkpoint.py --model_dir ./DeepSeek-V3 \ - --output_dir ./trtllm_checkpoint_deepseek_v3_32gpu_bf16 \ + --output_dir ./trtllm_checkpoint_deepseek_v3_16gpu_fp8 \ --dtype bfloat16 \ - --tp_size 32 \ + --use_fp8_weights \ + --tp_size 16 \ --workers 8 # using multiple workers can accelerate the conversion process +``` - -# Use Weight-Only Int8 quantization -python convert_checkpoint.py --model_dir ./DeepSeek-V3 \ - --output_dir ./trtllm_checkpoint_deepseek_v3_32gpu_bf16 \ - --dtype bfloat16 \ - --tp_size 32 \ - --use_weight_only \ - --weight_only_precision int8 \ - --workers 8 - -# Use Weight-Only Int4 quantization +To convert BF16 checkpoint: +```bash +# Convert Deepseek-v3 HF weights to TensorRT-LLM checkpoint in BF16. python convert_checkpoint.py --model_dir ./DeepSeek-V3 \ --output_dir ./trtllm_checkpoint_deepseek_v3_32gpu_bf16 \ --dtype bfloat16 \ --tp_size 32 \ - --use_weight_only \ - --weight_only_precision int4 \ - --workers 8 + --workers 8 # using multiple workers can accelerate the conversion process ``` We observed the checkpoint conversion time took hours, while using a significant amount of CPU memory, please adjust the `--workers` parameter to balance your time and memory consumption. After the checkpoint conversion, the TensorRT engine(s) can be built with the TensorRT-LLM checkpoint. +For FP8: +```bash +# Build FP8 engine +trtllm-build --checkpoint_dir ./trtllm_checkpoint_deepseek_v3_16gpu_fp8 \ + --output_dir ./trtllm_engines/deepseek_v3/fp8/tp16-sel4096-isl2048-bs4 \ + --max_batch_size 4 \ + --max_seq_len 4096 \ + --max_input_len 2048 \ + --use_paged_context_fmha enable \ + --workers 8 +``` + +For BF16: ```bash -# Build engine +# Build BF16 engine trtllm-build --checkpoint_dir ./trtllm_checkpoint_deepseek_v3_32gpu_bf16 \ --output_dir ./trtllm_engines/deepseek_v3/bf16/tp32-sel4096-isl2048-bs4 \ --gpt_attention_plugin bfloat16 \ @@ -124,13 +131,13 @@ trtllm-build --checkpoint_dir ./trtllm_checkpoint_deepseek_v3_32gpu_bf16 \ ***Caution: `--max_batch_size` and `--max_seq_len` are the main factors to determine how many GPU memory will be used during runtime, so later when try to run e.g., `summarize.py` or `mmlu.py` or `gptManagerBenchmark.cpp`may need adjust `--max_batch_size` and `--max_seq_len` accordingly to avoid OOM.(meaning rebuild TensorRT engine with smaller `--max_batch_size` and `--max_seq_len` if needed based on GPU memory size), there is beautiful technical log perf-best-practices.md (https://github.com/NVIDIA/TensorRT-LLM/blob/main/docs/source/performance/perf-best-practices.md) explained the mechanism.*** -Test the engine with [run.py](../run.py) script: +Test the FP8 engines with [run.py](../run.py) script: ``` # run.sh python3 ../run.py --input_text "Today is a nice day." \ --max_output_len 30 \ --tokenizer_dir ./DeepSeek-V3 \ - --engine_dir ./trtllm_engines/deepseek_v3/bf16/tp32-sel4096-isl2048-bs4 \ + --engine_dir ./trtllm_engines/deepseek_v3/fp8/tp16-sel4096-isl2048-bs4 \ --top_p 0.95 \ --temperature 0.3 @@ -139,7 +146,7 @@ python3 ../run.py --input_text "Today is a nice day." \ For multi-nodes inference, let's take Slurm as an example using above command (run.sh): ```bash -srun -N 4 -w node-[1-4] --gres=gpu:8 --ntasks-per-node 8 \ +srun -N 2 -w node-[1-2] --gres=gpu:8 --ntasks-per-node 8 \ --container-image tensorrt_llm/release:latest \ --container-mounts ${PWD}:/workspace \ sh /workspace/command/run.sh @@ -169,3 +176,84 @@ Input [Text 0]: "Today is a nice day." Output [Text 0 Beam 0]: " I am going to the park with my friends. We are going to play soccer. We are going" ``` +At last, we can evaluate the model with [mmlu.py](../mmlu.py) script: + + +```bash +# Download MMLU dataset +mkdir mmlu_data && cd mmlu_data +wget https://people.eecs.berkeley.edu/~hendrycks/data.tar && tar -xf data.tar +# Run MMLU evaluation +python3 mmlu.py \ + --hf_model_dir ${MODEL_DIR} \ + --engine_dir ./trtllm_engines/deepseek_v3/fp8/tp16-sel4096-isl2048-bs4 \ + --data_dir mmlu_data \ + --test_trt_llm 2>&1 | tee ${ENGINE_DIR}/test_with_mmlu.log +``` + +and the output will be like: + +``` +Average accuracy 0.926 - high_school_macroeconomics +Average accuracy 0.752 - high_school_mathematics +Average accuracy 0.954 - high_school_microeconomics +Average accuracy 0.848 - high_school_physics +Average accuracy 0.967 - high_school_psychology +Average accuracy 0.861 - high_school_statistics +Average accuracy 0.956 - high_school_us_history +Average accuracy 0.954 - high_school_world_history +Average accuracy 0.861 - human_aging +Average accuracy 0.931 - human_sexuality +Average accuracy 0.975 - international_law +Average accuracy 0.907 - jurisprudence +Average accuracy 0.920 - logical_fallacies +Average accuracy 0.848 - machine_learning +Average accuracy 0.951 - management +Average accuracy 0.957 - marketing +Average accuracy 0.950 - medical_genetics +Average accuracy 0.957 - miscellaneous +Average accuracy 0.870 - moral_disputes +Average accuracy 0.798 - moral_scenarios +Average accuracy 0.918 - nutrition +Average accuracy 0.916 - philosophy +Average accuracy 0.932 - prehistory +Average accuracy 0.869 - professional_accounting +Average accuracy 0.714 - professional_law +Average accuracy 0.956 - professional_medicine +Average accuracy 0.908 - professional_psychology +Average accuracy 0.800 - public_relations +Average accuracy 0.869 - security_studies +Average accuracy 0.960 - sociology +Average accuracy 0.950 - us_foreign_policy +Average accuracy 0.578 - virology +Average accuracy 0.930 - world_religions +Average accuracy 0.852 - math +Average accuracy 0.874 - health +Average accuracy 0.905 - physics +Average accuracy 0.936 - business +Average accuracy 0.958 - biology +Average accuracy 0.825 - chemistry +Average accuracy 0.888 - computer science +Average accuracy 0.912 - economics +Average accuracy 0.890 - engineering +Average accuracy 0.851 - philosophy +Average accuracy 0.917 - other +Average accuracy 0.932 - history +Average accuracy 0.944 - geography +Average accuracy 0.904 - politics +Average accuracy 0.936 - psychology +Average accuracy 0.949 - culture +Average accuracy 0.744 - law +Average accuracy 0.883 - STEM +Average accuracy 0.827 - humanities +Average accuracy 0.926 - social sciences +Average accuracy 0.898 - other (business, health, misc.) +Average accuracy: 0.877 +``` + +**Known Issue** + +1. The memory allocation for MoE is too large. + +This issue prevents running larger batch sizes and long sequence inputs. We will optimize and fix this issue soon. + diff --git a/examples/deepseek_v3/convert_checkpoint.py b/examples/deepseek_v3/convert_checkpoint.py index 7cbc7935d225..338bb6436565 100644 --- a/examples/deepseek_v3/convert_checkpoint.py +++ b/examples/deepseek_v3/convert_checkpoint.py @@ -77,23 +77,10 @@ def parse_arguments(): 'To shard it along hidden dimension, set embedding_sharding_dim=1' 'Note: embedding sharing is only enabled when embedding_sharding_dim=0') - parser.add_argument( - '--use_weight_only', - default=False, - action="store_true", - help='Quantize weights for the various GEMMs to INT4/INT8.' - 'See --weight_only_precision to set the precision') - parser.add_argument( - '--weight_only_precision', - const='int8', - type=str, - nargs='?', - default='int8', - choices=['int8', 'int4'], - help= - 'Define the precision for the weights when using weight-only quantization.' - 'You must also use --use_weight_only for that argument to have an impact.' - ) + parser.add_argument('--use_fp8_weights', + default=False, + action="store_true", + help='Use native FP8 weights of DeepSeek-V3.') parser.add_argument('--output_dir', type=str, @@ -131,36 +118,18 @@ def parse_arguments(): help= 'Only save the model config w/o read and converting weights, be careful, this is for debug only' ) - parser.add_argument( - '--disable_weight_only_quant_plugin', - default=False, - action="store_true", - help= - 'By default, using plugin implementation for weight quantization. Enabling disable_weight_only_quant_plugin flag will use ootb implementation instead of plugin.' - 'You must also use --use_weight_only for that argument to have an impact' - ) - # Add quantization related feature later + args = parser.parse_args() return args -def precision_to_config(precision, quant_config) -> QuantConfig: - '''update config dict for weight-only quantization - ''' - quant_config = QuantConfig() - precision_to_algo = {'int8': QuantAlgo.W8A16, 'int4': QuantAlgo.W4A16} - quant_config.quant_algo = precision_to_algo.get(precision) - return quant_config - - def args_to_quant_config(args: argparse.Namespace) -> QuantConfig: '''return config dict with quantization info based on the command line args ''' quant_config = QuantConfig() - if args.use_weight_only: - quant_config = precision_to_config(args.weight_only_precision, - quant_config) + if args.use_fp8_weights: + quant_config.quant_algo = QuantAlgo.FP8_CURRENT_SCALING return quant_config diff --git a/tensorrt_llm/auto_parallel/tensor_parallel/plugin_nodes/gpt_attention_node.py b/tensorrt_llm/auto_parallel/tensor_parallel/plugin_nodes/gpt_attention_node.py index b63a73159d93..a44f55c1deb1 100644 --- a/tensorrt_llm/auto_parallel/tensor_parallel/plugin_nodes/gpt_attention_node.py +++ b/tensorrt_llm/auto_parallel/tensor_parallel/plugin_nodes/gpt_attention_node.py @@ -47,9 +47,12 @@ class IdxEntry(Enum): MROPE_POSITION_DELTAS = auto() HOST_RUNTIME_PERF_KNOBS = auto() HOST_CONTEXT_PROGRESS = auto() - MLA_FUSED_Q_PROJ_TENSOR = auto() MLA_Q_B_PROJ_TENSOR = auto() MLA_KV_B_PROJ_TENSOR = auto() + MLA_K_B_PROJ_TRANS_TENSOR = auto() + MLA_Q_B_SCALE_TENSOR = auto() + MLA_KV_B_SCALE_TENSOR = auto() + MLA_K_B_TRANS_SCALE_TENSOR = auto() LOGN_SCALING = auto() @@ -78,6 +81,9 @@ def __init__(self, plugin_info): self.is_spec_decoding_enabled = bool( plugin_info.pfc_as_list['is_spec_decoding_enabled'][0]) self.is_mla_enabled = bool(plugin_info.pfc_as_list['is_mla_enabled'][0]) + self.is_ptp128c_enabled = bool( + plugin_info.pfc_as_list['is_ptp128c_enabled'][0]) + self.is_fp8_model = bool(plugin_info.pfc_as_list['is_fp8_model'][0]) self.use_logn_scaling = bool( plugin_info.pfc_as_list['use_logn_scaling'][0]) self.init_entry_to_index() @@ -157,12 +163,18 @@ def is_entry_used(self, entry: IdxEntry) -> bool: return True elif entry == IdxEntry.HOST_CONTEXT_PROGRESS: return True - elif entry == IdxEntry.MLA_FUSED_Q_PROJ_TENSOR: - return self.is_mla_enabled elif entry == IdxEntry.MLA_Q_B_PROJ_TENSOR: return self.is_mla_enabled elif entry == IdxEntry.MLA_KV_B_PROJ_TENSOR: return self.is_mla_enabled + elif entry == IdxEntry.MLA_K_B_PROJ_TRANS_TENSOR: + return self.is_mla_enabled and self.is_ptp128c_enabled and self.is_fp8_model + elif entry == IdxEntry.MLA_Q_B_SCALE_TENSOR: + return self.is_mla_enabled and self.is_ptp128c_enabled and self.is_fp8_model + elif entry == IdxEntry.MLA_KV_B_SCALE_TENSOR: + return self.is_mla_enabled and self.is_ptp128c_enabled and self.is_fp8_model + elif entry == IdxEntry.MLA_K_B_TRANS_SCALE_TENSOR: + return self.is_mla_enabled elif entry == IdxEntry.LOGN_SCALING: return self.use_logn_scaling else: diff --git a/tensorrt_llm/builder.py b/tensorrt_llm/builder.py index 32382b1ec8f4..d0f49bd0bfd0 100644 --- a/tensorrt_llm/builder.py +++ b/tensorrt_llm/builder.py @@ -1148,6 +1148,7 @@ def build(model: PretrainedModel, build_config: BuildConfig) -> Engine: disable_weight_only_quant_plugin = model.config.disable_weight_only_quant_plugin if hasattr( model.config, 'disable_weight_only_quant_plugin') else False use_fp8_rowwise = model.config.quant_mode.has_fp8_rowwise() + use_fp8_current_scaling = model.config.quant_mode.has_fp8_current_scaling() if build_config.plugin_config.manage_weights: if use_weight_only and disable_weight_only_quant_plugin: @@ -1167,6 +1168,11 @@ def build(model: PretrainedModel, build_config: BuildConfig) -> Engine: if use_fp8_rowwise: network.plugin_config.set_fp8_rowwise_quant_plugins(model.config.dtype) + + if use_fp8_current_scaling: + network.plugin_config.set_fp8_current_scaling_gemm_plugins( + model.config.dtype) + nccl_plugin = model.config.dtype if model.config.mapping.world_size > 1 else None network.plugin_config.set_nccl_plugin(nccl_plugin) diff --git a/tensorrt_llm/functional.py b/tensorrt_llm/functional.py index 7a20c1dd8fec..7c7085477e3b 100755 --- a/tensorrt_llm/functional.py +++ b/tensorrt_llm/functional.py @@ -4837,9 +4837,14 @@ def gpt_attention( qk_nope_head_dim: int = 0, qk_rope_head_dim: int = 0, v_head_dim: int = 0, - fused_q_proj: Optional[Tensor] = None, + is_ptp128c_enabled_flag: bool = False, + is_fp8_model_flag: bool = False, q_b_proj: Optional[Tensor] = None, kv_b_proj: Optional[Tensor] = None, + k_b_proj_trans: Optional[Tensor] = None, + q_b_scale: Optional[Tensor] = None, + kv_b_scale: Optional[Tensor] = None, + k_b_trans_scale: Optional[Tensor] = None, skip_attn=None, cp_group: List[int] = [0], cp_size: int = 1, @@ -5247,6 +5252,12 @@ def gpt_attention( v_head_dim = trt.PluginField("v_head_dim", np.array(v_head_dim, dtype=np.int32), trt.PluginFieldType.INT32) + is_ptp128c_enabled = trt.PluginField( + "is_ptp128c_enabled", np.array(is_ptp128c_enabled_flag, dtype=np.int8), + trt.PluginFieldType.INT8) + is_fp8_model = trt.PluginField("is_fp8_model", + np.array(is_fp8_model_flag, dtype=np.int8), + trt.PluginFieldType.INT8) p_dtype = default_net().plugin_config.gpt_attention_plugin pf_type = trt.PluginField( "type_id", np.array([int(str_dtype_to_trt(p_dtype))], np.int32), @@ -5366,7 +5377,8 @@ def gpt_attention( is_spec_decoding_enabled, spec_decoding_is_generation_length_variable, spec_decoding_max_generation_length, is_mla_enabled, q_lora_rank, kv_lora_rank, qk_nope_head_dim, qk_rope_head_dim, v_head_dim, - skip_attn_pf, cp_size, cp_rank, cp_group, use_logn_scaling + is_ptp128c_enabled, is_fp8_model, skip_attn_pf, cp_size, cp_rank, + cp_group, use_logn_scaling ]) attn_plug = attn_plg_creator.create_plugin("causal_attn", pfc) @@ -5463,10 +5475,15 @@ def gpt_attention( plug_inputs += [host_context_progress] if is_mla_enabled_flag: - assert fused_q_proj is not None assert q_b_proj is not None assert kv_b_proj is not None - plug_inputs += [fused_q_proj, q_b_proj, kv_b_proj] + assert k_b_proj_trans is not None + plug_inputs += [q_b_proj, kv_b_proj, k_b_proj_trans] + if is_ptp128c_enabled_flag and is_fp8_model_flag: + assert q_b_scale is not None + assert kv_b_scale is not None + assert k_b_trans_scale is not None + plug_inputs += [q_b_scale, kv_b_scale, k_b_trans_scale] if skip_attn is not None: plug_inputs += [skip_attn] diff --git a/tensorrt_llm/layers/attention.py b/tensorrt_llm/layers/attention.py index 4bfc86ca8c43..aad6bcadf738 100755 --- a/tensorrt_llm/layers/attention.py +++ b/tensorrt_llm/layers/attention.py @@ -1956,6 +1956,13 @@ def __init__( self.rotary_scaling = rotary_scaling self.shard_dim = 1 + self.is_ptp128c_enabled_flag = False + self.is_fp8_model_flag = False + if quant_mode.has_fp8_current_scaling(): + self.is_ptp128c_enabled_flag = True + if not quant_mode.has_fp8_quantize_weights_on_demand(): + self.is_fp8_model_flag = True + def yarn_get_mscale(scale=1, mscale=1): if scale <= 1: return 1.0 @@ -2015,20 +2022,46 @@ def yarn_get_mscale(scale=1, mscale=1): self.kv_a_layernorm = RmsNorm(kv_lora_rank, dtype=dtype, eps=eps) - self.fused_q_proj = Parameter( - shape=(self.num_attention_heads * - (self.kv_lora_rank + self.qk_rope_head_dim), - self.q_lora_rank), - dtype=dtype) + if quant_mode.has_fp8_current_scaling( + ) and not quant_mode.has_fp8_quantize_weights_on_demand(): + self.kv_b_proj_scale = Parameter( + shape=(int(self.num_attention_heads * self.qk_nope_head_dim / + 128 * 2), int(self.kv_lora_rank / 128)), + dtype='float32') + self.k_b_proj_trans_scale = Parameter( + shape=(int(self.num_attention_heads * self.kv_lora_rank / 128), + int(self.qk_nope_head_dim / 128)), + dtype='float32') + self.q_b_proj_scale = Parameter(shape=(int( + self.num_attention_heads * + (self.qk_nope_head_dim + self.qk_rope_head_dim) / 128), + int(self.q_lora_rank / 128)), + dtype='float32') + set_obj_attrs(self.q_b_proj_scale, { + "weight_loader": self.weight_loader, + }) + set_obj_attrs(self.k_b_proj_trans_scale, { + "weight_loader": self.weight_loader, + }) + set_obj_attrs(self.kv_b_proj_scale, { + "weight_loader": self.weight_loader, + }) + + mla_weight_type = dtype if quant_mode.has_fp8_quantize_weights_on_demand() else 'fp8' + self.kv_b_proj = Parameter( shape=(self.num_attention_heads * self.qk_nope_head_dim * 2, self.kv_lora_rank), - dtype=dtype) + dtype=mla_weight_type) + self.k_b_proj_trans = Parameter( + shape=(self.num_attention_heads * self.kv_lora_rank, + self.qk_nope_head_dim), + dtype=mla_weight_type) self.q_b_proj = Parameter( shape=(self.num_attention_heads * (self.qk_nope_head_dim + self.qk_rope_head_dim), self.q_lora_rank), - dtype=dtype) + dtype=mla_weight_type) self.dense = RowLinear(tp_size * self.num_attention_heads * self.v_head_dim, hidden_size, @@ -2036,15 +2069,15 @@ def yarn_get_mscale(scale=1, mscale=1): dtype=dtype, tp_group=tp_group, tp_size=tp_size) - set_obj_attrs(self.fused_q_proj, { - "weight_loader": self.weight_loader, - }) set_obj_attrs(self.q_b_proj, { "weight_loader": self.weight_loader, }) set_obj_attrs(self.kv_b_proj, { "weight_loader": self.weight_loader, }) + set_obj_attrs(self.k_b_proj_trans, { + "weight_loader": self.weight_loader, + }) def weight_loader(self, mapping: Mapping, param: Parameter, loaded_weight: torch.Tensor): @@ -2206,9 +2239,17 @@ def forward(self, qk_nope_head_dim=self.qk_nope_head_dim, qk_rope_head_dim=self.qk_rope_head_dim, v_head_dim=self.v_head_dim, - fused_q_proj=self.fused_q_proj.value, + is_ptp128c_enabled_flag=self.is_ptp128c_enabled_flag, + is_fp8_model_flag=self.is_fp8_model_flag, + k_b_proj_trans=self.k_b_proj_trans.value, q_b_proj=self.q_b_proj.value, - kv_b_proj=self.kv_b_proj.value) + kv_b_proj=self.kv_b_proj.value, + q_b_scale=self.q_b_proj_scale.value + if self.is_fp8_model_flag else None, + kv_b_scale=self.kv_b_proj_scale.value + if self.is_fp8_model_flag else None, + k_b_trans_scale=self.k_b_proj_trans_scale.value + if self.is_fp8_model_flag else None) context = self.dense(context) @@ -2227,77 +2268,96 @@ def split(v, tp_size, idx, dim=0): else: return torch.chunk(v, tp_size, dim=dim)[idx].contiguous() - if tllm_key.endswith("kv_b_proj"): + if tllm_key.endswith("kv_b_proj") or tllm_key.endswith( + "kv_b_proj_scale"): + if tllm_key.endswith("kv_b_proj"): + qk_nope_head_dim = self.qk_nope_head_dim + v_head_dim = self.v_head_dim + kv_lora_rank = self.kv_lora_rank + else: + qk_nope_head_dim = int(self.qk_nope_head_dim / 128) + v_head_dim = int(self.v_head_dim / 128) + kv_lora_rank = int(self.kv_lora_rank / 128) kv_b_proj = weights.unflatten(0, [ self.num_attention_heads * self.tp_size, - self.qk_nope_head_dim + self.v_head_dim + qk_nope_head_dim + v_head_dim ]) splited_kv_b_proj = split(kv_b_proj, self.tp_size, self.tp_rank, dim=0) k_nope_weight, v_weight = splited_kv_b_proj.split( - [self.qk_nope_head_dim, self.v_head_dim], + [qk_nope_head_dim, v_head_dim], dim=1, ) kv_b_proj_weight = torch.concat([ k_nope_weight.reshape( - self.num_attention_heads * self.qk_nope_head_dim, - self.kv_lora_rank), - v_weight.reshape(self.num_attention_heads * self.v_head_dim, - self.kv_lora_rank) + self.num_attention_heads * qk_nope_head_dim, kv_lora_rank), + v_weight.reshape(self.num_attention_heads * v_head_dim, + kv_lora_rank) ], dim=0) return {tllm_key: kv_b_proj_weight} - elif tllm_key.endswith("q_b_proj"): - q_b_proj = weights.unflatten(0, [ - self.num_attention_heads * self.tp_size, - self.qk_nope_head_dim + self.qk_rope_head_dim - ]) - splited_q_b_proj = split(q_b_proj, - self.tp_size, - self.tp_rank, - dim=0) - q_b_proj_weight = splited_q_b_proj.reshape( - self.num_attention_heads * - (self.qk_nope_head_dim + self.qk_rope_head_dim), - self.q_lora_rank) + elif tllm_key.endswith("q_b_proj") or tllm_key.endswith( + "q_b_proj_scale"): + if tllm_key.endswith("q_b_proj"): + q_b_proj = weights.unflatten(0, [ + self.num_attention_heads * self.tp_size, + self.qk_nope_head_dim + self.qk_rope_head_dim + ]) + splited_q_b_proj = split(q_b_proj, + self.tp_size, + self.tp_rank, + dim=0) + q_b_proj_weight = splited_q_b_proj.reshape( + self.num_attention_heads * + (self.qk_nope_head_dim + self.qk_rope_head_dim), + self.q_lora_rank) + else: + splited_q_b_proj = split(weights, + self.tp_size, + self.tp_rank, + dim=0) + q_b_proj_weight = splited_q_b_proj.reshape( + int(self.num_attention_heads * + (self.qk_nope_head_dim + self.qk_rope_head_dim) / 128), + int(self.q_lora_rank / 128)) return {tllm_key: q_b_proj_weight} - elif tllm_key.endswith("fused_q_proj"): - assert isinstance(weights, list) and len(weights) == 2 - q_b_proj = weights[0].unflatten(0, [ - self.num_attention_heads * self.tp_size, - self.qk_nope_head_dim + self.qk_rope_head_dim - ]) - splited_q_b_proj = split(q_b_proj, - self.tp_size, - self.tp_rank, - dim=0) - kv_b_proj = weights[1].unflatten(0, [ - self.num_attention_heads * self.tp_size, - self.qk_nope_head_dim + self.v_head_dim - ]) - splited_kv_b_proj = split(kv_b_proj, - self.tp_size, - self.tp_rank, - dim=0) - q_nope_weight, q_pe_weight = splited_q_b_proj.split( - [self.qk_nope_head_dim, self.qk_rope_head_dim], - dim=1, - ) - k_nope_weight, _ = splited_kv_b_proj.split( - [self.qk_nope_head_dim, self.v_head_dim], - dim=1, - ) - fused_q_nope_weight = torch.einsum( - 'hdq,hdk->hkq', - q_nope_weight, - k_nope_weight, - ) - fused_q_weight = torch.cat( - [fused_q_nope_weight, q_pe_weight], - dim=1, - ).flatten(start_dim=0, end_dim=1) - return {tllm_key: fused_q_weight} + elif tllm_key.endswith("k_b_proj_trans") or tllm_key.endswith( + "k_b_proj_trans_scale"): + if tllm_key.endswith("k_b_proj_trans"): + kv_b_proj = weights.unflatten(0, [ + self.num_attention_heads * self.tp_size, + self.qk_nope_head_dim + self.v_head_dim + ]) + splited_kv_b_proj = split(kv_b_proj, + self.tp_size, + self.tp_rank, + dim=0) + k_nope_weight, v_weight = splited_kv_b_proj.split( + [self.qk_nope_head_dim, self.v_head_dim], + dim=1, + ) + k_nope_weight_trans = k_nope_weight.transpose(2, 1).reshape( + self.num_attention_heads * self.kv_lora_rank, + self.qk_nope_head_dim) + else: + kv_b_proj_scale = weights.unflatten(0, [ + self.num_attention_heads * self.tp_size, + self.qk_nope_head_dim // 128 + self.v_head_dim // 128 + ]) + splited_kv_b_proj_scale = split(kv_b_proj_scale, + self.tp_size, + self.tp_rank, + dim=0) + k_nope_scale, v_scale = splited_kv_b_proj_scale.split( + [self.qk_nope_head_dim // 128, self.v_head_dim // 128], + dim=1, + ) + k_nope_weight_trans = k_nope_scale.transpose(2, 1).reshape( + self.num_attention_heads * self.kv_lora_rank // 128, + self.qk_nope_head_dim // 128) + + return {tllm_key: k_nope_weight_trans} else: return {tllm_key: weights} diff --git a/tensorrt_llm/layers/moe.py b/tensorrt_llm/layers/moe.py index c1e83c933207..d3273b670d47 100755 --- a/tensorrt_llm/layers/moe.py +++ b/tensorrt_llm/layers/moe.py @@ -238,6 +238,19 @@ def from_parameter(x): "max_low_rank", np.array(lora_max_low_rank, dtype=np.int32), trt.PluginFieldType.INT32) + # Customized plugin inputs for DeepSeek-V3 + use_deepseek = trt.PluginField( + "use_deepseek", + np.array([int(quant_mode.has_fp8_current_scaling())], dtype=np.int32), + trt.PluginFieldType.INT32) + + use_deepseek_with_native_fp8_weights = not quant_mode.has_fp8_quantize_weights_on_demand( + ) + use_deepseek_with_native_fp8_weights = trt.PluginField( + "use_deepseek_with_native_fp8_weights", + np.array([int(use_deepseek_with_native_fp8_weights)], dtype=np.int32), + trt.PluginFieldType.INT32) + pfc_inputs = [ p_remove_input_padding, p_num_experts, p_top_k, p_expert_hidden_size, p_expert_inter_size, p_activation_type, p_type_id, p_weight_type_id, @@ -250,6 +263,9 @@ def from_parameter(x): if use_lora: pfc_inputs += [p_lora_type_id, p_max_low_rank] + if quant_mode.has_fp8_current_scaling(): + pfc_inputs += [use_deepseek, use_deepseek_with_native_fp8_weights] + pfc = trt.PluginFieldCollection(pfc_inputs) # Create the plugin with our constant inputs to the constructor @@ -274,6 +290,12 @@ def from_parameter(x): assert expert_scale_2 plugin_inputs += [expert_scale_1, expert_scale_2] + if quant_mode.has_fp8_current_scaling( + ) and not quant_mode.has_fp8_quantize_weights_on_demand(): + assert expert_scale_1 + assert expert_scale_2 + plugin_inputs += [expert_scale_1, expert_scale_2] + # Add conditional inputs if quant_mode.has_fp8_qdq(): assert expert_scale_3 @@ -284,6 +306,12 @@ def from_parameter(x): assert output_dtype == trt.fp8 plugin_inputs += [expert_scale_4] + if quant_mode.has_fp8_current_scaling( + ) and hidden_states_raw.dtype == trt.fp8: + assert expert_scale_3 + assert expert_scale_4 + plugin_inputs += [expert_scale_3, expert_scale_4] + if use_lora: if quant_mode.has_fp8_qdq(): assert act_scale @@ -353,6 +381,11 @@ def __init__(self, in_features: int, out_features: int, self.tp_dim = tp_dim self.is_padded = False + if quant_mode.has_fp8_current_scaling( + ) and not quant_mode.has_fp8_quantize_weights_on_demand(): + self.dtype = 'fp8' + self.weight_dtype = 'fp8' + if quant_mode.is_weight_only(): bytes_per_col_scale = 2 if quant_mode.is_int4_weight_only() else 1 # We use a different shape here because the quantized weights have their own layout @@ -365,7 +398,7 @@ def __init__(self, in_features: int, out_features: int, self.register_parameter('per_channel_scale', None) self.weight = Parameter(shape=self.expert_shape, - dtype=weight_dtype, + dtype=self.weight_dtype, prefer_managed=True) if has_bias: @@ -379,12 +412,20 @@ def __init__(self, in_features: int, out_features: int, dtype=trt.float32) self.weights_scaling_factor = Parameter(shape=(experts_per_node, 1), dtype=trt.float32) + elif quant_mode.has_fp8_current_scaling( + ) and not quant_mode.has_fp8_quantize_weights_on_demand(): + self.register_parameter('activation_scaling_factor', None) + self.weights_scaling_factor = Parameter(shape=(experts_per_node, + out_features // 128, + in_features // 128), + dtype=trt.float32) else: self.register_parameter('activation_scaling_factor', None) self.register_parameter('weights_scaling_factor', None) def postprocess(self, tllm_key, weights, **kwargs): - if tllm_key.endswith("weight"): + if tllm_key.endswith("weight") or tllm_key.endswith( + "weights_scaling_factor"): if isinstance(weights, torch.Tensor): weights = [weights] if "fc" in tllm_key: @@ -395,9 +436,11 @@ def postprocess(self, tllm_key, weights, **kwargs): dim=-2) elif "proj" in tllm_key: weights = torch.stack(weights) - weights = weights.to(str_dtype_to_torch(self.dtype)) + if tllm_key.endswith("weight"): + weights = weights.to(str_dtype_to_torch(self.dtype)) - if not self.quant_mode.has_any_quant(): + if not self.quant_mode.has_any_quant( + ) or self.quant_mode.has_fp8_current_scaling(): return weights elif self.quant_mode.is_weight_only(): if "per_channel_scale" in tllm_key: @@ -476,14 +519,18 @@ def __init__(self, self.weight_dtype = trt.int8 elif quant_mode.has_fp8_qdq(): self.weight_dtype = trt.fp8 - + elif quant_mode.has_fp8_current_scaling( + ) and not quant_mode.has_fp8_quantize_weights_on_demand(): + self.weight_dtype = trt.fp8 rank_experts = self.mapping.ep_experts(self.num_experts) self.wrapper_tllm_to_externel_key_dict = { "mlp": "block_sparse_moe", "proj": [f"experts.{expert}.w2" for expert in rank_experts], "fc": [f"experts.{expert}.w3" for expert in rank_experts] + - [f"experts.{expert}.w1" for expert in rank_experts] + [f"experts.{expert}.w1" for expert in rank_experts], + "weights_scaling_factor": + "weight_scale_inv", } # Since output dimension is usually low (in the order of 10s), no TP at @@ -649,6 +696,32 @@ def forward_experts(self, hidden_states, routing, finished, raise RuntimeError( "Cannot output FP8 value without knowing quantization parameter" ) + elif self.quant_mode.has_fp8_current_scaling( + ) and not self.quant_mode.has_fp8_quantize_weights_on_demand(): + assert self.fc.weight.value.dtype == trt.fp8, ( + "mlp fc weight dtype should be fp8 if not quantize on demand.") + assert self.proj.weight.value.dtype == trt.fp8, ( + "mlp proj weight dtype should be fp8 if not quantize on demand." + ) + hidden_states_quant = hidden_states + if hidden_states_quant.dtype != trt.fp8: + pass + + dtype_quant = self.dtype + weight_dtype_quant = trt.fp8 + + fc1_weight_scales = self.fc.weights_scaling_factor.value + fc2_weight_scales = self.proj.weights_scaling_factor.value + # self.fc.activation_scaling_factor.value + # self.proj.activation_scaling_factor.value + + scale_1 = fc1_weight_scales + scale_2 = fc2_weight_scales + scale_3 = None # TODO(Jerry Shi): change to fc1_act_scales if cs is fused in RMSNorm + scale_4 = None # TODO(Jerry Shi): change to fc2_act_scales if cs is fused in Swiglu + scale_5 = None + + output_dtype_quant = self.dtype else: hidden_states_quant = hidden_states diff --git a/tensorrt_llm/models/deepseek_v2/convert.py b/tensorrt_llm/models/deepseek_v2/convert.py index 158b527fc70f..271481cbc038 100755 --- a/tensorrt_llm/models/deepseek_v2/convert.py +++ b/tensorrt_llm/models/deepseek_v2/convert.py @@ -214,15 +214,14 @@ def convert_layer(l): dim=1, ) - q_nope_weight, q_pe_weight = q_b_proj_weight.split( - [qk_nope_head_dim, qk_rope_head_dim], - dim=1, - ) k_nope_weight, v_weight = kv_b_proj_weight.split( [qk_nope_head_dim, v_head_dim], dim=1, ) + k_nope_weight_trans = k_nope_weight.transpose(2, 1).reshape( + num_heads // mapping.tp_size * kv_lora_rank, qk_nope_head_dim) + if q_lora_rank is None: q_b_proj_weight = q_b_proj_weight.reshape( num_heads * (qk_nope_head_dim + qk_rope_head_dim) // @@ -240,32 +239,21 @@ def convert_layer(l): ], dim=0) - # Fuse matrices for decompression - fused_q_nope_weight = torch.einsum( - 'hdq,hdk->hkq', - q_nope_weight, - k_nope_weight, - ) - fused_q_weight = torch.cat( - [fused_q_nope_weight, q_pe_weight], - dim=1, - ).flatten(start_dim=0, end_dim=1) - weights.update( get_tllm_linear_weight(fused_a_weight, trtllm_prex + 'attention.fused_a.')) weights.update( get_tllm_linear_weight(kv_a_layernorm_weight, trtllm_prex + 'attention.kv_a_layernorm.')) - weights.update( - get_param_weight(fused_q_weight, - trtllm_prex + 'attention.fused_q_proj')) weights.update( get_param_weight(q_b_proj_weight, trtllm_prex + 'attention.q_b_proj')) weights.update( get_param_weight(kv_b_proj_weight, trtllm_prex + 'attention.kv_b_proj')) + weights.update( + get_param_weight(k_nope_weight_trans, + trtllm_prex + 'attention.k_b_proj_trans')) weights.update( get_tllm_linear_weight( o_proj_weight, diff --git a/tensorrt_llm/models/deepseek_v2/model.py b/tensorrt_llm/models/deepseek_v2/model.py index 4fa7798b3263..1e30a1c85a33 100755 --- a/tensorrt_llm/models/deepseek_v2/model.py +++ b/tensorrt_llm/models/deepseek_v2/model.py @@ -69,7 +69,8 @@ def __init__(self, config: DeepSeekV2Config, layer_idx: int): rotary_scaling=config.rotary_scaling, tp_group=config.mapping.tp_group, tp_size=config.mapping.tp_size, - tp_rank=config.mapping.tp_rank) + tp_rank=config.mapping.tp_rank, + quant_mode=config.quant_mode) ### Added deepseek MoE and shared_experts ### First decoder layer: MLA + dense MLP + input_layernorm(RMSNorm) + post_attention_layernorm(RMSNorm) @@ -252,8 +253,11 @@ def from_hugging_face(cls, "q_a_layernrom": "q_a_layernorm", "kv_a_layernorm": "kv_a_layernorm", "q_b_proj": "q_b_proj.weight", + "q_b_proj_scale": "q_b_proj.weight_scale_inv", "kv_b_proj": "kv_b_proj.weight", - "fused_q_proj": ["q_b_proj.weight", "kv_b_proj.weight"], + "kv_b_proj_scale": "kv_b_proj.weight_scale_inv", + "k_b_proj_trans": "kv_b_proj.weight", + "k_b_proj_trans_scale": "kv_b_proj.weight_scale_inv", "shared_expert": "shared_experts", "e_score_correction_bias": "gate.e_score_correction_bias", } @@ -262,8 +266,11 @@ def from_hugging_face(cls, "fused_a": "kv_a_proj_with_mqa", "kv_a_layernorm": "kv_a_layernorm", "q_b_proj": "q_proj.weight", + "q_b_proj_scale": "q_proj.weight_scale_inv", "kv_b_proj": "kv_b_proj.weight", - "fused_q_proj": ["q_proj.weight", "kv_b_proj.weight"], + "kv_b_proj_scale": "kv_b_proj.weight_scale_inv", + "k_b_proj_trans": "kv_b_proj.weight", + "k_b_proj_trans_scale": "kv_b_proj.weight_scale_inv", "shared_expert": "shared_experts", "e_score_correction_bias": "gate.e_score_correction_bias", } diff --git a/tensorrt_llm/plugin/plugin.py b/tensorrt_llm/plugin/plugin.py index c6f4e873455d..2bd68f92f528 100644 --- a/tensorrt_llm/plugin/plugin.py +++ b/tensorrt_llm/plugin/plugin.py @@ -191,6 +191,14 @@ class PluginConfig(metaclass=PluginConfigMeta): "activation and per channel static scales for weights." "Note: It also requires same calibration in checkpoint." }) + _fp8_current_scaling_gemm_plugin: Optional[str] = field( + default=None, + init=False, + metadata={ + "help": + "The quantized GEMM for fp8, which uses per block dynamic scales for " + "activation and per block static scales for weights." + }) _qserve_gemm_plugin: Optional[str] = field( default=None, init=False, @@ -506,6 +514,10 @@ def set_fp8_rowwise_quant_plugins(self, dtype: str = "auto"): self.quantize_tensor_plugin = True return self + def set_fp8_current_scaling_gemm_plugins(self, dtype: str = "auto"): + self.fp8_current_scaling_gemm_plugin = dtype + return self + def set_context_fmha(self, context_fmha_type=ContextFMHAType.enabled): assert type(context_fmha_type) == ContextFMHAType self.context_fmha_type = context_fmha_type diff --git a/tensorrt_llm/quantization/functional.py b/tensorrt_llm/quantization/functional.py index 4a028ffe1077..275c750efc4b 100644 --- a/tensorrt_llm/quantization/functional.py +++ b/tensorrt_llm/quantization/functional.py @@ -212,6 +212,56 @@ def fp8_rowwise_gemm(input: Tensor, weights: Tensor, scales_a: Tensor, return _create_tensor(layer.get_output(0), layer) +def fp8_current_scaling_gemm( + input: Tensor, + weights: Tensor, + in_scales: Optional[Tensor] = None, + weight_scales: Optional[Tensor] = None, + need_quantize_acts_on_demand: bool = True, + need_quantize_weights_on_demand: bool = False) -> Tensor: + if not default_net().plugin_config.fp8_current_scaling_gemm_plugin: + raise TypeError( + "Fp8 current scaling GEMM is only supported with plugin") + else: + plg_creator = trt.get_plugin_registry().get_plugin_creator( + 'Fp8CurrentScalingGemm', '1', TRT_LLM_PLUGIN_NAMESPACE) + assert plg_creator is not None + + need_quantize_acts_on_demand = 1 if need_quantize_acts_on_demand else 0 + need_quantize_acts_on_demand = trt.PluginField( + "need_quantize_acts_on_demand", + np.array(need_quantize_acts_on_demand, dtype=np.int32), + trt.PluginFieldType.INT32) + + need_quantize_weights_on_demand = 1 if need_quantize_weights_on_demand else 0 + need_quantize_weights_on_demand = trt.PluginField( + "need_quantize_weights_on_demand", + np.array(need_quantize_weights_on_demand, dtype=np.int32), + trt.PluginFieldType.INT32) + + p_dtype = default_net().plugin_config.fp8_current_scaling_gemm_plugin + pf_type = trt.PluginField( + "type_id", np.array([int(str_dtype_to_trt(p_dtype))], np.int32), + trt.PluginFieldType.INT32) + + pfc = trt.PluginFieldCollection([ + need_quantize_acts_on_demand, need_quantize_weights_on_demand, + pf_type + ]) + fp8_cs_gemm_plug = plg_creator.create_plugin("fp8_current_scaling_gemm", + pfc) + plug_inputs = [input.trt_tensor, weights.trt_tensor] + if weight_scales is not None: + plug_inputs += [weight_scales.trt_tensor] + if in_scales is not None: + plug_inputs += [in_scales.trt_tensor] + layer = default_trtnet().add_plugin_v2(plug_inputs, fp8_cs_gemm_plug) + _add_plugin_info(layer, plg_creator, "fp8_current_scaling_gemm", pfc) + if not default_net().strongly_typed: + layer.get_input(1).set_dynamic_range(-448, 448) + return _create_tensor(layer.get_output(0), layer) + + def weight_only_quant_matmul(input: Tensor, weights: Tensor, scales: Tensor, diff --git a/tensorrt_llm/quantization/layers.py b/tensorrt_llm/quantization/layers.py index c91134602e39..2f4a581ca6a8 100644 --- a/tensorrt_llm/quantization/layers.py +++ b/tensorrt_llm/quantization/layers.py @@ -36,12 +36,13 @@ # isort: off from .functional import ( - dequantize, fp8_rowwise_gemm, fp8_rowwise_rms_norm, postprocess_fp8_rowwise, - postprocess_weight_only, postprocess_weight_only_groupwise, quantize, - quantize_fp8_per_token, quantize_per_token, quantize_tensor, - validate_group_size, smooth_quant_gemm, smooth_quant_layer_norm, - smooth_quant_rms_norm, weight_only_groupwise_quant_matmul, - weight_only_quant_matmul, qserve_gemm_per_group, qserve_gemm_per_channel) + dequantize, fp8_current_scaling_gemm, fp8_rowwise_gemm, + fp8_rowwise_rms_norm, postprocess_fp8_rowwise, postprocess_weight_only, + postprocess_weight_only_groupwise, quantize, quantize_fp8_per_token, + quantize_per_token, quantize_tensor, validate_group_size, smooth_quant_gemm, + smooth_quant_layer_norm, smooth_quant_rms_norm, + weight_only_groupwise_quant_matmul, weight_only_quant_matmul, + qserve_gemm_per_group, qserve_gemm_per_channel) # isort: on from .mode import QuantMode @@ -1620,6 +1621,177 @@ def postprocess(self, tllm_key, weights, **kwargs): return weights.to(str_dtype_to_torch(self.bias.dtype)) +class FP8CurrentScalingLinear(Linear): + + def __init__(self, + in_features, + out_features, + bias=False, + dtype=None, + tp_group=None, + tp_size=1, + gather_output=True, + prefer_managed_weight=True, + is_qkv=False, + quantize_weights_on_demand=False): + super().__init__(in_features, + out_features, + bias=bias, + dtype=dtype, + tp_group=tp_group, + tp_size=tp_size, + gather_output=gather_output, + prefer_managed_weight=prefer_managed_weight, + is_qkv=is_qkv) + self.quantize_weights_on_demand = quantize_weights_on_demand + if quantize_weights_on_demand: + self.weight = Parameter(shape=(self.out_features, self.in_features), + dtype=dtype, + prefer_managed=self.prefer_managed_weight) + else: + self.weight = Parameter(shape=(self.out_features, self.in_features), + dtype='fp8', + prefer_managed=self.prefer_managed_weight) + self.weights_scaling_factor = Parameter(shape=(math.ceil(self.out_features / 128), \ + math.ceil(self.in_features / 128)), dtype=trt.float32) + + # TODO: modify this for checkpoint loading + self.tllm_to_externel_key_dict = { + "weight": "weight", + } + + if not quantize_weights_on_demand: + self.tllm_to_externel_key_dict.update( + {"weights_scaling_factor": "weight_scale_inv"}) + + def forward(self, x, in_scales=None, lora_runtime_params=None): + weights_scaling_factor = None + w_in = self.weight.value + need_quantize_acts_on_demand = (x.dtype != trt.fp8) + need_quantize_weights_on_demand = self.quantize_weights_on_demand + if not need_quantize_acts_on_demand: + assert in_scales is not None, "FP8 input requires activation scales." + if not need_quantize_weights_on_demand: + weights_scaling_factor = self.weights_scaling_factor.value + + gemm_output = fp8_current_scaling_gemm(x, w_in, weights_scaling_factor, + in_scales, + need_quantize_acts_on_demand, + need_quantize_weights_on_demand) + + return self.collect_and_bias(gemm_output) + + def postprocess(self, tllm_key, weights, **kwargs): + # TODO: add checkpoint conversion logic + if tllm_key.endswith("scaling_factor"): + if isinstance(weights, list) and len(weights) == 2: + left = weights[0].to(torch.float32) + right = weights[1].to(torch.float32) + weights = torch.cat( + [left, right], + dim=0, + ) + return weights + else: + return weights.to(torch.float32) + elif tllm_key.endswith("weight"): + weight_dtype = str_dtype_to_torch( + self.dtype + ) if self.quantize_weights_on_demand else torch.float8_e4m3fn + if isinstance(weights, list) and len(weights) == 2: + left = weights[0].view(weight_dtype) + right = weights[1].view(weight_dtype) + weights = torch.cat( + [left, right], + dim=0, + ) + return weights + else: + return weights.view(weight_dtype) + elif tllm_key.endswith("bias"): + return weights.to(str_dtype_to_torch(self.bias.dtype)) + + +class FP8CurrentScalingRowLinear(RowLinear): + + def __init__(self, + in_features, + out_features, + bias=False, + dtype=None, + tp_group=None, + tp_size=1, + prefer_managed_weight=True, + is_expert=False, + quantize_weights_on_demand=False): + super().__init__(in_features, + out_features, + bias=bias, + dtype=dtype, + tp_group=tp_group, + tp_size=tp_size, + prefer_managed_weight=prefer_managed_weight, + is_expert=is_expert) + self.quantize_weights_on_demand = quantize_weights_on_demand + if quantize_weights_on_demand: + self.weight = Parameter( + shape=(self.out_features, self.in_features), + dtype=dtype, + prefer_managed=self.prefer_managed_weight, + ) + else: + self.weight = Parameter( + shape=(self.out_features, self.in_features), + dtype="fp8", + prefer_managed=self.prefer_managed_weight, + ) + self.weights_scaling_factor = Parameter(shape=(math.ceil(self.out_features / 128), \ + math.ceil(self.in_features / 128)), dtype=trt.float32) + + self.tllm_to_externel_key_dict = { + "weight": "weight", + } + + if not quantize_weights_on_demand: + self.tllm_to_externel_key_dict.update( + {"weights_scaling_factor": "weight_scale_inv"}) + + def forward(self, + x, + in_scales=None, + lora_runtime_params=None, + all_reduce_params=None): + + weights_scaling_factor = None + w_in = self.weight.value + + need_quantize_acts_on_demand = (x.dtype != trt.fp8) + need_quantize_weights_on_demand = self.quantize_weights_on_demand + if not need_quantize_acts_on_demand: + assert in_scales is not None, "FP8 input requires activation scales." + if not need_quantize_weights_on_demand: + weights_scaling_factor = self.weights_scaling_factor.value + gemm_output = fp8_current_scaling_gemm(x, w_in, weights_scaling_factor, + in_scales, + need_quantize_acts_on_demand, + need_quantize_weights_on_demand) + + return self.collect_and_bias(gemm_output, + all_reduce_params=all_reduce_params) + + def postprocess(self, tllm_key, weights, **kwargs): + # TODO: add checkpoint conversion logic + if tllm_key.endswith("scaling_factor"): + return weights.to(torch.float32) + elif tllm_key.endswith("weight"): + weight_dtype = str_dtype_to_torch( + self.dtype + ) if self.quantize_weights_on_demand else torch.float8_e4m3fn + return weights.view(weight_dtype) + elif tllm_key.endswith("bias"): + return weights.to(str_dtype_to_torch(self.bias.dtype)) + + class Fp8RowwiseMLP(Module): def __init__( diff --git a/tensorrt_llm/quantization/mode.py b/tensorrt_llm/quantization/mode.py index 07f830474125..a89d0138ae40 100644 --- a/tensorrt_llm/quantization/mode.py +++ b/tensorrt_llm/quantization/mode.py @@ -36,6 +36,7 @@ class QuantAlgo(StrEnum, metaclass=BaseEnumMeta): W4A8_QSERVE_PER_CHANNEL = auto() FP8 = auto() FP8_PER_CHANNEL_PER_TOKEN = auto() + FP8_CURRENT_SCALING = auto() INT8 = auto() MIXED_PRECISION = auto() NO_QUANT = auto() @@ -78,6 +79,10 @@ class QuantMode(IntFlag): FP8_QDQ = auto() # FP8 rowwise FP8_ROWWISE = auto() + # FP8 current scaling + FP8_1x128_128x128_CURRENT_SCALING = auto() + # FP8 current scaling + FP8_QUANTIZE_WEIGHTS_ON_DEMAND = auto() # W4A8 qserve W4A8_QSERVE = auto() @@ -130,6 +135,12 @@ def has_act_or_weight_quant(self): def has_per_token_dynamic_scaling(self): return self._any(self.PER_TOKEN) + def has_fp8_current_scaling(self): + return self._any(self.FP8_1x128_128x128_CURRENT_SCALING) + + def has_fp8_quantize_weights_on_demand(self): + return self._any(self.FP8_QUANTIZE_WEIGHTS_ON_DEMAND) + def has_act_static_scaling(self): return not self.has_per_token_dynamic_scaling( ) and not self.has_fp8_rowwise() @@ -162,7 +173,8 @@ def has_any_quant(self): return self._any(self.INT4_WEIGHTS | self.INT8_WEIGHTS | self.ACTIVATIONS | self.INT8_KV_CACHE | self.FP8_KV_CACHE - | self.FP8_QDQ | self.FP8_ROWWISE) + | self.FP8_QDQ | self.FP8_ROWWISE + | self.FP8_1x128_128x128_CURRENT_SCALING) def set_int8_kv_cache(self): return self | self.INT8_KV_CACHE @@ -186,6 +198,8 @@ def from_description(quantize_weights=False, use_int8_kv_cache=False, use_fp8_kv_cache=False, use_fp8_qdq=False, + use_fp8_current_scaling=False, + use_fp8_quantize_weight_on_demand=False, use_fp8_rowwise=False, use_w4a8_qserve=False): @@ -196,12 +210,14 @@ def raise_error(): f"{per_token=}, " f"{per_channel=}, " f"{per_group=}, " - f"{use_int4_weights=}" - f"{use_int8_kv_cache=}" - f"{use_fp8_kv_cache=}" - f"{use_fp8_qdq=}" - f"{use_fp8_rowwise=}" - f"{use_w4a8_qserve=}") + f"{use_int4_weights=}, " + f"{use_int8_kv_cache=}, " + f"{use_fp8_kv_cache=}, " + f"{use_fp8_qdq=}, " + f"{use_fp8_current_scaling=}, " + f"{use_fp8_quantize_weight_on_demand=}," + f"{use_fp8_rowwise=}, " + f"{use_w4a8_qserve=}, ") # We must quantize weights when we quantize activations. if quantize_activations and not quantize_weights: @@ -246,6 +262,12 @@ def raise_error(): if use_fp8_rowwise: mode = mode | QuantMode.FP8_ROWWISE | QuantMode.PER_TOKEN | QuantMode.PER_CHANNEL + if use_fp8_current_scaling: + mode = mode | QuantMode.FP8_1x128_128x128_CURRENT_SCALING + + if use_fp8_quantize_weight_on_demand: + mode = mode | QuantMode.FP8_QUANTIZE_WEIGHTS_ON_DEMAND + # W4A8 QServe if use_w4a8_qserve: mode = mode | QuantMode.W4A8_QSERVE @@ -319,6 +341,10 @@ def from_quant_algo( quant_mode = QuantMode.from_description(use_fp8_qdq=True) elif quant_algo == QuantAlgo.FP8_PER_CHANNEL_PER_TOKEN: quant_mode = QuantMode.from_description(use_fp8_rowwise=True) + elif quant_algo == QuantAlgo.FP8_CURRENT_SCALING: + quant_mode = QuantMode.from_description( + use_fp8_current_scaling=True, + use_fp8_quantize_weight_on_demand=False) else: quant_mode = QuantMode(0) @@ -345,6 +371,8 @@ def to_dict(self): self.has_fp8_qdq(), 'enable_fp8_rowwise': self.has_fp8_rowwise(), + 'enable_fp8_current_scaling': + self.has_fp8_current_scaling(), 'fp8_kv_cache': self.has_fp8_kv_cache(), 'use_weight_only': diff --git a/tensorrt_llm/quantization/quantize.py b/tensorrt_llm/quantization/quantize.py index 17efc23a9b00..ba7636867f10 100644 --- a/tensorrt_llm/quantization/quantize.py +++ b/tensorrt_llm/quantization/quantize.py @@ -9,16 +9,20 @@ from ..layers.moe import MixtureOfExperts from ..models.modeling_utils import LayerQuantConfig, QuantConfig from ..parameter import Parameter -from .layers import (FP8Linear, FP8RowLinear, Fp8RowwiseAttention, - Fp8RowwiseGatedMLP, Fp8RowwiseMLP, Fp8RowwiseRmsNorm, - Int8SmoothQuantLinear, Int8SmoothQuantRowLinear, - QServeAttention, QServeGatedMLP, QServeMLP, QServeRmsNorm, - SmoothQuantAttention, SmoothQuantGatedMLP, - SmoothQuantLayerNorm, SmoothQuantMLP, SmoothQuantRmsNorm, - WeightOnlyGroupwiseQuantColumnLinear, - WeightOnlyGroupwiseQuantRowLinear, - WeightOnlyQuantColumnLinear, WeightOnlyQuantEmbedding, - WeightOnlyQuantRowLinear) + +# isort: off +from .layers import ( + FP8CurrentScalingLinear, FP8CurrentScalingRowLinear, FP8Linear, + FP8RowLinear, FP8RowLinear, Fp8RowwiseAttention, Fp8RowwiseGatedMLP, + Fp8RowwiseMLP, Fp8RowwiseRmsNorm, Int8SmoothQuantLinear, + Int8SmoothQuantRowLinear, QServeAttention, QServeGatedMLP, QServeMLP, + QServeRmsNorm, SmoothQuantAttention, SmoothQuantGatedMLP, + SmoothQuantLayerNorm, SmoothQuantMLP, SmoothQuantRmsNorm, + WeightOnlyGroupwiseQuantColumnLinear, WeightOnlyGroupwiseQuantRowLinear, + WeightOnlyQuantColumnLinear, WeightOnlyQuantEmbedding, + WeightOnlyQuantRowLinear) +# isort: on + from .mode import W8A8_SQ_PLUGIN_LIST, QuantAlgo, QuantMode @@ -35,7 +39,6 @@ def quantize_layers( '*position_embedding', '*block_embedding', '*shared_expert_gate', - '*fused_a', ] for name, module, parent in model.named_modules_with_parent(): @@ -71,8 +74,12 @@ def quantize_layers( if isinstance(module, ColumnLinear): init_params[ "out_features"] = module.out_features * module.tp_size + if quant_config.quant_mode.has_fp8_quantize_weights_on_demand(): + init_params["quantize_weights_on_demand"] = True elif isinstance(module, RowLinear): init_params["in_features"] = module.in_features * module.tp_size + if quant_config.quant_mode.has_fp8_quantize_weights_on_demand(): + init_params["quantize_weights_on_demand"] = True if preprocess_init_params is not None: preprocess_init_params(init_params, name, module) quant_layer = quant_cls(**init_params) @@ -227,6 +234,22 @@ def fp8_quantize(model, quant_config: QuantConfig): return model +def fp8_current_scaling_quantize(model, quant_config: QuantConfig): + assert quant_config.quant_mode.has_fp8_current_scaling() + + quant_map = { + ColumnLinear: FP8CurrentScalingLinear, + RowLinear: FP8CurrentScalingRowLinear, + } + + model = quantize_layers( + model, + quant_config, + quant_map, + ) + return model + + def fp8_rowwise_quantize(model, quant_config: QuantConfig): assert quant_config.quant_mode.has_fp8_rowwise() @@ -542,6 +565,8 @@ def quantize(model, quant_config: Union[QuantConfig, LayerQuantConfig]): if layer_quant_mode.has_fp8_qdq(): module = fp8_quantize(module, layer_quant_cfg) + elif layer_quant_mode.has_fp8_current_scaling(): + module = fp8_current_scaling_quantize(module, layer_quant_cfg) elif layer_quant_mode.has_fp8_rowwise(): module = fp8_rowwise_quantize(module, layer_quant_cfg) elif layer_quant_mode.is_qserve_w4a8(): diff --git a/tests/attention/test_deepseek_v2_attention.py b/tests/attention/test_deepseek_v2_attention.py index 7fc723ca9d4d..56cd2c54ec64 100644 --- a/tests/attention/test_deepseek_v2_attention.py +++ b/tests/attention/test_deepseek_v2_attention.py @@ -1138,7 +1138,7 @@ def _construct_execution( # position_ids=position_ids_tensor, # q_a_proj=q_a_proj_tensor, # q_a_layernorm=q_a_layernorm_tensor, - fused_q_proj=kv_a_proj_with_mqa_tensor, + k_b_proj_trans=kv_a_proj_with_mqa_tensor, q_b_proj=q_b_proj_tensor, # kv_a_proj_with_mqa=kv_a_proj_with_mqa_tensor, # kv_a_layernorm=kv_a_layernorm_tensor,