Blockwise and Groupwise GEMM for Blackwell and Improvements for Hopper (#2139)
- Blockwise and Groupwise GEMM improvements for Hopper. - Blockwise and Groupwise GEMM for Blackwell. - Blockwise Grouped GEMM for Hopper. - Static ScalePromotionInterval for Hopper FP8 GEMMs. Co-authored-by: dePaul Miller <23461061+depaulmillz@users.noreply.github.com>
This commit is contained in:
dePaul Miller
189
include/cutlass/detail/sm100_blockwise_scale_layout.hpp
Normal file
189
include/cutlass/detail/sm100_blockwise_scale_layout.hpp
Normal file
@@ -0,0 +1,189 @@
|
||||
/***************************************************************************************************
|
||||
* Copyright (c) 2025 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
* SPDX-License-Identifier: BSD-3-Clause
|
||||
*
|
||||
* Redistribution and use in source and binary forms, with or without
|
||||
* modification, are permitted provided that the following conditions are met:
|
||||
*
|
||||
* 1. Redistributions of source code must retain the above copyright notice, this
|
||||
* list of conditions and the following disclaimer.
|
||||
*
|
||||
* 2. Redistributions in binary form must reproduce the above copyright notice,
|
||||
* this list of conditions and the following disclaimer in the documentation
|
||||
* and/or other materials provided with the distribution.
|
||||
*
|
||||
* 3. Neither the name of the copyright holder nor the names of its
|
||||
* contributors may be used to endorse or promote products derived from
|
||||
* this software without specific prior written permission.
|
||||
*
|
||||
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
* AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
* IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
||||
* DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
||||
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
||||
* DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
||||
* SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
||||
* CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
||||
* OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
*
|
||||
**************************************************************************************************/
|
||||
|
||||
|
||||
|
||||
/*! \file
|
||||
\brief Block Wise Scale configs specific for SM100 Blockwise/Groupwise MMA
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "cutlass/layout/matrix.h"
|
||||
|
||||
#include "cute/int_tuple.hpp"
|
||||
#include "cute/atom/mma_traits_sm100.hpp"
|
||||
|
||||
namespace cutlass::detail{
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
using namespace cute;
|
||||
|
||||
template<int SFVecSizeM, int SFVecSizeN, int SFVecSizeK, UMMA::Major majorSFA = UMMA::Major::MN, UMMA::Major majorSFB = UMMA::Major::MN>
|
||||
struct Sm100BlockwiseScaleConfig {
|
||||
|
||||
using ShapeSFA = Shape<Shape<Int<SFVecSizeM>, int32_t>, Shape<Int<SFVecSizeK>, int32_t>, int32_t>;
|
||||
using ShapeSFB = Shape<Shape<Int<SFVecSizeN>, int32_t>, Shape<Int<SFVecSizeK>, int32_t>, int32_t>;
|
||||
|
||||
using StrideSFA = conditional_t<majorSFA == UMMA::Major::MN,
|
||||
Stride<Stride<_0,_1>,Stride<_0,int32_t>, int32_t>,
|
||||
Stride<Stride<_0,int32_t>,Stride<_0,_1>, int32_t>>;
|
||||
|
||||
using StrideSFB = conditional_t<majorSFB == UMMA::Major::MN,
|
||||
Stride<Stride<_0,_1>,Stride<_0,int32_t>, int32_t>,
|
||||
Stride<Stride<_0,int32_t>,Stride<_0,_1>, int32_t>>;
|
||||
|
||||
using LayoutSFA = Layout<ShapeSFA, StrideSFA>;
|
||||
using LayoutSFB = Layout<ShapeSFB, StrideSFB>;
|
||||
|
||||
CUTE_HOST_DEVICE
|
||||
static constexpr auto
|
||||
deduce_layoutSFA() {
|
||||
return LayoutSFA{};
|
||||
}
|
||||
|
||||
template<typename CtaShape_MNK>
|
||||
CUTE_HOST_DEVICE
|
||||
static constexpr auto
|
||||
smem_atom_layoutSFA(CtaShape_MNK cta_shape_mnk) {
|
||||
static_assert(cute::is_static_v<CtaShape_MNK>, "Expect static CTA shape");
|
||||
auto strides = [&]() CUTLASS_LAMBDA_FUNC_INLINE {
|
||||
auto [M, N, K] = cta_shape_mnk;
|
||||
if constexpr (majorSFA == UMMA::Major::MN) {
|
||||
return make_stride(make_stride(_0{}, _1{}), make_stride(_0{}, Int<cute::ceil_div(size<0>(CtaShape_MNK{}), SFVecSizeM)>{}));
|
||||
}
|
||||
else {
|
||||
return make_stride(make_stride(_0{}, Int<cute::ceil_div(size<2>(CtaShape_MNK{}), SFVecSizeK)>{}), make_stride(_0{}, _1{}));
|
||||
}
|
||||
}();
|
||||
|
||||
auto [M, N, K] = cta_shape_mnk;
|
||||
return make_layout(
|
||||
make_shape(make_shape(Int<SFVecSizeM>{}, Int<cute::ceil_div(size<0>(CtaShape_MNK{}), SFVecSizeM)>{}),
|
||||
make_shape(Int<SFVecSizeK>{}, Int<cute::ceil_div(size<2>(CtaShape_MNK{}), SFVecSizeK)>{})),
|
||||
strides
|
||||
);
|
||||
}
|
||||
|
||||
|
||||
CUTE_HOST_DEVICE
|
||||
static constexpr auto
|
||||
deduce_layoutSFB() {
|
||||
return LayoutSFB{};
|
||||
}
|
||||
|
||||
template<typename CtaShape_MNK>
|
||||
CUTE_HOST_DEVICE
|
||||
static constexpr auto
|
||||
smem_atom_layoutSFB(CtaShape_MNK cta_shape_mnk) {
|
||||
static_assert(cute::is_static_v<CtaShape_MNK>, "Expect static CTA shape");
|
||||
auto strides = [&]() CUTLASS_LAMBDA_FUNC_INLINE {
|
||||
if constexpr (majorSFA == UMMA::Major::MN) {
|
||||
return make_stride(make_stride(_0{}, _1{}), make_stride(_0{}, Int<cute::ceil_div(size<1>(CtaShape_MNK{}), SFVecSizeN)>{}));
|
||||
}
|
||||
else {
|
||||
return make_stride(make_stride(_0{}, Int<cute::ceil_div(size<2>(CtaShape_MNK{}), SFVecSizeK)>{}), make_stride(_0{}, _1{}));
|
||||
}
|
||||
}();
|
||||
|
||||
auto [M, N, K] = cta_shape_mnk;
|
||||
return make_layout(
|
||||
make_shape(make_shape(Int<SFVecSizeN>{}, Int<cute::ceil_div(size<1>(CtaShape_MNK{}), SFVecSizeN)>{}),
|
||||
make_shape(Int<SFVecSizeK>{}, Int<cute::ceil_div(size<2>(CtaShape_MNK{}), SFVecSizeK)>{})),
|
||||
strides
|
||||
);
|
||||
}
|
||||
|
||||
// The following function is provided for user fill dynamic problem size to the layout_SFA.
|
||||
template <class ProblemShape>
|
||||
CUTE_HOST_DEVICE
|
||||
static constexpr auto
|
||||
tile_atom_to_shape_SFA(ProblemShape problem_shape) {
|
||||
auto problem_shape_MNKL = append<4>(problem_shape, 1);
|
||||
|
||||
auto strides = [&]() CUTLASS_LAMBDA_FUNC_INLINE {
|
||||
auto [M, N, K, L] = problem_shape_MNKL;
|
||||
if constexpr (majorSFA == UMMA::Major::MN) {
|
||||
return make_stride(make_stride(_0{}, _1{}), make_stride(_0{}, cute::ceil_div(M, SFVecSizeM)));
|
||||
}
|
||||
else {
|
||||
return make_stride(make_stride(_0{}, cute::ceil_div(K, SFVecSizeK)), make_stride(_0{}, _1{}));
|
||||
}
|
||||
}();
|
||||
|
||||
auto [M, N, K, L] = problem_shape_MNKL;
|
||||
auto mk_layout = make_layout(
|
||||
make_shape(make_shape(Int<SFVecSizeM>{}, cute::ceil_div(M, SFVecSizeM)),
|
||||
make_shape(Int<SFVecSizeK>{}, cute::ceil_div(K, SFVecSizeK))),
|
||||
strides
|
||||
);
|
||||
|
||||
return make_layout(append(shape(mk_layout), L), append(stride(mk_layout), size(filter_zeros(mk_layout))));
|
||||
}
|
||||
|
||||
// The following function is provided for user fill dynamic problem size to the layout_SFB.
|
||||
template <class ProblemShape>
|
||||
CUTE_HOST_DEVICE
|
||||
static constexpr auto
|
||||
tile_atom_to_shape_SFB(ProblemShape problem_shape) {
|
||||
auto problem_shape_MNKL = append<4>(problem_shape, 1);
|
||||
|
||||
auto strides = [&]() CUTLASS_LAMBDA_FUNC_INLINE {
|
||||
auto [M, N, K, L] = problem_shape_MNKL;
|
||||
|
||||
if constexpr (majorSFB == UMMA::Major::MN) {
|
||||
return make_stride(make_stride(_0{}, _1{}), make_stride(_0{}, cute::ceil_div(N, SFVecSizeN)));
|
||||
}
|
||||
else {
|
||||
return make_stride(make_stride(_0{}, cute::ceil_div(K, SFVecSizeK)), make_stride(_0{}, _1{}));
|
||||
}
|
||||
}();
|
||||
|
||||
auto [M, N, K, L] = problem_shape_MNKL;
|
||||
auto nk_layout = make_layout(
|
||||
make_shape(make_shape(Int<SFVecSizeN>{}, cute::ceil_div(N, SFVecSizeN)),
|
||||
make_shape(Int<SFVecSizeK>{}, cute::ceil_div(K, SFVecSizeK))),
|
||||
strides
|
||||
);
|
||||
|
||||
return make_layout(append(shape(nk_layout), L), append(stride(nk_layout), size(filter_zeros(nk_layout))));
|
||||
}
|
||||
|
||||
};
|
||||
|
||||
template<class MmaTileShape_MNK>
|
||||
constexpr auto sm100_trivial_blockwise_scale_config(MmaTileShape_MNK) {
|
||||
return Sm100BlockwiseScaleConfig<size<0>(MmaTileShape_MNK{}), size<1>(MmaTileShape_MNK{}), size<2>(MmaTileShape_MNK{})>{};
|
||||
}
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace cutlass::detail
|
||||
@@ -0,0 +1,304 @@
|
||||
/***************************************************************************************************
|
||||
* Copyright (c) 2025 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
* SPDX-License-Identifier: BSD-3-Clause
|
||||
*
|
||||
* Redistribution and use in source and binary forms, with or without
|
||||
* modification, are permitted provided that the following conditions are met:
|
||||
*
|
||||
* 1. Redistributions of source code must retain the above copyright notice, this
|
||||
* list of conditions and the following disclaimer.
|
||||
*
|
||||
* 2. Redistributions in binary form must reproduce the above copyright notice,
|
||||
* this list of conditions and the following disclaimer in the documentation
|
||||
* and/or other materials provided with the distribution.
|
||||
*
|
||||
* 3. Neither the name of the copyright holder nor the names of its
|
||||
* contributors may be used to endorse or promote products derived from
|
||||
* this software without specific prior written permission.
|
||||
*
|
||||
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
* AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
* IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
||||
* DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
||||
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
||||
* DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
||||
* SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
||||
* CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
||||
* OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
*
|
||||
**************************************************************************************************/
|
||||
#pragma once
|
||||
|
||||
#include "cutlass/gemm/collective/builders/sm100_common.inl"
|
||||
#include "cutlass/gemm/collective/builders/sm100_pipeline_carveout.inl"
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace cutlass::gemm::collective {
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace detail {
|
||||
|
||||
// Returns the maximum number of smem tiles that can be used with a given smem capacity, or overrides with manual count.
|
||||
template<
|
||||
int CapacityBytes,
|
||||
class ElementA,
|
||||
class ElementB,
|
||||
class ElementScalar,
|
||||
class ScaleShapeMNK,
|
||||
class TileShapeMNK,
|
||||
class MainloopPipelineStorage,
|
||||
class TransformLoadPipelineStorage,
|
||||
class TransformPipelineStorage,
|
||||
int stages
|
||||
>
|
||||
constexpr int
|
||||
sm100_compute_stage_count_or_override_blockwise(StageCount<stages> stage_count) {
|
||||
return stages;
|
||||
}
|
||||
|
||||
// Returns the maximum number of smem tiles that can be used with a given smem capacity, or overrides with manual count.
|
||||
template<
|
||||
int CapacityBytes,
|
||||
class ElementA,
|
||||
class ElementB,
|
||||
class ElementScalar,
|
||||
class ScaleShapeMNK,
|
||||
class TileShapeMNK,
|
||||
class MainloopPipelineStorage,
|
||||
class TransformLoadPipelineStorage,
|
||||
class TransformPipelineStorage,
|
||||
int stages
|
||||
>
|
||||
constexpr int
|
||||
sm100_compute_stage_count_or_override_blockwise(cute::Int<stages> stage_count) {
|
||||
return stages;
|
||||
}
|
||||
|
||||
// Returns the maximum number of smem tiles that can be used with a given smem capacity, or overrides with manual count.
|
||||
template<
|
||||
int CapacityBytes,
|
||||
class ElementA,
|
||||
class ElementB,
|
||||
class ElementScalar,
|
||||
class ScaleShapeMNK,
|
||||
class TileShapeMNK,
|
||||
class MainloopPipelineStorage,
|
||||
class TransformLoadPipelineStorage,
|
||||
class TransformPipelineStorage,
|
||||
int carveout_bytes>
|
||||
constexpr int
|
||||
sm100_compute_stage_count_or_override_blockwise(StageCountAutoCarveout<carveout_bytes> stage_count) {
|
||||
// For F8/F6/F4 sub-bytes, ElementA/B will be passed in as uint8_t
|
||||
// For Planar Complex, ElementA/B will be passed in as cutlass::complex<ElementARaw>
|
||||
// Each stage include (CollectiveMma::SharedStorage)
|
||||
// 1. smem for A and smem for B (CollectiveMma::SharedStorage::TensorStorage)
|
||||
// 2. one of each of the pipelines
|
||||
constexpr auto pipeline_bytes = sizeof(MainloopPipelineStorage) +
|
||||
sizeof(TransformLoadPipelineStorage) + sizeof(TransformPipelineStorage);
|
||||
|
||||
constexpr auto a_bits = cute::sizeof_bits_v<ElementA>;
|
||||
constexpr auto b_bits = cute::sizeof_bits_v<ElementB>;
|
||||
constexpr auto scale_bits = cute::sizeof_bits_v<ElementScalar>;
|
||||
|
||||
constexpr int stage_bytes =
|
||||
cutlass::bits_to_bytes(a_bits * size<0>(TileShapeMNK{}) * size<2>(TileShapeMNK{})) +
|
||||
cutlass::bits_to_bytes(b_bits * size<1>(TileShapeMNK{}) * size<2>(TileShapeMNK{})) +
|
||||
cutlass::bits_to_bytes(scale_bits * size<0>(ScaleShapeMNK{}) * size<2>(ScaleShapeMNK{})) +
|
||||
cutlass::bits_to_bytes(scale_bits * size<1>(ScaleShapeMNK{}) * size<2>(ScaleShapeMNK{})) +
|
||||
static_cast<int>(pipeline_bytes);
|
||||
|
||||
return (CapacityBytes - carveout_bytes) / stage_bytes;
|
||||
}
|
||||
|
||||
} // namespace detail
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <
|
||||
class ElementA,
|
||||
class GmemLayoutATagPair,
|
||||
int AlignmentA,
|
||||
class ElementB,
|
||||
class GmemLayoutBTagPair,
|
||||
int AlignmentB,
|
||||
class ElementAccumulator,
|
||||
class TileShape_MNK,
|
||||
class ClusterShape_MNK,
|
||||
class StageCountType,
|
||||
class KernelScheduleType
|
||||
>
|
||||
struct CollectiveBuilder<
|
||||
arch::Sm100,
|
||||
arch::OpClassTensorOp,
|
||||
ElementA,
|
||||
GmemLayoutATagPair,
|
||||
AlignmentA,
|
||||
ElementB,
|
||||
GmemLayoutBTagPair,
|
||||
AlignmentB,
|
||||
ElementAccumulator,
|
||||
TileShape_MNK, // (MmaAtomShapeM, MmaAtomShapeN, TileK)
|
||||
ClusterShape_MNK, // Static cluster shape or dynamic (int, int, _1)
|
||||
StageCountType,
|
||||
KernelScheduleType,
|
||||
cute::enable_if_t<
|
||||
not cute::is_tuple_v<ElementA> && not cute::is_tuple_v<ElementB> &&
|
||||
not cute::is_complex_v<ElementA> && not cute::is_complex_v<ElementB> &&
|
||||
cute::is_tuple_v<GmemLayoutATagPair> && cute::is_tuple_v<GmemLayoutBTagPair> &&
|
||||
// Dense Gemm
|
||||
cute::is_base_of_v<KernelScheduleSm100Blockwise, KernelScheduleType> &&
|
||||
// Alignment check
|
||||
detail::sm1xx_gemm_is_aligned<ElementA, AlignmentA, ElementB, AlignmentB, KernelScheduleType>()>>
|
||||
{
|
||||
static_assert(cute::is_static_v<TileShape_MNK>, "TileShape has to be static");
|
||||
static_assert(detail::check_input_datatypes<ElementA, ElementB>(), "Incorrect input types");
|
||||
|
||||
using GmemLayoutATag = cute::remove_cvref_t<decltype(get<0>(GmemLayoutATagPair{}))>;
|
||||
using GmemLayoutSFATag = cute::remove_cvref_t<decltype(get<1>(GmemLayoutATagPair{}))>;
|
||||
using GmemLayoutBTag = cute::remove_cvref_t<decltype(get<0>(GmemLayoutBTagPair{}))>;
|
||||
using GmemLayoutSFBTag = cute::remove_cvref_t<decltype(get<1>(GmemLayoutBTagPair{}))>;
|
||||
|
||||
static_assert(cute::depth(GmemLayoutSFATag{}) == 2 and cute::depth(GmemLayoutSFBTag{}) == 2,
|
||||
"Expect SFA and SFB layout to be depth of two with shape ((SFVecMN, restMN),(SFVecK, restK), L)");
|
||||
static_assert(size<1,0>(GmemLayoutSFATag{}) == size<1, 0>(GmemLayoutSFBTag{}),
|
||||
"SFA and SFB must have equivalent SF vector sizes along K");
|
||||
|
||||
static constexpr cute::UMMA::Major UmmaMajorA = cutlass::gemm::collective::detail::tag_to_umma_major_A<GmemLayoutATag>();
|
||||
static constexpr cute::UMMA::Major UmmaMajorB = cutlass::gemm::collective::detail::tag_to_umma_major_B<GmemLayoutBTag>();
|
||||
|
||||
// Data type used by MMA instruction
|
||||
using ElementAMma = decltype(cutlass::gemm::collective::detail::sm100_kernel_input_element_to_mma_input_element<ElementA>());
|
||||
using ElementBMma = decltype(cutlass::gemm::collective::detail::sm100_kernel_input_element_to_mma_input_element<ElementB>());
|
||||
|
||||
static constexpr bool is_2sm = cute::is_base_of_v<KernelSchedule2Sm, KernelScheduleType> ||
|
||||
(not cute::is_base_of_v<KernelSchedule1Sm, KernelScheduleType> &&
|
||||
not cute::is_base_of_v<KernelSchedule2Sm, KernelScheduleType> &&
|
||||
cute::is_static_v<ClusterShape_MNK> &&
|
||||
cute::get<0>(ClusterShape_MNK{}) % 2 == 0 );
|
||||
|
||||
static_assert(detail::sm100_gemm_check_for_f8f6f4_mix8bit_requirement<ElementAMma, ElementBMma,
|
||||
TileShape_MNK, ClusterShape_MNK,
|
||||
UmmaMajorA, UmmaMajorB, KernelScheduleType, is_2sm>(),
|
||||
"TileSize and MNK Major does not met with MMA Mix 8-bit TMA load requirement" );
|
||||
using TiledMma = decltype(detail::sm100_make_trivial_tiled_mma<
|
||||
ElementAMma, ElementBMma, ElementAccumulator,
|
||||
decltype(cute::product_each(TileShape_MNK{})), ClusterShape_MNK,
|
||||
UmmaMajorA, UmmaMajorB, KernelScheduleType>());
|
||||
|
||||
using ElementAMma_SmemAllocType = cute::conditional_t<cute::sizeof_bits_v<ElementAMma> < 8, uint8_t, ElementAMma>;
|
||||
using ElementBMma_SmemAllocType = cute::conditional_t<cute::sizeof_bits_v<ElementBMma> < 8, uint8_t, ElementBMma>;
|
||||
|
||||
using AtomThrID = typename TiledMma::AtomThrID;
|
||||
|
||||
using AtomThrShapeMNK = cute::Shape<decltype(cute::shape<0>(typename TiledMma::ThrLayoutVMNK{})), _1, _1>;
|
||||
using CtaTileShape_MNK = decltype(cute::shape_div(TileShape_MNK{}, AtomThrShapeMNK{}));
|
||||
|
||||
// ((MMA_TILE_M,MMA_TILE_K), MMA_M, MMA_K)
|
||||
using MmaShapeA_MK = decltype(partition_shape_A(TiledMma{}, make_shape(cute::size<0>(TileShape_MNK{}),
|
||||
cute::size<2>(TileShape_MNK{}))));
|
||||
// ((MMA_TILE_N,MMA_TILE_K), MMA_N, MMA_K)
|
||||
using MmaShapeB_NK = decltype(partition_shape_B(TiledMma{}, make_shape(cute::size<1>(TileShape_MNK{}),
|
||||
cute::size<2>(TileShape_MNK{}))));
|
||||
|
||||
using BlockTileA_M = decltype(cute::size<0,0>(MmaShapeA_MK{}) * cute::size<1>(MmaShapeA_MK{}));
|
||||
using BlockTileA_K = decltype(cute::size<0,1>(MmaShapeA_MK{}) * cute::size<2>(MmaShapeA_MK{}));
|
||||
using BlockTileB_N = decltype(cute::size<0,0>(MmaShapeB_NK{}) * cute::size<1>(MmaShapeB_NK{}));
|
||||
using BlockTileB_K = decltype(cute::size<0,1>(MmaShapeB_NK{}) * cute::size<2>(MmaShapeB_NK{}));
|
||||
|
||||
static_assert(BlockTileA_K{} == BlockTileB_K{}, "Block tile Ks should be equal");
|
||||
|
||||
using SmemShape_M = decltype(shape_div(shape<0>(TileShape_MNK{}), shape_div(shape<0>(TileShape_MNK{}), size<0>(TileShape_MNK{}) / size(AtomThrID{}))));
|
||||
using SmemShape_N = decltype(shape_div(shape<1>(TileShape_MNK{}), shape_div(shape<1>(TileShape_MNK{}), size<1>(TileShape_MNK{}) / size(AtomThrID{}))));
|
||||
using SmemShape_K = decltype(cute::get<2>(TileShape_MNK{}));
|
||||
|
||||
using GmemTiledCopyA = decltype(cutlass::gemm::collective::detail::sm100_cluster_shape_to_tma_atom_A(
|
||||
ClusterShape_MNK{}, AtomThrID{}));
|
||||
using GmemTiledCopyB = decltype(cutlass::gemm::collective::detail::sm100_cluster_shape_to_tma_atom_B(
|
||||
ClusterShape_MNK{}, AtomThrID{}));
|
||||
|
||||
using SmemLayoutAtomA = decltype(cutlass::gemm::collective::detail::sm100_smem_selector<
|
||||
UmmaMajorA, ElementAMma_SmemAllocType, SmemShape_M, SmemShape_K>());
|
||||
using SmemLayoutAtomB = decltype(cutlass::gemm::collective::detail::sm100_smem_selector<
|
||||
UmmaMajorB, ElementBMma_SmemAllocType, SmemShape_N, SmemShape_K>());
|
||||
static constexpr uint32_t TotalTmemRows = 128;
|
||||
static constexpr uint32_t Sm100TmemCapacityColumns = 512;
|
||||
static constexpr uint32_t TotalTmem = TotalTmemRows * Sm100TmemCapacityColumns;
|
||||
static constexpr uint32_t AccumulatorPipelineStageCount = (is_2sm || (!is_2sm && size(shape<0,0>(MmaShapeA_MK{}) > 64))) ?
|
||||
TotalTmem / (cute::size<0>(CtaTileShape_MNK{}) * cute::size<1>(CtaTileShape_MNK{}))
|
||||
: (Sm100TmemCapacityColumns / cute::size<1>(CtaTileShape_MNK{})) * 2; // 1SM MMA_M = 64 case
|
||||
static_assert(AccumulatorPipelineStageCount > 0, "Accumulator pipeline stage count must be positive. This error probably means that TileShape_MNK and/or TiledMma::ThrLayoutVMNK are wrong.");
|
||||
|
||||
// Calculate scheduler pipeline stages. Having one more stage than the accumulator allows more latency hiding.
|
||||
using StrideA = cutlass::gemm::TagToStrideA_t<GmemLayoutATag>;
|
||||
using InternalStrideA = cute::remove_pointer_t<StrideA>;
|
||||
// Grouped GEMM (where Stride type is Stride*) does not use CLC based scheduler.
|
||||
// SchedulerPipelineStageCount could be set to zero for Grouped GEMM, but we shouldn't define CLC Pipeline's barrier arrays of size zero.
|
||||
static constexpr uint32_t SchedulerPipelineStageCount = cute::is_same_v<InternalStrideA, StrideA> ? (AccumulatorPipelineStageCount + 1) : 1;
|
||||
|
||||
static constexpr uint32_t KernelSmemCarveout = detail::Sm100DenseGemmTmaUmmaCarveout<
|
||||
ClusterShape_MNK,
|
||||
AccumulatorPipelineStageCount,
|
||||
SchedulerPipelineStageCount,
|
||||
detail::CLCResponseSize,
|
||||
false
|
||||
>::KernelSmemCarveout;
|
||||
// Reduce SMEM capacity available for buffers considering barrier allocations.
|
||||
static constexpr int Sm100ReducedSmemCapacityBytes = cutlass::gemm::collective::detail::sm100_smem_capacity_bytes - KernelSmemCarveout;
|
||||
|
||||
using SmemTileShape = cute::Shape<BlockTileA_M, BlockTileB_N, BlockTileA_K>;
|
||||
using MainloopPipelineStorage = typename cutlass::PipelineTmaUmmaAsync<1>::SharedStorage;
|
||||
using TransformLoadPipelineStorage = typename cutlass::PipelineAsync<1>::SharedStorage;
|
||||
using TransformPipelineStorage = typename cutlass::PipelineUmmaAsync<1>::SharedStorage;
|
||||
|
||||
static constexpr int ScaleGranularityM = size<0,0>(GmemLayoutSFATag{});
|
||||
static constexpr int ScaleGranularityN = size<0,0>(GmemLayoutSFBTag{});
|
||||
static constexpr int ScaleGranularityK = size<1,0>(GmemLayoutSFBTag{});
|
||||
|
||||
static_assert(size<0>(CtaTileShape_MNK{}) >= ScaleGranularityM, "Scale Granularity must be smaller than or equal to the tile shape");
|
||||
static_assert(size<1>(CtaTileShape_MNK{}) >= ScaleGranularityN, "Scale Granularity must be smaller than or equal to the tile shape");
|
||||
static_assert(size<2>(CtaTileShape_MNK{}) >= ScaleGranularityK, "Scale Granularity must be smaller than or equal to the tile shape");
|
||||
|
||||
using BlockTileScale_M = Int<size<0>(TileShape_MNK{}) / ScaleGranularityM>;
|
||||
using BlockTileScale_N = Int<size<1>(TileShape_MNK{}) / ScaleGranularityN>;
|
||||
using BlockTileScale_K = Int<size<2>(TileShape_MNK{}) / ScaleGranularityK>;
|
||||
|
||||
using ScaleTileShape = cute::Shape<BlockTileScale_M, BlockTileScale_N, BlockTileScale_K>;
|
||||
|
||||
static constexpr int PipelineStages = cutlass::gemm::collective::detail::sm100_compute_stage_count_or_override_blockwise<
|
||||
Sm100ReducedSmemCapacityBytes, ElementAMma_SmemAllocType, ElementBMma_SmemAllocType,
|
||||
ElementAccumulator, ScaleTileShape, SmemTileShape, MainloopPipelineStorage,
|
||||
TransformLoadPipelineStorage, TransformPipelineStorage>(StageCountType{});
|
||||
static_assert(PipelineStages > 0, "Smem usage is too high. Can't create any SMEM buffers for A, B, and scales.");
|
||||
|
||||
using DispatchPolicy = cutlass::gemm::MainloopSm100TmaUmmaWarpSpecializedBlockwiseScaling<
|
||||
PipelineStages,
|
||||
SchedulerPipelineStageCount,
|
||||
AccumulatorPipelineStageCount,
|
||||
ClusterShape_MNK>;
|
||||
|
||||
using CollectiveOp = cutlass::gemm::collective::CollectiveMma<
|
||||
DispatchPolicy,
|
||||
TileShape_MNK,
|
||||
ElementA,
|
||||
cute::tuple<cutlass::gemm::TagToStrideA_t<GmemLayoutATag>, cutlass::gemm::TagToStrideA_t<GmemLayoutSFATag>>,
|
||||
ElementB,
|
||||
cute::tuple<cutlass::gemm::TagToStrideB_t<GmemLayoutBTag>, cutlass::gemm::TagToStrideB_t<GmemLayoutSFBTag>>,
|
||||
TiledMma,
|
||||
GmemTiledCopyA,
|
||||
SmemLayoutAtomA,
|
||||
void,
|
||||
cute::identity,
|
||||
GmemTiledCopyB,
|
||||
SmemLayoutAtomB,
|
||||
void,
|
||||
cute::identity
|
||||
>;
|
||||
};
|
||||
|
||||
|
||||
} // namespace cutlass::gemm::collective
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
@@ -1046,8 +1046,7 @@ template <
|
||||
class TileShape_MNK,
|
||||
class ClusterShape_MNK,
|
||||
class StageCountType,
|
||||
int ScaleGranularityM_,
|
||||
int ScaleGranularityN_
|
||||
class KernelScheduleType
|
||||
>
|
||||
struct CollectiveBuilder<
|
||||
arch::Sm90,
|
||||
@@ -1062,11 +1061,16 @@ struct CollectiveBuilder<
|
||||
TileShape_MNK,
|
||||
ClusterShape_MNK,
|
||||
StageCountType,
|
||||
KernelTmaWarpSpecializedCooperativeFP8BlockScaledAccum<ScaleGranularityM_, ScaleGranularityN_>,
|
||||
KernelScheduleType,
|
||||
cute::enable_if_t<
|
||||
not detail::is_use_rmem_A<ElementA, GmemLayoutATag, ElementB, GmemLayoutBTag>()>
|
||||
cute::is_same_v<decltype(KernelScheduleType::ScaleGranularityM), decltype(KernelScheduleType::ScaleGranularityN)> and
|
||||
not detail::is_use_rmem_A<ElementA, GmemLayoutATag, ElementB, GmemLayoutBTag>()
|
||||
>
|
||||
> {
|
||||
using KernelScheduleType = KernelTmaWarpSpecializedCooperativeFP8BlockScaledAccum<ScaleGranularityM_, ScaleGranularityN_>;
|
||||
|
||||
static constexpr auto ScaleGranularityM_ = KernelScheduleType::ScaleGranularityM;
|
||||
static constexpr auto ScaleGranularityN_ = KernelScheduleType::ScaleGranularityN;
|
||||
static constexpr auto ScalePromotionInterval_ = KernelScheduleType::ScalePromotionInterval;
|
||||
|
||||
static_assert(is_static<TileShape_MNK>::value);
|
||||
static_assert(is_static<ClusterShape_MNK>::value);
|
||||
@@ -1076,12 +1080,12 @@ struct CollectiveBuilder<
|
||||
static_assert(detail::is_aligned<ElementA, AlignmentA, ElementB, AlignmentB, detail::tma_alignment_bytes>(),
|
||||
"Should meet TMA alignment requirement\n");
|
||||
|
||||
static constexpr bool IsArrayOfPointersGemm = (cute::is_any_of_v<KernelScheduleType,
|
||||
KernelPtrArrayTmaWarpSpecializedCooperative,
|
||||
KernelPtrArrayTmaWarpSpecializedPingpong>);
|
||||
static constexpr bool IsArrayOfPointersGemm = (
|
||||
cute::is_base_of_v<KernelPtrArrayTmaWarpSpecializedCooperative, KernelScheduleType> ||
|
||||
cute::is_base_of_v<KernelPtrArrayTmaWarpSpecializedPingpong, KernelScheduleType>);
|
||||
|
||||
static constexpr bool IsFP8Input = detail::is_input_fp8<ElementA, ElementB>();
|
||||
static_assert((!IsFP8Input || !IsArrayOfPointersGemm),
|
||||
"KernelTmaWarpSpecializedCooperativeFP8BlockScaledAccum is only compatible with FP8 Blocked Scaled version right now.");
|
||||
static_assert(IsFP8Input, "Warp Specialized gemm with FP8 BlockScaled Accumulator is only compatible with FP8 Blocked Scaled version right now.");
|
||||
|
||||
// For fp32 types, map to tf32 MMA value type
|
||||
using ElementAMma = cute::conditional_t<cute::is_same_v<ElementA, float>, tfloat32_t, ElementA>;
|
||||
@@ -1091,10 +1095,9 @@ struct CollectiveBuilder<
|
||||
static constexpr cute::GMMA::Major GmmaMajorA = detail::gmma_ss_tag_to_major_A<ElementAMma, GmemLayoutATag>();
|
||||
static constexpr cute::GMMA::Major GmmaMajorB = detail::gmma_ss_tag_to_major_B<ElementBMma, GmemLayoutBTag>();
|
||||
|
||||
static constexpr bool IsCooperative = cute::is_any_of_v<KernelScheduleType,
|
||||
KernelTmaWarpSpecializedCooperative,
|
||||
KernelPtrArrayTmaWarpSpecializedCooperative,
|
||||
KernelTmaWarpSpecializedCooperativeFP8BlockScaledAccum<ScaleGranularityM_, ScaleGranularityN_>>;
|
||||
static constexpr bool IsCooperative = cute::is_base_of_v<KernelTmaWarpSpecializedCooperative, KernelScheduleType> ||
|
||||
cute::is_base_of_v<KernelPtrArrayTmaWarpSpecializedCooperative, KernelScheduleType>;
|
||||
|
||||
using AtomLayoutMNK = cute::conditional_t<IsCooperative,
|
||||
Layout<Shape<_2,_1,_1>>, Layout<Shape<_1,_1,_1>>>;
|
||||
|
||||
@@ -1121,7 +1124,9 @@ struct CollectiveBuilder<
|
||||
|
||||
static constexpr int PipelineStages = detail::compute_stage_count_with_blockwise_scale<detail::sm90_smem_capacity_bytes - KernelSmemCarveout,
|
||||
ElementAMma, ElementBMma, ElementBlockScale, TileShape_MNK, ScaleMsPerTile, ScaleNsPerTile>(StageCountType{});
|
||||
using DispatchPolicy = MainloopSm90TmaGmmaWarpSpecializedBlockScalingFP8<PipelineStages, ClusterShape_MNK, KernelScheduleType, ScaleGranularityM_, ScaleGranularityN_>;
|
||||
using DispatchPolicy = cute::conditional_t<IsArrayOfPointersGemm,
|
||||
MainloopSm90ArrayTmaGmmaWarpSpecializedBlockScaling<PipelineStages, ClusterShape_MNK, KernelScheduleType, ScaleGranularityM_, ScaleGranularityN_, ScalePromotionInterval_>,
|
||||
MainloopSm90TmaGmmaWarpSpecializedBlockScalingFP8<PipelineStages, ClusterShape_MNK, KernelScheduleType, ScaleGranularityM_, ScaleGranularityN_, ScalePromotionInterval_>>;
|
||||
|
||||
using SmemCopyAtomA = void;
|
||||
using SmemCopyAtomB = void;
|
||||
|
||||
@@ -43,6 +43,7 @@
|
||||
#include "cutlass/gemm/collective/builders/sm100_umma_builder.inl"
|
||||
#include "cutlass/gemm/collective/builders/sm100_9xBF16_umma_builder.inl"
|
||||
#include "cutlass/gemm/collective/builders/sm100_blockscaled_umma_builder.inl"
|
||||
#include "cutlass/gemm/collective/builders/sm100_blockwise_umma_builder.inl"
|
||||
#endif
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
@@ -50,6 +50,7 @@
|
||||
#include "cutlass/gemm/collective/sm90_mma_tma_gmma_ss_warpspecialized_fp8.hpp"
|
||||
#include "cutlass/gemm/collective/sm90_mma_array_tma_gmma_ss_warpspecialized_fp8.hpp"
|
||||
#include "cutlass/gemm/collective/sm90_mma_tma_gmma_ss_warpspecialized_fp8_blockwise_scaling.hpp"
|
||||
#include "cutlass/gemm/collective/sm90_mma_array_tma_gmma_ss_warpspecialized_fp8_blockwise_scaling.hpp"
|
||||
|
||||
#if !defined(__CUDACC_RTC__)
|
||||
#include "cutlass/gemm/collective/sm100_mma_warpspecialized.hpp"
|
||||
@@ -59,5 +60,6 @@
|
||||
|
||||
#include "cutlass/gemm/collective/sm100_blockscaled_mma_warpspecialized.hpp"
|
||||
#include "cutlass/gemm/collective/sm100_blockscaled_mma_array_warpspecialized.hpp"
|
||||
#include "cutlass/gemm/collective/sm100_mma_warpspecialized_blockwise_scaling.hpp"
|
||||
#endif // !defined(__CUDACC_RTC__)
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
@@ -223,6 +223,30 @@ public:
|
||||
mma_count_ = 0;
|
||||
}
|
||||
}
|
||||
|
||||
/// scale (multiply_add) the results from the MMA accumulators to main accumulator without checking the counter.
|
||||
CUTLASS_DEVICE
|
||||
void scale(ElementAccumulator const &scale) {
|
||||
scale_core(scale);
|
||||
}
|
||||
|
||||
template <
|
||||
class EngineScale,
|
||||
class LayoutScale>
|
||||
CUTLASS_DEVICE
|
||||
void scale(const cute::Tensor<EngineScale, LayoutScale> &scale) {
|
||||
scale_core(scale);
|
||||
}
|
||||
|
||||
template <
|
||||
class EngineScaleA,
|
||||
class LayoutScaleA,
|
||||
class EngineScaleB,
|
||||
class LayoutScaleB>
|
||||
CUTLASS_DEVICE
|
||||
void scale(const cute::Tensor<EngineScaleA, LayoutScaleA> &scaleA, const cute::Tensor<EngineScaleB, LayoutScaleB> &scaleB) {
|
||||
scale_core(scaleA, scaleB);
|
||||
}
|
||||
|
||||
/// scale (multiply_add) the residue results from the MMA accumulators to main accumulator if needed.
|
||||
CUTLASS_DEVICE
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -204,6 +204,8 @@ public:
|
||||
using PipelineState = cutlass::PipelineState<DispatchPolicy::Stages>;
|
||||
using PipelineParams = typename MainloopPipeline::Params;
|
||||
|
||||
static constexpr int NumProducerThreadEvents = 1;
|
||||
|
||||
using SmemLayoutAtomScale = Layout<Shape<decltype(cute::shape<0>(SwappedSmemLayoutAtomA{})), cute::Int<1>>>;
|
||||
using ScaleTileShape = decltype(make_shape(shape<0>(TileShape{}), shape<1>(SmemLayoutAtomScale{})));
|
||||
|
||||
@@ -1354,6 +1356,18 @@ public:
|
||||
static_assert(cutlass::detail::dependent_false<KernelSchedule>, "Conversion mode not handled in tensormaps_fence_acquire.");
|
||||
}
|
||||
}
|
||||
|
||||
template <class InputTensors, class ProblemShape_MNKL>
|
||||
CUTLASS_DEVICE
|
||||
InputTensors
|
||||
tensors_perform_update(
|
||||
InputTensors const& input_tensors,
|
||||
[[maybe_unused]] Params const& mainloop_params,
|
||||
[[maybe_unused]] ProblemShape_MNKL problem_shape_mnkl,
|
||||
[[maybe_unused]] int32_t next_batch) {
|
||||
return input_tensors;
|
||||
}
|
||||
|
||||
};
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
@@ -116,6 +116,9 @@ struct CollectiveMma<
|
||||
|
||||
using PipelineParams = typename MainloopPipeline::Params;
|
||||
using CtaShape_MNK = decltype(shape_div(TileShape{}, ClusterShape{}));
|
||||
|
||||
static constexpr int NumProducerThreadEvents = 1;
|
||||
|
||||
static_assert(rank(SmemLayoutAtomA{}) == 2, "SmemLayoutAtom must be rank 2 (M/N, K)");
|
||||
static_assert((size<0>(TileShape{}) % size<0>(SmemLayoutAtomA{})) == 0, "SmemLayoutAtom must evenly divide tile shape.");
|
||||
static_assert((size<2>(TileShape{}) % size<1>(SmemLayoutAtomA{})) == 0, "SmemLayoutAtom must evenly divide tile shape.");
|
||||
@@ -749,6 +752,16 @@ struct CollectiveMma<
|
||||
cute::tma_descriptor_fence_acquire(get<1>(input_tensormaps));
|
||||
}
|
||||
|
||||
template <class InputTensors, class ProblemShape_MNKL>
|
||||
CUTLASS_DEVICE
|
||||
InputTensors
|
||||
tensors_perform_update(
|
||||
InputTensors const& input_tensors,
|
||||
[[maybe_unused]] Params const& mainloop_params,
|
||||
[[maybe_unused]] ProblemShape_MNKL problem_shape_mnkl,
|
||||
[[maybe_unused]] int32_t next_batch) {
|
||||
return input_tensors;
|
||||
}
|
||||
|
||||
};
|
||||
|
||||
|
||||
@@ -759,6 +759,18 @@ struct CollectiveMma<
|
||||
cute::tma_descriptor_fence_acquire(get<0>(input_tensormaps));
|
||||
cute::tma_descriptor_fence_acquire(get<1>(input_tensormaps));
|
||||
}
|
||||
|
||||
template <class InputTensors, class ProblemShape_MNKL>
|
||||
CUTLASS_DEVICE
|
||||
InputTensors
|
||||
tensors_perform_update(
|
||||
InputTensors const& input_tensors,
|
||||
[[maybe_unused]] Params const& mainloop_params,
|
||||
[[maybe_unused]] ProblemShape_MNKL problem_shape_mnkl,
|
||||
[[maybe_unused]] int32_t next_batch) {
|
||||
return input_tensors;
|
||||
}
|
||||
|
||||
};
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -59,6 +59,7 @@ template <
|
||||
class KernelSchedule,
|
||||
int ScaleGranularityM_,
|
||||
int ScaleGranularityN_,
|
||||
int ScalePromotionInterval_,
|
||||
class TileShape_,
|
||||
class ElementA_,
|
||||
class StrideA_,
|
||||
@@ -74,7 +75,7 @@ template <
|
||||
class SmemCopyAtomB_,
|
||||
class TransformB_>
|
||||
struct CollectiveMma<
|
||||
MainloopSm90TmaGmmaWarpSpecializedBlockScalingFP8<Stages, ClusterShape, KernelSchedule, ScaleGranularityM_, ScaleGranularityN_>,
|
||||
MainloopSm90TmaGmmaWarpSpecializedBlockScalingFP8<Stages, ClusterShape, KernelSchedule, ScaleGranularityM_, ScaleGranularityN_, ScalePromotionInterval_>,
|
||||
TileShape_,
|
||||
ElementA_,
|
||||
StrideA_,
|
||||
@@ -93,7 +94,7 @@ struct CollectiveMma<
|
||||
//
|
||||
// Type Aliases
|
||||
//
|
||||
using DispatchPolicy = MainloopSm90TmaGmmaWarpSpecializedBlockScalingFP8<Stages, ClusterShape, KernelSchedule, ScaleGranularityM_, ScaleGranularityN_>;
|
||||
using DispatchPolicy = MainloopSm90TmaGmmaWarpSpecializedBlockScalingFP8<Stages, ClusterShape, KernelSchedule, ScaleGranularityM_, ScaleGranularityN_, ScalePromotionInterval_>;
|
||||
using TileShape = TileShape_;
|
||||
using ElementA = ElementA_;
|
||||
using StrideA = StrideA_;
|
||||
@@ -122,6 +123,8 @@ struct CollectiveMma<
|
||||
|
||||
static constexpr int ScaleGranularityM = ScaleGranularityM_ == 0 ? size<0>(TileShape{}) : ScaleGranularityM_;
|
||||
static constexpr int ScaleGranularityN = ScaleGranularityN_ == 0 ? size<1>(TileShape{}) : ScaleGranularityN_;
|
||||
static constexpr int ScalePromotionInterval = ScalePromotionInterval_;
|
||||
static_assert(ScalePromotionInterval % 4 == 0, "ScalePromotionInterval must be a multiple of 4.");
|
||||
static constexpr int ScaleMsPerTile = size<0>(TileShape{}) / ScaleGranularityM;
|
||||
static constexpr int ScaleNsPerTile = size<1>(TileShape{}) / ScaleGranularityN;
|
||||
|
||||
@@ -281,7 +284,9 @@ struct CollectiveMma<
|
||||
constexpr int min_tma_aligned_elements_B = tma_alignment_bits / cutlass::sizeof_bits<ElementB>::value;
|
||||
implementable = implementable && cutlass::detail::check_alignment<min_tma_aligned_elements_B>(cute::make_shape(N,K,L), StrideB{});
|
||||
/* MMA promotion interval should be a multiple of 4, since each mainloop iteration would issue 4 MMA instructions. */
|
||||
implementable = implementable && (args.mma_promotion_interval % 4 == 0);
|
||||
constexpr int pipe_k = size<2>(TileShape{}) / tile_size<2>(TiledMma{});
|
||||
implementable = implementable && (args.mma_promotion_interval % 4 == 0) && (args.mma_promotion_interval == ScalePromotionInterval);
|
||||
implementable = implementable && (pipe_k % 4 == 0) && (pipe_k <= args.mma_promotion_interval);
|
||||
|
||||
if (!implementable) {
|
||||
CUTLASS_TRACE_HOST(" CAN IMPLEMENT: Problem Size doesn't meet the minimum alignment requirements for TMA.\n");
|
||||
@@ -481,6 +486,38 @@ struct CollectiveMma<
|
||||
}
|
||||
}
|
||||
|
||||
template<
|
||||
class EngineAccum,
|
||||
class LayoutAccum,
|
||||
class ScaleFactor
|
||||
>
|
||||
CUTLASS_DEVICE
|
||||
void scale_if_needed(GmmaFP8Accumulation<EngineAccum, LayoutAccum>& accumulation, ScaleFactor scaleFactor) {
|
||||
if constexpr (ScalePromotionInterval != 4) {
|
||||
accumulation.scale_if_needed(scaleFactor);
|
||||
}
|
||||
else {
|
||||
// avoid unnecessary tests when granularity is the finnest
|
||||
accumulation.scale(scaleFactor);
|
||||
}
|
||||
}
|
||||
template<
|
||||
class EngineAccum,
|
||||
class LayoutAccum,
|
||||
class ScaleFactor1,
|
||||
class ScaleFactor2
|
||||
>
|
||||
CUTLASS_DEVICE
|
||||
void scale_if_needed(GmmaFP8Accumulation<EngineAccum, LayoutAccum>& accumulation, ScaleFactor1 scaleFactor1, ScaleFactor2 scaleFactor2) {
|
||||
if constexpr (ScalePromotionInterval != 4) {
|
||||
accumulation.scale_if_needed(scaleFactor1, scaleFactor2);
|
||||
}
|
||||
else {
|
||||
// avoid unnecessary tests when granularity is the finnest
|
||||
accumulation.scale(scaleFactor1, scaleFactor2);
|
||||
}
|
||||
}
|
||||
|
||||
/// Perform a collective-scoped matrix multiply-accumulate
|
||||
/// Consumer Perspective
|
||||
template <
|
||||
@@ -575,7 +612,7 @@ struct CollectiveMma<
|
||||
|
||||
tiled_mma.accumulate_ = GMMA::ScaleOut::Zero;
|
||||
|
||||
GmmaFP8Accumulation accumulation(accum, mainloop_params.mma_promotion_interval, size<2>(tCrA));
|
||||
GmmaFP8Accumulation accumulation(accum, ScalePromotionInterval, size<2>(tCrA));
|
||||
warpgroup_fence_operand(accumulation());
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int k_tile_prologue = prologue_mma_count; k_tile_prologue > 0; --k_tile_prologue)
|
||||
@@ -584,7 +621,13 @@ struct CollectiveMma<
|
||||
auto barrier_token = pipeline.consumer_try_wait(smem_pipe_read);
|
||||
pipeline.consumer_wait(smem_pipe_read, barrier_token);
|
||||
|
||||
if (accumulation.prepare_if_needed()) {
|
||||
if constexpr (ScalePromotionInterval != 4) {
|
||||
if (accumulation.prepare_if_needed()) {
|
||||
tiled_mma.accumulate_ = GMMA::ScaleOut::Zero;
|
||||
}
|
||||
}
|
||||
else {
|
||||
// Always zero out the accumulator for finest granularity
|
||||
tiled_mma.accumulate_ = GMMA::ScaleOut::Zero;
|
||||
}
|
||||
|
||||
@@ -624,16 +667,16 @@ struct CollectiveMma<
|
||||
// Block scale the accumulators with reg tensor `tCrScaleAViewAsC` and `tCrScaleBViewAsC`
|
||||
if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile == 1) {
|
||||
ElementBlockScale scale_ab = tCrScaleAViewAsC.data()[0];
|
||||
accumulation.scale_if_needed(scale_ab);
|
||||
scale_if_needed(accumulation, scale_ab);
|
||||
}
|
||||
if constexpr (ScaleMsPerTile > 1 && ScaleNsPerTile == 1) {
|
||||
accumulation.scale_if_needed(tCrScaleAViewAsC);
|
||||
scale_if_needed(accumulation, tCrScaleAViewAsC);
|
||||
}
|
||||
if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile > 1) {
|
||||
accumulation.scale_if_needed(tCrScaleBViewAsC);
|
||||
scale_if_needed(accumulation, tCrScaleBViewAsC);
|
||||
}
|
||||
if constexpr (ScaleMsPerTile > 1 && ScaleNsPerTile > 1) {
|
||||
accumulation.scale_if_needed(tCrScaleAViewAsC, tCrScaleBViewAsC);
|
||||
scale_if_needed(accumulation, tCrScaleAViewAsC, tCrScaleBViewAsC);
|
||||
}
|
||||
|
||||
++smem_pipe_read;
|
||||
@@ -677,7 +720,13 @@ struct CollectiveMma<
|
||||
}
|
||||
}
|
||||
|
||||
if (accumulation.prepare_if_needed()) {
|
||||
if constexpr (ScalePromotionInterval != 4) {
|
||||
if (accumulation.prepare_if_needed()) {
|
||||
tiled_mma.accumulate_ = GMMA::ScaleOut::Zero;
|
||||
}
|
||||
}
|
||||
else {
|
||||
// Always zero out the accumulator for finest granularity
|
||||
tiled_mma.accumulate_ = GMMA::ScaleOut::Zero;
|
||||
}
|
||||
|
||||
@@ -699,16 +748,16 @@ struct CollectiveMma<
|
||||
// Block scale the accumulators with reg tensor `tCrScaleAViewAsC` and `tCrScaleBViewAsC`
|
||||
if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile == 1) {
|
||||
ElementBlockScale scale_ab = tCrScaleAViewAsC.data()[0];
|
||||
accumulation.scale_if_needed(scale_ab);
|
||||
scale_if_needed(accumulation, scale_ab);
|
||||
}
|
||||
if constexpr (ScaleMsPerTile > 1 && ScaleNsPerTile == 1) {
|
||||
accumulation.scale_if_needed(tCrScaleAViewAsC);
|
||||
scale_if_needed(accumulation, tCrScaleAViewAsC);
|
||||
}
|
||||
if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile > 1) {
|
||||
accumulation.scale_if_needed(tCrScaleBViewAsC);
|
||||
scale_if_needed(accumulation, tCrScaleBViewAsC);
|
||||
}
|
||||
if constexpr (ScaleMsPerTile > 1 && ScaleNsPerTile > 1) {
|
||||
accumulation.scale_if_needed(tCrScaleAViewAsC, tCrScaleBViewAsC);
|
||||
scale_if_needed(accumulation, tCrScaleAViewAsC, tCrScaleBViewAsC);
|
||||
}
|
||||
|
||||
pipeline.consumer_release(smem_pipe_release); // UNLOCK smem_pipe_release, done _computing_ on it
|
||||
@@ -718,18 +767,21 @@ struct CollectiveMma<
|
||||
++smem_pipe_release;
|
||||
}
|
||||
|
||||
if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile == 1) {
|
||||
ElementBlockScale scale_ab = tCrScaleAViewAsC.data()[0];
|
||||
accumulation.scale_residue_if_needed(scale_ab);
|
||||
}
|
||||
if constexpr (ScaleMsPerTile > 1 && ScaleNsPerTile == 1) {
|
||||
accumulation.scale_residue_if_needed(tCrScaleAViewAsC);
|
||||
}
|
||||
if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile > 1) {
|
||||
accumulation.scale_residue_if_needed(tCrScaleBViewAsC);
|
||||
}
|
||||
if constexpr (ScaleMsPerTile > 1 && ScaleNsPerTile > 1) {
|
||||
accumulation.scale_residue_if_needed(tCrScaleAViewAsC, tCrScaleBViewAsC);
|
||||
if constexpr (ScalePromotionInterval != 4) {
|
||||
// residues only exists when granularity is not the finnest
|
||||
if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile == 1) {
|
||||
ElementBlockScale scale_ab = tCrScaleAViewAsC.data()[0];
|
||||
accumulation.scale_residue_if_needed(scale_ab);
|
||||
}
|
||||
if constexpr (ScaleMsPerTile > 1 && ScaleNsPerTile == 1) {
|
||||
accumulation.scale_residue_if_needed(tCrScaleAViewAsC);
|
||||
}
|
||||
if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile > 1) {
|
||||
accumulation.scale_residue_if_needed(tCrScaleBViewAsC);
|
||||
}
|
||||
if constexpr (ScaleMsPerTile > 1 && ScaleNsPerTile > 1) {
|
||||
accumulation.scale_residue_if_needed(tCrScaleAViewAsC, tCrScaleBViewAsC);
|
||||
}
|
||||
}
|
||||
|
||||
warpgroup_fence_operand(accumulation());
|
||||
|
||||
@@ -120,10 +120,38 @@ template<
|
||||
// `ScaleGranularityM`/`ScaleGranularityN` specifies scaling granularity along M/N, while zero-value
|
||||
// `ScaleGranularityM`/`ScaleGranularityN` indicates that scaling granularity is
|
||||
// `size<0>(TileShape_MNK{})`/`size<1>(TileShape_MNK{})` along M/N.
|
||||
int ScaleGranularityM = 0,
|
||||
int ScaleGranularityN = 0
|
||||
int ScaleGranularityM_ = 0,
|
||||
int ScaleGranularityN_ = 0,
|
||||
// `ScalePromotionInterval` specifies the interval to promote the accumulator for scaling
|
||||
// It is required to be a multiple of 4 and specified in terms of number of MMA instructions
|
||||
// in the reduction dimension. i.e for FP8 kernels, it is
|
||||
// ScalePromotionInterval * MMA_K = ScalePromotionInterval * 32 = 128 elements in K by default
|
||||
int ScalePromotionInterval_ = 4
|
||||
|
||||
>
|
||||
struct KernelTmaWarpSpecializedCooperativeFP8BlockScaledAccum: KernelTmaWarpSpecializedCooperative { };
|
||||
struct KernelTmaWarpSpecializedCooperativeFP8BlockScaledAccum: KernelTmaWarpSpecializedCooperative {
|
||||
constexpr static int ScaleGranularityM = ScaleGranularityM_;
|
||||
constexpr static int ScaleGranularityN = ScaleGranularityN_;
|
||||
constexpr static int ScalePromotionInterval = ScalePromotionInterval_;
|
||||
};
|
||||
|
||||
template<
|
||||
// `ScaleGranularityM`/`ScaleGranularityN` specifies scaling granularity along M/N, while zero-value
|
||||
// `ScaleGranularityM`/`ScaleGranularityN` indicates that scaling granularity is
|
||||
// `size<0>(TileShape_MNK{})`/`size<1>(TileShape_MNK{})` along M/N.
|
||||
int ScaleGranularityM_,
|
||||
int ScaleGranularityN_,
|
||||
// `ScalePromotionInterval` specifies the interval to promote the accumulator for scaling
|
||||
// It is required to be a multiple of 4 and specified in terms of number of MMA instructions
|
||||
// in the reduction dimension. i.e for FP8 kernels, it is
|
||||
// ScalePromotionInterval * MMA_K = ScalePromotionInterval * 32 = 128 elements in K by default
|
||||
int ScalePromotionInterval_ = 4
|
||||
>
|
||||
struct KernelPtrArrayTmaWarpSpecializedCooperativeFP8BlockScaledAccum: KernelPtrArrayTmaWarpSpecializedCooperative {
|
||||
constexpr static int ScaleGranularityM = ScaleGranularityM_;
|
||||
constexpr static int ScaleGranularityN = ScaleGranularityN_;
|
||||
constexpr static int ScalePromotionInterval = ScalePromotionInterval_;
|
||||
};
|
||||
|
||||
// Policies to opt into mixed type GEMMs
|
||||
struct KernelTmaWarpSpecializedMixedInput : KernelTmaWarpSpecialized { };
|
||||
@@ -310,12 +338,17 @@ template<
|
||||
// `ScaleGranularityM`/`ScaleGranularityN` indicates that scaling granularity is
|
||||
// `size<0>(TileShape_MNK{})`/`size<1>(TileShape_MNK{})` along M/N.
|
||||
int ScaleGranularityM = 0,
|
||||
int ScaleGranularityN = 0
|
||||
int ScaleGranularityN = 0,
|
||||
// `ScalePromotionInterval` specifies the interval to promote the accumulator for scaling
|
||||
// It is required to be a multiple of 4 and specified in terms of number of MMA instructions
|
||||
// in the reduction dimension. i.e for FP8 kernels, it is
|
||||
// ScalePromotionInterval * MMA_K = ScalePromotionInterval * 32 = 128 elements in K by default
|
||||
int ScalePromotionInterval = 4
|
||||
>
|
||||
struct MainloopSm90TmaGmmaWarpSpecializedBlockScalingFP8
|
||||
: MainloopSm90TmaGmmaWarpSpecialized<Stages_, ClusterShape_, KernelSchedule> {
|
||||
static_assert(
|
||||
cute::is_same_v<KernelSchedule, KernelTmaWarpSpecializedCooperativeFP8BlockScaledAccum<ScaleGranularityM, ScaleGranularityN>>,
|
||||
cute::is_same_v<KernelSchedule, KernelTmaWarpSpecializedCooperativeFP8BlockScaledAccum<ScaleGranularityM, ScaleGranularityN, ScalePromotionInterval>>,
|
||||
"KernelSchedule must be one of the warp specialized policies");
|
||||
};
|
||||
|
||||
@@ -327,6 +360,7 @@ template<
|
||||
>
|
||||
struct MainloopSm90ArrayTmaGmmaWarpSpecialized {
|
||||
constexpr static int Stages = Stages_;
|
||||
constexpr static int PipelineAsyncMmaStages = 1;
|
||||
using ClusterShape = ClusterShape_;
|
||||
using ArchTag = arch::Sm90;
|
||||
using Schedule = KernelSchedule;
|
||||
@@ -391,6 +425,26 @@ struct MainloopSm90ArrayTmaGmmaWarpSpecializedMixedInput {
|
||||
"KernelSchedule must be one of the Ptr-Array or Grouped Gemm TMA Warp Specialized Cooperative policies");
|
||||
};
|
||||
|
||||
// n-buffer in smem (Hopper TMA), pipelined with Hopper GMMA and TMA, Warp specialized dynamic schedule
|
||||
// For FP8 kernels with Block Scaling
|
||||
template<
|
||||
int Stages_,
|
||||
class ClusterShape_ = Shape<_1,_1,_1>,
|
||||
class KernelSchedule = KernelPtrArrayTmaWarpSpecializedCooperative,
|
||||
// `ScaleGranularityM`/`ScaleGranularityN` specifies scaling granularity along M/N, while zero-value
|
||||
// `ScaleGranularityM`/`ScaleGranularityN` indicates that scaling granularity is
|
||||
// `size<0>(TileShape_MNK{})`/`size<1>(TileShape_MNK{})` along M/N.
|
||||
int ScaleGranularityM = 0,
|
||||
int ScaleGranularityN = 0,
|
||||
int ScalePromotionInterval = 4
|
||||
>
|
||||
struct MainloopSm90ArrayTmaGmmaWarpSpecializedBlockScaling
|
||||
: MainloopSm90ArrayTmaGmmaWarpSpecialized<Stages_, ClusterShape_, KernelSchedule> {
|
||||
static_assert(
|
||||
cute::is_same_v<KernelSchedule, KernelPtrArrayTmaWarpSpecializedCooperativeFP8BlockScaledAccum<ScaleGranularityM, ScaleGranularityN>>,
|
||||
"KernelSchedule must be one of the warp specialized policies");
|
||||
};
|
||||
|
||||
|
||||
template<
|
||||
int SchedulerPipelineStageCount_,
|
||||
@@ -411,6 +465,14 @@ struct KernelTmaWarpSpecializedBlockScaledSm100 final {
|
||||
static constexpr int AccumulatorPipelineStageCount = AccumulatorPipelineStageCount_;
|
||||
};
|
||||
|
||||
template<
|
||||
int SchedulerPipelineStageCount_,
|
||||
int AccumulatorPipelineStageCount_
|
||||
>
|
||||
struct KernelTmaWarpSpecializedMmaTransformSm100 final {
|
||||
static constexpr int SchedulerPipelineStageCount = SchedulerPipelineStageCount_;
|
||||
static constexpr int AccumulatorPipelineStageCount = AccumulatorPipelineStageCount_;
|
||||
};
|
||||
|
||||
|
||||
// InputTransform GEMM
|
||||
@@ -484,6 +546,13 @@ struct KernelScheduleSm100PtrArrayDenseGemm : KernelScheduleSm100DenseGemm {};
|
||||
struct KernelPtrArrayTmaWarpSpecialized1SmSm100 final : KernelSchedule1Sm, KernelScheduleSm100PtrArrayDenseGemm {};
|
||||
struct KernelPtrArrayTmaWarpSpecialized2SmSm100 final : KernelSchedule2Sm, KernelScheduleSm100PtrArrayDenseGemm {};
|
||||
|
||||
|
||||
///////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
// SM100 Blockwise GEMM Dispatch Policies
|
||||
///////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
struct KernelScheduleSm100Blockwise : KernelScheduleSm100 {};
|
||||
struct KernelTmaWarpSpecializedBlockwise1SmSm100 final : KernelSchedule1Sm, KernelScheduleSm100Blockwise {};
|
||||
|
||||
///////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
// SM100 Planar Complex GEMM Dispatch Policies
|
||||
///////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
@@ -530,6 +599,9 @@ struct KernelTmaWarpSpecialized1SmMxf4Sm100 final : KernelSchedule1Sm, KernelSch
|
||||
struct KernelTmaWarpSpecialized2SmMxf4Sm100 final : KernelSchedule2Sm, KernelScheduleMxNvf4Sm100 { };
|
||||
struct KernelTmaWarpSpecialized1SmMxf8f6f4Sm100 final : KernelSchedule1Sm, KernelScheduleMxf8f6f4Sm100 { };
|
||||
struct KernelTmaWarpSpecialized2SmMxf8f6f4Sm100 final : KernelSchedule2Sm, KernelScheduleMxf8f6f4Sm100 { };
|
||||
///////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
// SM100 BlockScaled Ptr Array Dense GEMM Dispatch Policies
|
||||
///////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
// BlockScaled Dense GEMM + (Ptr Array or Group GEMM)
|
||||
struct KernelSchedulePtrArrayBlockScaledGemmSm100 : KernelScheduleBlockScaledGemmSm100 {};
|
||||
struct KernelSchedulePtrArrayMxNvf4Sm100 : KernelSchedulePtrArrayBlockScaledGemmSm100 {};
|
||||
@@ -544,8 +616,6 @@ struct KernelPtrArrayTmaWarpSpecialized2SmMxf4Sm100 final : KernelSchedule2Sm, K
|
||||
struct KernelPtrArrayTmaWarpSpecialized1SmMxf8f6f4Sm100 final : KernelSchedule1Sm, KernelSchedulePtrArrayMxf8f6f4Sm100 { };
|
||||
struct KernelPtrArrayTmaWarpSpecialized2SmMxf8f6f4Sm100 final : KernelSchedule2Sm, KernelSchedulePtrArrayMxf8f6f4Sm100 { };
|
||||
|
||||
|
||||
|
||||
// n-buffer in smem, pipelined with Blackwell UMMA and TMA, Warp specialized dynamic schedule
|
||||
template<
|
||||
int Stages_,
|
||||
@@ -561,7 +631,20 @@ struct MainloopSm100TmaUmmaWarpSpecialized {
|
||||
constexpr static bool IsOverlappingAccum = false;
|
||||
};
|
||||
|
||||
|
||||
// n-buffer in smem, pipelined with Blackwell UMMA and TMA, Warp specialized dynamic schedule
|
||||
template<
|
||||
int Stages_,
|
||||
int SchedulerPipelineStageCount_,
|
||||
int AccumulatorPipelineStageCount_,
|
||||
class ClusterShape_ = Shape<_1,_1,_1>
|
||||
>
|
||||
struct MainloopSm100TmaUmmaWarpSpecializedBlockwiseScaling {
|
||||
constexpr static int Stages = Stages_;
|
||||
using ClusterShape = ClusterShape_;
|
||||
using ArchTag = arch::Sm100;
|
||||
using Schedule = KernelTmaWarpSpecializedMmaTransformSm100<SchedulerPipelineStageCount_, AccumulatorPipelineStageCount_>;
|
||||
constexpr static bool IsOverlappingAccum = false;
|
||||
};
|
||||
|
||||
// n-buffer in smem, pipelined with Blackwell UMMA and TMA, Warp specialized dynamic schedule
|
||||
template<
|
||||
|
||||
@@ -65,6 +65,7 @@ struct IsCutlass3ArrayKernel<ProblemShape, cute::void_t<typename ProblemShape::U
|
||||
#include "cutlass/gemm/kernel/sm90_gemm_array_tma_warpspecialized_cooperative.hpp"
|
||||
|
||||
#include "cutlass/gemm/kernel/sm100_gemm_tma_warpspecialized.hpp"
|
||||
#include "cutlass/gemm/kernel/sm100_gemm_tma_warpspecialized_mma_transform.hpp"
|
||||
#include "cutlass/gemm/kernel/sm100_gemm_array_tma_warpspecialized.hpp"
|
||||
#include "cutlass/gemm/kernel/sm100_gemm_tma_warpspecialized_input_transform.hpp"
|
||||
#include "cutlass/gemm/kernel/sm100_gemm_array_tma_warpspecialized_input_transform.hpp"
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -128,10 +128,11 @@ public:
|
||||
using TileSchedulerParams = typename TileScheduler::Params;
|
||||
|
||||
static constexpr uint32_t NumLoadWarpGroups = 1;
|
||||
static constexpr uint32_t NumMmaThreads = CUTE_STATIC_V(size(TiledMma{}));
|
||||
static constexpr uint32_t NumMmaThreads = size(TiledMma{});
|
||||
static constexpr uint32_t NumMmaWarpGroups = NumMmaThreads / NumThreadsPerWarpGroup;
|
||||
static constexpr uint32_t MaxThreadsPerBlock = NumMmaThreads + (NumLoadWarpGroups * NumThreadsPerWarpGroup);
|
||||
static constexpr uint32_t MinBlocksPerMultiprocessor = 1;
|
||||
static constexpr uint32_t NumProducerThreads = CollectiveMainloop::NumProducerThreadEvents;
|
||||
|
||||
/// Register requirement for Load and Math WGs
|
||||
static constexpr uint32_t LoadRegisterRequirement = 40;
|
||||
@@ -434,7 +435,8 @@ public:
|
||||
mainloop_pipeline_params.role = MainloopPipeline::ThreadCategory::Consumer;
|
||||
}
|
||||
mainloop_pipeline_params.is_leader = warp_group_thread_idx == 0;
|
||||
mainloop_pipeline_params.num_consumers = size(TiledMma{});
|
||||
mainloop_pipeline_params.num_consumers = NumMmaThreads;
|
||||
mainloop_pipeline_params.num_producers = NumProducerThreads;
|
||||
mainloop_pipeline_params.transaction_bytes = params.mainloop.tma_transaction_bytes;
|
||||
MainloopPipeline mainloop_pipeline(shared_storage.pipelines.mainloop, mainloop_pipeline_params, ClusterShape{});
|
||||
|
||||
@@ -575,6 +577,7 @@ public:
|
||||
auto k_tile_iter = cute::make_coord_iterator(idx2crd(work_k_tile_start, shape<3>(gA_mkl)), shape<3>(gA_mkl));
|
||||
|
||||
if (did_batch_change) {
|
||||
load_inputs = collective_mainloop.tensors_perform_update(load_inputs, params.mainloop, problem_shape_MNKL, curr_batch);
|
||||
collective_mainloop.tensormaps_fence_acquire(input_tensormaps);
|
||||
}
|
||||
|
||||
|
||||
@@ -131,6 +131,7 @@ public:
|
||||
static constexpr uint32_t NumMmaWarpGroups = 2;
|
||||
static constexpr uint32_t MaxThreadsPerBlock = CUTE_STATIC_V(size(TiledMma{})) + (NumMmaWarpGroups * NumThreadsPerWarpGroup);
|
||||
static constexpr uint32_t MinBlocksPerMultiprocessor = 1;
|
||||
static constexpr uint32_t NumProducerThreads = CollectiveMainloop::NumProducerThreadEvents;
|
||||
|
||||
/// Register requirement for Load and Math WGs
|
||||
static constexpr uint32_t LoadRegisterRequirement = 40;
|
||||
@@ -443,6 +444,7 @@ public:
|
||||
}
|
||||
mainloop_pipeline_params.is_leader = warp_group_thread_idx == 0;
|
||||
mainloop_pipeline_params.num_consumers = NumThreadsPerWarpGroup;
|
||||
mainloop_pipeline_params.num_producers = NumProducerThreads;
|
||||
mainloop_pipeline_params.transaction_bytes = params.mainloop.tma_transaction_bytes;
|
||||
MainloopPipeline mainloop_pipeline(shared_storage.pipelines.mainloop, mainloop_pipeline_params, ClusterShape{});
|
||||
|
||||
@@ -607,6 +609,7 @@ public:
|
||||
auto k_tile_iter = cute::make_coord_iterator(idx2crd(work_k_tile_start, shape<3>(gA_mkl)), shape<3>(gA_mkl));
|
||||
|
||||
if (did_batch_change) {
|
||||
load_inputs = collective_mainloop.tensors_perform_update(load_inputs, params.mainloop, problem_shape_MNKL, curr_batch);
|
||||
collective_mainloop.tensormaps_fence_acquire(input_tensormaps);
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user