[Common] Split cast/gated kernels by scaling mode #2248
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,97 @@ | ||
| /************************************************************************* | ||
| * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. | ||
| * | ||
| * See LICENSE for license information. | ||
| ************************************************************************/ | ||
|
|
||
| /*! \file common.cuh | ||
| * \brief Common functions in quantize. | ||
|
Contributor
There was a problem hiding this comment. style: file comment says "in quantize" but this lives in |
||
| */ | ||
|
|
||
| #ifndef TRANSFORMER_ENGINE_QUANTIZE_CORE_COMMON_CUH_ | ||
|
Contributor
There was a problem hiding this comment. syntax: header guard mismatch – file is |
||
| #define TRANSFORMER_ENGINE_QUANTIZE_CORE_COMMON_CUH_ | ||
|
|
||
| #include <cuda.h> | ||
| #include <cudaTypedefs.h> | ||
| #include <cuda_runtime.h> | ||
| #include <transformer_engine/transformer_engine.h> | ||
|
|
||
| #include "../../common.h" | ||
| #include "../../utils.cuh" | ||
|
|
||
| namespace transformer_engine { | ||
| namespace dispatch { | ||
| namespace common { | ||
| inline bool full_tile_1D_tensor(const Tensor *const t, const size_t elems_per_block) { | ||
| const size_t N = product(t->data.shape); | ||
| const bool isFullTile = (N % elems_per_block == 0); | ||
| return isFullTile; | ||
| } | ||
|
|
||
| inline bool dimensions_supported_by_TMA(const Tensor *const t) { | ||
| const size_t cols = t->flat_last_dim(); | ||
| constexpr size_t TMA_bytes = 16; | ||
| const size_t alignment_requirement = (TMA_bytes * 8) / typeToNumBits(t->dtype()); | ||
| return cols % alignment_requirement == 0; | ||
| } | ||
|
|
||
| namespace kernel { | ||
|
|
||
| constexpr size_t THREADS_PER_BLOCK = 256; | ||
| template <int nvec, typename OType> | ||
| __global__ void __launch_bounds__(THREADS_PER_BLOCK) | ||
| reduce_dbias_kernel(OType *const dbias_output, const float *const dbias_partial, | ||
| const size_t rows, const size_t cols) { | ||
| using ComputeVec = Vec<float, nvec>; | ||
| using OutputVec = Vec<OType, nvec>; | ||
|
|
||
| const size_t thread_id = blockIdx.x * blockDim.x + threadIdx.x; | ||
|
|
||
| if (thread_id * nvec >= cols) { | ||
| return; | ||
| } | ||
|
|
||
| const float *const thread_in_base = dbias_partial + thread_id * nvec; | ||
| OType *const thread_out_base = dbias_output + thread_id * nvec; | ||
|
|
||
| ComputeVec ldg_vec; | ||
| ComputeVec acc_vec; | ||
| acc_vec.clear(); | ||
| for (int i = 0; i < rows; ++i) { | ||
| ldg_vec.load_from(thread_in_base + i * cols); | ||
| #pragma unroll | ||
| for (int e = 0; e < nvec; ++e) { | ||
| acc_vec.data.elt[e] += ldg_vec.data.elt[e]; | ||
| } | ||
| } | ||
|
|
||
| OutputVec stg_vec; | ||
| #pragma unroll | ||
| for (int e = 0; e < nvec; ++e) { | ||
| stg_vec.data.elt[e] = static_cast<OType>(acc_vec.data.elt[e]); | ||
| } | ||
| stg_vec.store_to(thread_out_base); | ||
| } | ||
| } // namespace kernel | ||
|
|
||
| template <typename IType> | ||
| void reduce_dbias(const float *workspace_ptr, Tensor *dbias, const size_t rows, const size_t cols, | ||
|
Comment on lines
+77
to
+78
Contributor
There was a problem hiding this comment. style: template parameter |
||
| cudaStream_t stream) { | ||
| using namespace kernel; | ||
| constexpr size_t reduce_dbias_store_bytes = 8; // stg.64 | ||
| constexpr size_t reduce_dbias_nvec = reduce_dbias_store_bytes / sizeof(IType); | ||
|
|
||
| NVTE_CHECK(cols % reduce_dbias_nvec == 0, "Unsupported shape."); | ||
| const size_t reduce_dbias_num_blocks = DIVUP(cols, THREADS_PER_BLOCK * reduce_dbias_nvec); | ||
|
|
||
| reduce_dbias_kernel<reduce_dbias_nvec, IType> | ||
| <<<reduce_dbias_num_blocks, THREADS_PER_BLOCK, 0, stream>>>( | ||
| reinterpret_cast<IType *>(dbias->data.dptr), workspace_ptr, rows, cols); | ||
| NVTE_CHECK_CUDA(cudaGetLastError()); | ||
| } | ||
|
|
||
| } // namespace common | ||
| } // namespace dispatch | ||
| } // namespace transformer_engine | ||
|
|
||
| #endif // TRANSFORMER_ENGINE_QUANTIZE_CORE_COMMON_CUH_ | ||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,56 @@ | ||
| /************************************************************************* | ||
| * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. | ||
| * | ||
| * See LICENSE for license information. | ||
| ************************************************************************/ | ||
|
|
||
| /*! \file dequantize.cuh | ||
| * \brief Dequantize dispatcher. | ||
| */ | ||
|
|
||
| #ifndef TRANSFORMER_ENGINE_DISPATCH_DEQUANTIZE_CUH_ | ||
| #define TRANSFORMER_ENGINE_DISPATCH_DEQUANTIZE_CUH_ | ||
|
|
||
| #include <transformer_engine/transformer_engine.h> | ||
|
|
||
| #include "../../common.h" | ||
| #include "../fp8/dequantize_fp8.cuh" | ||
| #include "../mxfp8/dequantize_mxfp8.cuh" | ||
| #include "../nvfp4/dequantize_nvfp4.cuh" | ||
|
|
||
| namespace transformer_engine { | ||
| namespace dispatch { | ||
|
|
||
| inline void dequantize_helper(const Tensor &input, Tensor *output, cudaStream_t stream) { | ||
| CheckInputTensor(input, "cast_input"); | ||
| CheckOutputTensor(*output, "cast_output"); | ||
|
|
||
| switch (input.scaling_mode) { | ||
| case NVTE_DELAYED_TENSOR_SCALING: { | ||
| NVTE_CHECK(is_fp8_dtype(input.data.dtype), "Input must have FP8 type."); | ||
| NVTE_CHECK(!is_fp8_dtype(output->data.dtype), "Output must be in higher precision."); | ||
| NVTE_CHECK(output->data.shape == input.data.shape, "Input and output shapes need to match."); | ||
| fp8::dequantize(input, output, stream); | ||
Oleg-Goncharov marked this conversation as resolved.
Show resolved
Hide resolved
|
||
| break; | ||
| } | ||
| case NVTE_MXFP8_1D_SCALING: { | ||
| if (is_supported_by_CC_100()) { | ||
| mxfp8::dequantize(input, output, stream); | ||
| } else { | ||
| NVTE_ERROR("MXFP8 Dequantization is NOT supported by architectures < 10.0"); | ||
| } | ||
| break; | ||
| } | ||
| case NVTE_NVFP4_1D_SCALING: { | ||
| nvfp4::dequantize(input, output, stream); | ||
| break; | ||
| } | ||
| default: | ||
| NVTE_ERROR("Not implemented scaling mode: " + to_string(input.scaling_mode) + "."); | ||
| } | ||
| } | ||
|
|
||
| } // namespace dispatch | ||
| } // namespace transformer_engine | ||
|
|
||
| #endif // TRANSFORMER_ENGINE_DISPATCH_DEQUANTIZE_CUH_ | ||
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
style: indentation of the continuation line changed from 4 to 1 space, which deviates from the project's
.clang-format(ContinuationIndentWidth: 4)