v4.2 release. (#2587)
* Fix default cluster callback values to 1 to avoid profiler failure when these values are not set in command line. * v4.2 release.
This commit is contained in:
Junkai-Wu
@@ -132,7 +132,7 @@ auto sm100_make_simt_gmem_tiled_copy_SFA() {
|
||||
return make_tiled_copy(
|
||||
SmemScalingCopyAtomA{},
|
||||
Layout<Shape<_32>>{}, // 32 threads
|
||||
Layout<Shape<Shape<Int<ScaleGranularityM>, Int<ElementsPerSFACopy>>>, Stride<Stride<_0, _1>>>{});
|
||||
Layout<Shape<Int<ElementsPerSFACopy>>>{});
|
||||
}
|
||||
else {
|
||||
using SmemScalingCopyAtomA = Copy_Atom<SM80_CP_ASYNC_CACHEALWAYS<Element>, Element>;
|
||||
@@ -166,7 +166,7 @@ auto sm100_make_simt_gmem_tiled_copy_SFB() {
|
||||
return make_tiled_copy(
|
||||
SmemScalingCopyAtomB{},
|
||||
Layout<Shape<_32>>{}, // 32 threads
|
||||
Layout<Shape<Shape<Int<ScaleGranularityN>, Int<ElementsPerSFBCopy>>>, Stride<Stride<_0, _1>>>{});
|
||||
Layout<Shape<Int<ElementsPerSFBCopy>>>{});
|
||||
}
|
||||
else {
|
||||
using SmemScalingCopyAtomB = Copy_Atom<SM80_CP_ASYNC_CACHEALWAYS<Element>, Element>;
|
||||
|
||||
@@ -569,6 +569,73 @@ sm100_make_trivial_fastFP32_tiled_mma() {
|
||||
}
|
||||
}
|
||||
|
||||
//Setting mma for Mixed input gemm. Here, ElementAMma should be TACompute
|
||||
template<
|
||||
class ElementAMma,
|
||||
class ElementBMma,
|
||||
class ElementAccumulator,
|
||||
class TileShape_MNK,
|
||||
class ClusterShape_MNK,
|
||||
UMMA::Major UmmaMajorA,
|
||||
UMMA::Major UmmaMajorB,
|
||||
class KernelScheduleType
|
||||
>
|
||||
constexpr auto
|
||||
sm100_make_trivial_mixed_input_tiled_mma() {
|
||||
constexpr int M = cute::size<0>(TileShape_MNK{});
|
||||
constexpr int N = cute::size<1>(TileShape_MNK{});
|
||||
//MMA 1Sm requested
|
||||
if constexpr (cute::is_base_of_v<KernelSchedule1Sm, KernelScheduleType> ) {
|
||||
if constexpr (UmmaMajorA == cute::UMMA::Major::K && !cute::is_base_of_v<KernelTmaWarpSpecializedMixedInputSmemSm100, KernelScheduleType>) {
|
||||
if constexpr (cute::is_same_v<ElementBMma, cutlass::half_t> || cute::is_same_v<ElementBMma, cutlass::bfloat16_t>) {
|
||||
return make_tiled_mma(cute::SM100_MMA_F16BF16_TS<ElementAMma, ElementBMma, ElementAccumulator,
|
||||
M, N, UmmaMajorA, UmmaMajorB>{});
|
||||
}
|
||||
if constexpr (cute::is_same_v<ElementBMma, cutlass::float_e4m3_t>) {
|
||||
return make_tiled_mma(cute::SM100_MMA_F8F6F4_TS<ElementAMma, ElementBMma, ElementAccumulator,
|
||||
M, N, UmmaMajorA, UmmaMajorB>{});
|
||||
}
|
||||
}
|
||||
else { // If A needs to be transposed by MMA, fall back to SMEM from A MMA instructions
|
||||
if constexpr (cute::is_same_v<ElementBMma, cutlass::half_t> || cute::is_same_v<ElementBMma, cutlass::bfloat16_t>) {
|
||||
return make_tiled_mma(cute::SM100_MMA_F16BF16_SS<ElementAMma, ElementBMma, ElementAccumulator,
|
||||
M, N, UmmaMajorA, UmmaMajorB>{});
|
||||
}
|
||||
if constexpr (cute::is_same_v<ElementBMma, cutlass::float_e4m3_t>) {
|
||||
return make_tiled_mma(
|
||||
cute::MMA_Traits<
|
||||
cute::SM100_MMA_F8F6F4_SS,
|
||||
ElementAMma,
|
||||
ElementBMma,
|
||||
ElementAccumulator,
|
||||
cute::C<M>,
|
||||
cute::C<N>,
|
||||
cute::integral_constant<UMMA::Major, UmmaMajorA>,
|
||||
cute::integral_constant<UMMA::Major, UmmaMajorB>,
|
||||
cute::integral_constant<UMMA::ScaleIn, cute::UMMA::ScaleIn::One>,
|
||||
cute::integral_constant<UMMA::ScaleIn, cute::UMMA::ScaleIn::One>>{});
|
||||
}
|
||||
}
|
||||
}
|
||||
//MMA 2Sm requested
|
||||
else if constexpr (cute::is_base_of_v<KernelSchedule2Sm, KernelScheduleType>) {
|
||||
if constexpr (UmmaMajorA == cute::UMMA::Major::K && !cute::is_base_of_v<KernelTmaWarpSpecializedMixedInputSmemSm100, KernelScheduleType>) {
|
||||
if constexpr (cute::is_same_v<ElementBMma, cutlass::half_t> || cute::is_same_v<ElementBMma, cutlass::bfloat16_t>) {
|
||||
return make_tiled_mma(cute::SM100_MMA_F16BF16_2x1SM_TS<ElementAMma, ElementBMma, ElementAccumulator,
|
||||
M, N, UmmaMajorA, UmmaMajorB>{});
|
||||
}
|
||||
if constexpr (cute::is_same_v<ElementBMma, cutlass::float_e4m3_t>) {
|
||||
return make_tiled_mma(cute::SM100_MMA_F8F6F4_2x1SM_TS<ElementAMma, ElementBMma, ElementAccumulator,
|
||||
M, N, UmmaMajorA, UmmaMajorB>{});
|
||||
}
|
||||
}
|
||||
}
|
||||
else {
|
||||
static_assert(cutlass::detail::dependent_false<TileShape_MNK> == 0,
|
||||
"Unsupported policy for SM100 collective builder.");
|
||||
}
|
||||
}
|
||||
|
||||
template<
|
||||
class CtaShape_MNK
|
||||
>
|
||||
|
||||
@@ -0,0 +1,173 @@
|
||||
/***************************************************************************************************
|
||||
* Copyright (c) 2024 - 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"
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace cutlass::gemm::collective {
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <
|
||||
class ElementA,
|
||||
class GmemLayoutATag,
|
||||
int AlignmentA,
|
||||
class ElementB,
|
||||
class GmemLayoutBTag,
|
||||
int AlignmentB,
|
||||
class ElementAccumulator,
|
||||
class TileShape_MNK,
|
||||
class ClusterShape_MNK,
|
||||
class StageCountType,
|
||||
class BuilderScheduleTag
|
||||
>
|
||||
struct CollectiveBuilder<
|
||||
arch::Sm100,
|
||||
arch::OpClassTensorOp,
|
||||
ElementA,
|
||||
GmemLayoutATag,
|
||||
AlignmentA,
|
||||
ElementB,
|
||||
GmemLayoutBTag,
|
||||
AlignmentB,
|
||||
ElementAccumulator,
|
||||
TileShape_MNK, // (MmaAtomShapeM, MmaAtomShapeN, TileK)
|
||||
ClusterShape_MNK, // Static cluster shape or dynamic (int, int, _1)
|
||||
StageCountType,
|
||||
BuilderScheduleTag,
|
||||
cute::enable_if_t<cute::is_same_v<KernelWarpSpecialized1SmSm100, BuilderScheduleTag> ||
|
||||
(cute::is_same_v<KernelScheduleAuto, BuilderScheduleTag> &&
|
||||
(((sizeof(ElementA) * AlignmentA) % cutlass::gemm::collective::detail::tma_alignment_bytes != 0) ||
|
||||
((sizeof(ElementB) * AlignmentB) % cutlass::gemm::collective::detail::tma_alignment_bytes != 0)))>
|
||||
>
|
||||
{
|
||||
static_assert(cute::is_static_v<TileShape_MNK>, "TileShape has to be static");
|
||||
|
||||
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::sm1xx_kernel_input_element_to_mma_input_element<ElementA>());
|
||||
using ElementBMma = decltype(cutlass::gemm::collective::detail::sm1xx_kernel_input_element_to_mma_input_element<ElementB>());
|
||||
|
||||
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 TiledMma = decltype(detail::sm100_make_trivial_tiled_mma<
|
||||
ElementAMma, ElementBMma, ElementAccumulator,
|
||||
decltype(cute::product_each(TileShape_MNK{})), ClusterShape_MNK,
|
||||
UmmaMajorA, UmmaMajorB, BuilderScheduleTag>());
|
||||
|
||||
using AtomThrID = typename TiledMma::AtomThrID;
|
||||
|
||||
// ((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{}))));
|
||||
|
||||
// Assigning 4 warps for mainloop load
|
||||
static constexpr int NumLoadThreads = 128;
|
||||
|
||||
using AlignmentTypeA = cute::uint_byte_t<static_cast<int>(sizeof(ElementA)) * AlignmentA>;
|
||||
using GmemCopyAtomA = cute::Copy_Atom<SM80_CP_ASYNC_CACHEALWAYS_ZFILL<AlignmentTypeA>, ElementA>;
|
||||
using GmemTiledCopyA = decltype(detail::make_simt_gmem_tiled_copy<
|
||||
GmemCopyAtomA, NumLoadThreads, AlignmentA, TagToStrideA_t<GmemLayoutATag>,
|
||||
decltype(cute::get<0>(TileShape_MNK{})), decltype(cute::get<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 SmemLayoutAtomA = decltype(cutlass::gemm::collective::detail::sm100_smem_selector<
|
||||
UmmaMajorA, ElementAMma_SmemAllocType, BlockTileA_M, BlockTileA_K>());
|
||||
|
||||
using AlignmentTypeB = cute::uint_byte_t<static_cast<int>(sizeof(ElementB)) * AlignmentB>;
|
||||
using GmemCopyAtomB = cute::Copy_Atom<SM80_CP_ASYNC_CACHEALWAYS_ZFILL<AlignmentTypeB>, ElementB>;
|
||||
using GmemTiledCopyB = decltype(detail::make_simt_gmem_tiled_copy<
|
||||
GmemCopyAtomB, NumLoadThreads, AlignmentB, TagToStrideB_t<GmemLayoutBTag>,
|
||||
decltype(cute::get<1>(TileShape_MNK{})), decltype(cute::get<2>(TileShape_MNK{}))>());
|
||||
|
||||
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{}));
|
||||
using SmemLayoutAtomB = decltype(cutlass::gemm::collective::detail::sm100_smem_selector<
|
||||
UmmaMajorB, ElementBMma_SmemAllocType, BlockTileB_N, BlockTileB_K>());
|
||||
|
||||
static constexpr uint32_t AccumulatorPipelineStageCount = 2;
|
||||
// Calculate scheduler pipeline stages. Having one more stage than the accumulator allows more latency hiding.
|
||||
static constexpr uint32_t SchedulerPipelineStageCount = AccumulatorPipelineStageCount + 1;
|
||||
|
||||
// AccumulatorPipeline = PipelineUmmaAsync
|
||||
static constexpr auto AccumulatorPipelineStorage = sizeof(typename cutlass::PipelineUmmaAsync<AccumulatorPipelineStageCount>::SharedStorage);
|
||||
// CLCPipeline = PipelineCLCFetchAsync
|
||||
static constexpr auto CLCPipelineStorage = sizeof(typename cutlass::PipelineCLCFetchAsync<SchedulerPipelineStageCount, ClusterShape_MNK>::SharedStorage);
|
||||
// CLC (scheduler) response
|
||||
static constexpr auto CLCResponseStorage = SchedulerPipelineStageCount * detail::CLCResponseSize;
|
||||
// Smem usage that's not part of CollectiveEpilogue::SharedStorage & CollectiveMainloop::SharedStorage
|
||||
static constexpr auto KernelSmemCarveout = static_cast<int>( AccumulatorPipelineStorage +
|
||||
CLCPipelineStorage +
|
||||
CLCResponseStorage);
|
||||
// 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::PipelineUmmaConsumerAsync<1>::SharedStorage;
|
||||
|
||||
static constexpr int PipelineStages = cutlass::gemm::collective::detail::sm100_compute_stage_count_or_override<
|
||||
Sm100ReducedSmemCapacityBytes, ElementAMma_SmemAllocType, ElementBMma_SmemAllocType, SmemTileShape, MainloopPipelineStorage>(StageCountType{});
|
||||
|
||||
using CollectiveOp = cutlass::gemm::collective::CollectiveMma<
|
||||
cutlass::gemm::MainloopSm100UmmaCpAsyncWarpSpecialized<
|
||||
PipelineStages,
|
||||
SchedulerPipelineStageCount,
|
||||
AccumulatorPipelineStageCount,
|
||||
ClusterShape_MNK>,
|
||||
TileShape_MNK,
|
||||
ElementA,
|
||||
cutlass::gemm::TagToStrideA_t<GmemLayoutATag>,
|
||||
ElementB,
|
||||
cutlass::gemm::TagToStrideB_t<GmemLayoutBTag>,
|
||||
TiledMma,
|
||||
GmemTiledCopyA,
|
||||
SmemLayoutAtomA,
|
||||
void,
|
||||
cute::identity,
|
||||
GmemTiledCopyB,
|
||||
SmemLayoutAtomB,
|
||||
void,
|
||||
cute::identity
|
||||
>;
|
||||
};
|
||||
|
||||
} // namespace cutlass::gemm::collective
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
@@ -0,0 +1,346 @@
|
||||
/***************************************************************************************************
|
||||
* Copyright (c) 2024 - 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"
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
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 ElementAMma,
|
||||
class ElementScale,
|
||||
class ElementZero,
|
||||
class ElementB,
|
||||
class CtaTileShape_MNK,
|
||||
class TiledMma,
|
||||
class KernelScheduleType,
|
||||
UMMA::Major UmmaMajorA,
|
||||
int ScaleGranularityK,
|
||||
int stages
|
||||
>
|
||||
constexpr cute::tuple<int, int, int>
|
||||
sm100_compute_stage_count_or_override_mixed_input(StageCount<stages> stage_count) {
|
||||
constexpr int Load2TransformStageCount = stages;
|
||||
constexpr int Transform2MmaStageCount = stages;
|
||||
constexpr int AccumulatorStageCount = stages;
|
||||
return cute::make_tuple(Load2TransformStageCount, Transform2MmaStageCount, AccumulatorStageCount);
|
||||
}
|
||||
|
||||
template<
|
||||
int CapacityBytes,
|
||||
class ElementA,
|
||||
class ElementAMma,
|
||||
class ElementScale,
|
||||
class ElementZero,
|
||||
class ElementB,
|
||||
class CtaTileShape_MNK,
|
||||
class TiledMma,
|
||||
class KernelScheduleType,
|
||||
UMMA::Major UmmaMajorA,
|
||||
int ScaleGranularityK,
|
||||
int carveout_bytes
|
||||
>
|
||||
constexpr cute::tuple<int, int, int>
|
||||
sm100_compute_stage_count_or_override_mixed_input(StageCountAutoCarveout<carveout_bytes> stage_count) {
|
||||
|
||||
constexpr int CtaM = get<0>(CtaTileShape_MNK{});
|
||||
constexpr int CtaN = get<1>(CtaTileShape_MNK{});
|
||||
static_assert(CtaN <= 128, "Can't support CtaN>128 tiles");
|
||||
constexpr int CtaK = get<2>(CtaTileShape_MNK{});
|
||||
using AtomThrID = typename TiledMma::AtomThrID;
|
||||
|
||||
constexpr int TmemColumns = 512;
|
||||
|
||||
constexpr bool IsAComputeinTmem = UmmaMajorA == cute::UMMA::Major::K && !cute::is_base_of_v<KernelTmaWarpSpecializedMixedInputSmemSm100, KernelScheduleType>;
|
||||
constexpr bool IsAComputeinSmem = !IsAComputeinTmem;
|
||||
|
||||
// Detect 2x2 TMEM layout
|
||||
constexpr int TmemAccWordsPerDP = (CtaM == 64 && size(AtomThrID{}) == 2) ? CtaN/2 : CtaN;
|
||||
constexpr int TmemAWordsPerDP = CtaK / 2;
|
||||
|
||||
constexpr int AccumulatorStageCount = (IsAComputeinTmem) ? ((TmemAccWordsPerDP == 128) ? 2 : 3) : (TmemColumns / TmemAccWordsPerDP);
|
||||
|
||||
constexpr int SmemCapacityAfterMma2AccumCarveout = CapacityBytes - (carveout_bytes + AccumulatorStageCount * 32);
|
||||
|
||||
constexpr int TmemInAStageCount_Potential = (IsAComputeinTmem) ? (TmemColumns - AccumulatorStageCount * TmemAccWordsPerDP) / TmemAWordsPerDP : 10000;
|
||||
|
||||
// Mainload2Transform Pipeline
|
||||
constexpr auto load2transform_pipeline_bytes = sizeof(typename cutlass::PipelineTmaTransformAsync<1>::SharedStorage);
|
||||
constexpr auto a_bits = cute::sizeof_bits_v<ElementA>; // ElementA introduce here
|
||||
constexpr auto s_bits = cute::is_void_v<ElementScale> ? 0 : cute::sizeof_bits_v<ElementScale>;
|
||||
constexpr auto z_bits = cute::is_void_v<ElementZero> ? 0 : cute::sizeof_bits_v<ElementZero>;
|
||||
|
||||
constexpr auto load2mma_pipeline_bytes = sizeof(typename cutlass::PipelineTmaUmmaAsync<1>::SharedStorage);
|
||||
constexpr auto b_bits = cute::sizeof_bits_v<ElementB>; // ElementB introduce here
|
||||
|
||||
constexpr int ab_stage_bytes =
|
||||
cutlass::bits_to_bytes(a_bits * size<0>(CtaTileShape_MNK{}) * size<2>(CtaTileShape_MNK{})) +
|
||||
cutlass::bits_to_bytes(s_bits * size<0>(CtaTileShape_MNK{}) * size<2>(CtaTileShape_MNK{}) / ScaleGranularityK) +
|
||||
cutlass::bits_to_bytes(z_bits * size<0>(CtaTileShape_MNK{}) * size<2>(CtaTileShape_MNK{}) / ScaleGranularityK) +
|
||||
cutlass::bits_to_bytes(b_bits * size<1>(CtaTileShape_MNK{}) / size(AtomThrID{}) * size<2>(CtaTileShape_MNK{})) +
|
||||
static_cast<int>(load2transform_pipeline_bytes) + static_cast<int>(load2mma_pipeline_bytes);
|
||||
|
||||
// Transform2Mma Pipeline
|
||||
constexpr auto transform2mma_pipeline_bytes = sizeof(typename cutlass::PipelineUmmaConsumerAsync<1>::SharedStorage);
|
||||
constexpr auto a_compute_bits = cute::sizeof_bits_v<ElementAMma>;
|
||||
constexpr int ab_compute_stage_bytes =
|
||||
cutlass::bits_to_bytes(a_compute_bits * int(IsAComputeinSmem) * size<0>(CtaTileShape_MNK{}) * size<2>(CtaTileShape_MNK{})) + // If ACompute is in TMEM, Acompute buffer has 0 bytes.
|
||||
static_cast<int>(transform2mma_pipeline_bytes);
|
||||
|
||||
constexpr int ABComputeStageCount_Potential = SmemCapacityAfterMma2AccumCarveout / (ab_stage_bytes + ab_compute_stage_bytes);
|
||||
|
||||
// The number of SMEM buffers for A, B. ACompute (if in SMEM), BCompute should be at least Transform2MmaStageCount
|
||||
constexpr int Transform2MmaStageCount = std::min(TmemInAStageCount_Potential, ABComputeStageCount_Potential);
|
||||
|
||||
constexpr int SmemCapacityAfterABComputeCarveout = SmemCapacityAfterMma2AccumCarveout - (Transform2MmaStageCount * ab_compute_stage_bytes);
|
||||
|
||||
// Can we boost the number of buffers for A and B?
|
||||
constexpr int Load2TransformStageCount = SmemCapacityAfterABComputeCarveout / ab_stage_bytes;
|
||||
|
||||
static_assert(Load2TransformStageCount >= 2 && Transform2MmaStageCount >= 2 && AccumulatorStageCount >= 2, "Not enough SMEM or TMEM capacity for selected tile size");
|
||||
return cute::make_tuple(Load2TransformStageCount, Transform2MmaStageCount, AccumulatorStageCount);
|
||||
}
|
||||
|
||||
} // namespace detail
|
||||
|
||||
template <typename LayoutScale>
|
||||
constexpr int get_ScaleGranularityK() {
|
||||
if constexpr (cute::is_void_v<LayoutScale>) {
|
||||
return 1;
|
||||
} else {
|
||||
return size<1,0>(LayoutScale{});
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
// Mixed Input MMA kernels builder
|
||||
template <
|
||||
class ElementAOptionalTuple,
|
||||
class GmemLayoutATagTuple,
|
||||
int AlignmentA,
|
||||
class ElementBOptionalTuple,
|
||||
class GmemLayoutBTag,
|
||||
int AlignmentB,
|
||||
class ElementAccumulator,
|
||||
class TileShape_MNK, // The Cluster-level TileShape
|
||||
class ClusterShape_MNK,
|
||||
class StageCountType,
|
||||
class KernelScheduleType
|
||||
>
|
||||
struct CollectiveBuilder<
|
||||
arch::Sm100,
|
||||
arch::OpClassTensorOp,
|
||||
ElementAOptionalTuple, // ElementA
|
||||
GmemLayoutATagTuple, // LayoutA
|
||||
AlignmentA,
|
||||
ElementBOptionalTuple, // ElementB
|
||||
GmemLayoutBTag, // LayoutB
|
||||
AlignmentB,
|
||||
ElementAccumulator,
|
||||
TileShape_MNK, // (MmaAtomShapeM, MmaAtomShapeN, TileK)
|
||||
ClusterShape_MNK, // Static cluster shape or dynamic (int, int, int)
|
||||
StageCountType,
|
||||
KernelScheduleType,
|
||||
cute::enable_if_t<
|
||||
(cute::is_base_of_v<KernelScheduleSm100MixedInputGemm, KernelScheduleType>) &&
|
||||
((sizeof(float) * AlignmentA) % detail::tma_alignment_bytes == 0) &&
|
||||
((sizeof(float) * AlignmentB) % detail::tma_alignment_bytes == 0)>>
|
||||
{
|
||||
using GmemLayoutATag = detail::deduce_mixed_width_dtype_t<0, GmemLayoutATagTuple>;
|
||||
using GmemLayoutScaleTag = detail::deduce_mixed_width_dtype_t<1, GmemLayoutATagTuple>;
|
||||
|
||||
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>();
|
||||
|
||||
using ElementA = detail::deduce_mixed_width_dtype_t<0, ElementAOptionalTuple>;
|
||||
using ElementB = detail::deduce_mixed_width_dtype_t<0, ElementBOptionalTuple>;
|
||||
using ElementScale = detail::deduce_mixed_width_dtype_t<1, ElementAOptionalTuple>;
|
||||
using ElementZero = detail::deduce_mixed_width_dtype_t<2, ElementAOptionalTuple>;
|
||||
|
||||
static constexpr bool NeitherIsTuple = !cute::is_tuple<ElementAOptionalTuple>::value && !cute::is_tuple<ElementBOptionalTuple>::value;
|
||||
static constexpr bool IsANarrow = cute::sizeof_bits_v<ElementA> < cute::sizeof_bits_v<ElementB>;
|
||||
static constexpr bool IsMixedInput = cute::sizeof_bits_v<ElementA> != cute::sizeof_bits_v<ElementB>;
|
||||
static_assert(IsMixedInput, "Mixed Input GEMM Kernel doesn't support regular gemm.");
|
||||
|
||||
static_assert((cute::is_tuple<ElementAOptionalTuple>::value ^ cute::is_tuple<ElementBOptionalTuple>::value ||
|
||||
(NeitherIsTuple && (cute::sizeof_bits<ElementA>::value != cute::sizeof_bits<ElementB>::value))),
|
||||
"Either A OR B must be a tuple or the widths of A and B must be different.");
|
||||
using ElementPairA = cute::conditional_t<IsMixedInput && IsANarrow && NeitherIsTuple, cute::tuple<ElementA>, ElementAOptionalTuple>;
|
||||
using ElementPairB = cute::conditional_t<IsMixedInput && !IsANarrow && NeitherIsTuple, cute::tuple<ElementB>, ElementBOptionalTuple>;
|
||||
static constexpr bool IsATransformed = cute::is_tuple<ElementPairA>::value;
|
||||
static_assert(IsATransformed, "A matrix should be transformed.");
|
||||
|
||||
// For fp32 types, map to tf32 MMA value type.
|
||||
using ElementMma = cute::conditional_t<cute::is_same_v<ElementB, float>, tfloat32_t, ElementB>;
|
||||
|
||||
|
||||
using ElementAMma = ElementMma;
|
||||
using ElementBMma = ElementMma;
|
||||
|
||||
static constexpr int IsSubbyteA = cute::sizeof_bits_v<ElementA> < 8;
|
||||
using TmaElementA = cute::conditional_t<IsSubbyteA, uint8_t, ElementA>;
|
||||
|
||||
static constexpr int ScalingFactor = 1;
|
||||
|
||||
using TiledMma = decltype(detail::sm100_make_trivial_mixed_input_tiled_mma<ElementAMma, ElementB, ElementAccumulator, TileShape_MNK, ClusterShape_MNK, UmmaMajorA, UmmaMajorB, KernelScheduleType>());
|
||||
using AtomThrID = typename TiledMma::AtomThrID;
|
||||
using AtomThrShapeMNK = Shape<decltype(shape<0>(typename TiledMma::ThrLayoutVMNK{})), _1, _1>;
|
||||
using CtaTileShape_MNK = decltype(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 GmemTiledCopyA = decltype(detail::sm90_cluster_shape_to_tma_atom(cute::size<1>(ClusterShape_MNK{})));
|
||||
using GmemTiledCopyB = decltype(detail::sm100_cluster_shape_to_tma_atom_B(ClusterShape_MNK{}, AtomThrID{}));
|
||||
|
||||
// Input transform kernel can not use TMA 2SM instructions.
|
||||
using SmemLayoutAtomA = decltype(cutlass::gemm::collective::detail::sm100_smem_selector<UmmaMajorA, ElementA,
|
||||
BlockTileA_M, BlockTileA_K>());
|
||||
using SmemLayoutAtomACompute = decltype(cutlass::gemm::collective::detail::sm100_smem_selector<UmmaMajorA, ElementAMma, BlockTileA_M, BlockTileA_K>());
|
||||
using SmemLayoutAtomPairA = cutlass::gemm::collective::detail::CollectiveMmaEmulatedLayoutAtomType<
|
||||
SmemLayoutAtomA, SmemLayoutAtomACompute>;
|
||||
static constexpr int MMA_M = cute::size<0,0>(MmaShapeA_MK{});
|
||||
using CopyAtomPairA = cutlass::gemm::collective::detail::CollectiveMmaEmulatedCopyType<
|
||||
Copy_Atom<AutoVectorizingCopyWithAssumedAlignment<128>, ElementA>,
|
||||
cute::conditional_t<(UmmaMajorA == cute::UMMA::Major::K && !cute::is_base_of_v<KernelTmaWarpSpecializedMixedInputSmemSm100, KernelScheduleType>),
|
||||
cute::conditional_t<(MMA_M == 64 && size(AtomThrID{}) == 1), SM100_TMEM_STORE_16dp256b1x, SM100_TMEM_STORE_32dp32b8x>, // TS Implementation
|
||||
Copy_Atom<AutoVectorizingCopyWithAssumedAlignment<128>, ElementA>> // SS Implementation
|
||||
>;
|
||||
|
||||
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{}));
|
||||
|
||||
// Input transform kernel can not use TMA 2SM instructions.
|
||||
using SmemLayoutAtomB = decltype(cutlass::gemm::collective::detail::sm100_smem_selector<UmmaMajorB, ElementB,
|
||||
BlockTileB_N, BlockTileB_K>());
|
||||
using SmemLayoutAtomBCompute = decltype(cutlass::gemm::collective::detail::sm100_smem_selector<UmmaMajorB, ElementBMma, BlockTileB_N, BlockTileB_K>());
|
||||
using SmemLayoutAtomPairB = cutlass::gemm::collective::detail::CollectiveMmaEmulatedLayoutAtomType<
|
||||
SmemLayoutAtomB, SmemLayoutAtomBCompute>;
|
||||
using CopyAtomPairB = cutlass::gemm::collective::detail::CollectiveMmaEmulatedCopyType<
|
||||
Copy_Atom<AutoVectorizingCopyWithAssumedAlignment<128>, ElementB>,
|
||||
Copy_Atom<AutoVectorizingCopyWithAssumedAlignment<128>, ElementMma>
|
||||
>;
|
||||
|
||||
//Creating the stride of Transformed Input
|
||||
using StrideA = cutlass::gemm::TagToStrideA_t<GmemLayoutATag>;
|
||||
using LayoutScale = cutlass::gemm::TagToStrideA_t<GmemLayoutScaleTag>;
|
||||
|
||||
using VoidShapeScale = Shape<Shape<Int<128>, _1>, Shape<Int<64>, _1>, _1>; //Dummy Value to create a dummy ScaleConfig
|
||||
using VoidStrideScale = Stride<Stride<_0,_1>,Stride<_0, _1>, _1>;
|
||||
using VoidLayoutScale = Layout<VoidShapeScale, VoidStrideScale>;
|
||||
|
||||
using NonVoidLayoutScale = cute::conditional_t<
|
||||
cute::is_void_v<LayoutScale>, VoidLayoutScale, LayoutScale>;
|
||||
|
||||
using StridePairA = decltype(cute::make_tuple(StrideA{}, NonVoidLayoutScale{}));
|
||||
|
||||
// SmemCarveout
|
||||
static constexpr int SchedulerPipelineStageCount = 3;
|
||||
static constexpr bool IsArrayOfPointersGemm = (cute::is_base_of_v<KernelScheduleSm100PtrArrayFastFP32Gemm, KernelScheduleType>);
|
||||
|
||||
// CLCPipeline = PipelineCLCFetchAsync
|
||||
static constexpr auto CLCPipelineStorage = sizeof(typename cutlass::PipelineCLCFetchAsync<SchedulerPipelineStageCount, ClusterShape_MNK>::SharedStorage);
|
||||
// CLC (scheduler) response
|
||||
static constexpr auto CLCResponseStorage = SchedulerPipelineStageCount * detail::CLCResponseSize;
|
||||
// CLC Throttle pipeline storage
|
||||
static constexpr auto CLCThrottlePipelineStorage = sizeof(typename cutlass::PipelineAsync<SchedulerPipelineStageCount>::SharedStorage);
|
||||
// Tmem dealloc
|
||||
static constexpr auto TmemDeallocStorage = sizeof(cutlass::arch::ClusterBarrier);
|
||||
// Tmem ptr storage
|
||||
static constexpr auto TmemBasePtrsStorage = sizeof(uint32_t);
|
||||
// Tensormap Storage
|
||||
static constexpr size_t TensorMapStorage = IsArrayOfPointersGemm ? sizeof(cute::TmaDescriptor) * 2 /* for A and B */ : 0;
|
||||
|
||||
// Smem usage that's not part of CollectiveEpilogue::SharedStorage & CollectiveMainloop::SharedStorage
|
||||
static constexpr auto KernelSmemCarveout = static_cast<int>( CLCPipelineStorage +
|
||||
CLCResponseStorage +
|
||||
CLCThrottlePipelineStorage +
|
||||
TmemDeallocStorage +
|
||||
TmemBasePtrsStorage +
|
||||
TensorMapStorage);
|
||||
|
||||
// Reduce SMEM capacity available for buffers considering extra B smem and barrier smem allocations
|
||||
static constexpr int Sm100ReducedSmemCapacityBytes = detail::sm100_smem_capacity_bytes - KernelSmemCarveout;
|
||||
|
||||
static constexpr int ScaleGranularityK = get_ScaleGranularityK<LayoutScale>();
|
||||
|
||||
static constexpr auto stage_info = cutlass::gemm::collective::detail::sm100_compute_stage_count_or_override_mixed_input<
|
||||
Sm100ReducedSmemCapacityBytes, TmaElementA, ElementAMma, ElementScale, ElementZero, ElementB, CtaTileShape_MNK, TiledMma, KernelScheduleType, UmmaMajorA, ScaleGranularityK>(StageCountType{});
|
||||
|
||||
static constexpr int Load2TransformPipelineStageCount = get<0>(stage_info);
|
||||
static constexpr int Transform2MmaPipelineStageCount = get<1>(stage_info);
|
||||
static constexpr int AccumulatorPipelineStageCount = get<2>(stage_info);
|
||||
|
||||
static_assert(!IsArrayOfPointersGemm, "mixed input does not support grouped gemm on Blackwell");
|
||||
|
||||
using DispatchPolicy = cutlass::gemm::MainloopSm100TmaUmmaWarpSpecializedMixedInput<
|
||||
Load2TransformPipelineStageCount,
|
||||
Transform2MmaPipelineStageCount,
|
||||
SchedulerPipelineStageCount,
|
||||
AccumulatorPipelineStageCount,
|
||||
ClusterShape_MNK
|
||||
>;
|
||||
using CollectiveOp = cutlass::gemm::collective::CollectiveMma<
|
||||
DispatchPolicy,
|
||||
TileShape_MNK,
|
||||
ElementPairA,
|
||||
StridePairA,
|
||||
ElementPairB,
|
||||
cutlass::gemm::TagToStrideB_t<GmemLayoutBTag>,
|
||||
TiledMma,
|
||||
GmemTiledCopyA,
|
||||
SmemLayoutAtomPairA,
|
||||
CopyAtomPairA,
|
||||
cute::identity,
|
||||
GmemTiledCopyB,
|
||||
SmemLayoutAtomPairB,
|
||||
CopyAtomPairB,
|
||||
cute::identity
|
||||
>;
|
||||
};
|
||||
|
||||
} // namespace cutlass::gemm::collective
|
||||
@@ -184,6 +184,7 @@ struct CollectiveBuilder<
|
||||
not cute::is_complex_v<ElementA> && not cute::is_complex_v<ElementB> &&
|
||||
// Dense Gemm / PtrArrayDenseGemm
|
||||
(
|
||||
(not cute::is_same_v<KernelWarpSpecialized1SmSm100, BuilderScheduleTag>) &&
|
||||
(cute::is_base_of_v<KernelScheduleSm100DenseGemm, BuilderScheduleTag> ||
|
||||
cute::is_same_v<KernelScheduleAuto, BuilderScheduleTag>)) &&
|
||||
// Alignment check
|
||||
|
||||
@@ -0,0 +1,546 @@
|
||||
/***************************************************************************************************
|
||||
* Copyright (c) 2024 - 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/detail/sm103_blockscaled_layout.hpp"
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
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 TileShapeMNK,
|
||||
class TileShapeSFA,
|
||||
class TileShapeSFB,
|
||||
int stages
|
||||
>
|
||||
constexpr int
|
||||
sm103_compute_stage_count_or_override_blockscaled(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 TileShapeMNK,
|
||||
class TileShapeSFA,
|
||||
class TileShapeSFB,
|
||||
int carveout_bytes
|
||||
>
|
||||
constexpr auto
|
||||
sm103_compute_stage_count_or_override_blockscaled(StageCountAutoCarveout<carveout_bytes> stage_count) {
|
||||
// For F8F6F4 MMA sub-bytes, ElementA/B will be passed in as uint8_t
|
||||
// Each stage include (CollectiveMma::SharedStorage)
|
||||
// 1. smem for A and smem for B (CollectiveMma::SharedStorage::TensorStorage)
|
||||
// 2. one MainloopPipeline = PipelineTmaUmmaAsync (CollectiveMma::SharedStorage::SharedStorage)
|
||||
// 3. smem for SFB and smem for SFB (CollectiveMma::SharedStorage::TensorStorage, independent of input size b.c. sizeof(sf) is fixed)
|
||||
constexpr auto mainloop_pipeline_bytes = sizeof(typename cutlass::PipelineTmaUmmaAsync<1>::SharedStorage);
|
||||
constexpr auto a_bits = cute::sizeof_bits_v<ElementA>;
|
||||
constexpr auto b_bits = cute::sizeof_bits_v<ElementB>;
|
||||
constexpr auto stage_sfa_bytes = size(filter_zeros(TileShapeSFA{}));
|
||||
constexpr auto stage_sfb_bytes = size(filter_zeros(TileShapeSFB{}));
|
||||
|
||||
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{})) +
|
||||
static_cast<int>(mainloop_pipeline_bytes * 2 + stage_sfa_bytes + stage_sfb_bytes);
|
||||
|
||||
constexpr int ab_buffer = (CapacityBytes - carveout_bytes) / stage_bytes;
|
||||
constexpr int sb_buffer = ab_buffer + (CapacityBytes - carveout_bytes - ab_buffer * stage_bytes) / (mainloop_pipeline_bytes + stage_sfa_bytes + stage_sfb_bytes);
|
||||
return make_tuple(ab_buffer, sb_buffer);
|
||||
}
|
||||
|
||||
template<
|
||||
class ElementAMma,
|
||||
class ElementBMma,
|
||||
class ElementAccumulator,
|
||||
class ElementSF,
|
||||
class TileShape_MNK,
|
||||
class ClusterShape_MNK,
|
||||
UMMA::Major UmmaMajorA,
|
||||
UMMA::Major UmmaMajorB,
|
||||
int SFVectorSize
|
||||
>
|
||||
constexpr auto
|
||||
sm103_make_blockscaled_1sm_tiled_mma() {
|
||||
using AtomLayout_MNK = Layout<ClusterShape_MNK>;
|
||||
constexpr int M = cute::size<0>(TileShape_MNK{});
|
||||
static_assert(M == 128 || M == 256, "Invalid TileShape_M.");
|
||||
|
||||
constexpr int N = cute::size<1>(TileShape_MNK{});
|
||||
static_assert(N % 64 == 0 && N <= 256, "Invalid TileShape_N.");
|
||||
|
||||
if constexpr (cute::is_same_v<ElementAMma, cutlass::float_e2m1_t> ||
|
||||
cute::is_same_v<ElementAMma, cutlass::type_erased_dynamic_float4_t>) {
|
||||
return make_tiled_mma(cute::SM103::SM103_MXF4_ULTRA_SS_VS<ElementAMma, ElementBMma, ElementAccumulator, ElementSF,
|
||||
M, N, SFVectorSize, UmmaMajorA, UmmaMajorB>{});
|
||||
}
|
||||
else {
|
||||
static_assert(cutlass::detail::dependent_false<ElementAMma>,
|
||||
"Unsupported configuration for SM103 collective builder.");
|
||||
}
|
||||
}
|
||||
|
||||
template<
|
||||
class ElementAMma,
|
||||
class ElementBMma,
|
||||
class ElementAccumulator,
|
||||
class ElementSF,
|
||||
class TileShape_MNK,
|
||||
class ClusterShape_MNK,
|
||||
UMMA::Major UmmaMajorA,
|
||||
UMMA::Major UmmaMajorB,
|
||||
int SFVectorSize
|
||||
>
|
||||
constexpr auto
|
||||
sm103_make_blockscaled_2sm_tiled_mma() {
|
||||
using AtomLayout_MNK = Layout<decltype(shape_div(ClusterShape_MNK{}, Shape<_2,_1,_1>{}))>;
|
||||
|
||||
constexpr int M = cute::size<0>(TileShape_MNK{});
|
||||
static_assert(M == 128 || M == 256, "Invalid TileShape_M.");
|
||||
|
||||
constexpr int N = cute::size<1>(TileShape_MNK{});
|
||||
static_assert(N % 64 == 0 && N <= 256, "Invalid TileShape_N.");
|
||||
|
||||
if constexpr (cute::is_same_v<ElementAMma, cutlass::float_e2m1_t> ||
|
||||
cute::is_same_v<ElementAMma, cutlass::type_erased_dynamic_float4_t>) {
|
||||
return make_tiled_mma(cute::SM103::SM103_MXF4_ULTRA_2x1SM_SS_VS<ElementAMma, ElementBMma, ElementAccumulator, ElementSF,
|
||||
M, N, SFVectorSize, UmmaMajorA, UmmaMajorB>{});
|
||||
}
|
||||
else {
|
||||
static_assert(cutlass::detail::dependent_false<ElementAMma>,
|
||||
"Unsupported configuration for SM103 collective builder.");
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
template<
|
||||
class ElementAMma,
|
||||
class ElementBMma,
|
||||
class ElementAccumulator,
|
||||
class ElementSF,
|
||||
class ClusterTileShape_MNK,
|
||||
class ClusterShape_MNK,
|
||||
UMMA::Major UmmaMajorA,
|
||||
UMMA::Major UmmaMajorB,
|
||||
class BuilderScheduleTag
|
||||
>
|
||||
constexpr auto
|
||||
sm103_make_blockscaled_tiled_mma() {
|
||||
constexpr uint32_t SFVectorSize = find_vector_size<BuilderScheduleTag>();
|
||||
|
||||
// MMA_2SM requested
|
||||
if constexpr (cute::is_base_of_v<KernelSchedule2Sm, BuilderScheduleTag>) {
|
||||
return sm103_make_blockscaled_2sm_tiled_mma<ElementAMma, ElementBMma, ElementAccumulator, ElementSF,
|
||||
ClusterTileShape_MNK, ClusterShape_MNK, UmmaMajorA, UmmaMajorB, SFVectorSize>();
|
||||
}
|
||||
// MMA_1SM requested
|
||||
else if constexpr (cute::is_base_of_v<KernelSchedule1Sm, BuilderScheduleTag>) {
|
||||
return sm103_make_blockscaled_1sm_tiled_mma<ElementAMma, ElementBMma, ElementAccumulator, ElementSF,
|
||||
ClusterTileShape_MNK, ClusterShape_MNK, UmmaMajorA, UmmaMajorB, SFVectorSize>();
|
||||
}
|
||||
// Auto scheduling requested
|
||||
else if constexpr (cute::is_same_v<BuilderScheduleTag, KernelScheduleAuto>) {
|
||||
if constexpr (cute::get<0>(ClusterShape_MNK{}) % 2 == 0) {
|
||||
return sm103_make_blockscaled_2sm_tiled_mma<ElementAMma, ElementBMma, ElementAccumulator, ElementSF,
|
||||
ClusterTileShape_MNK, ClusterShape_MNK, UmmaMajorA, UmmaMajorB, SFVectorSize>();
|
||||
}
|
||||
else {
|
||||
return sm103_make_blockscaled_1sm_tiled_mma<ElementAMma, ElementBMma, ElementAccumulator, ElementSF,
|
||||
ClusterTileShape_MNK, ClusterShape_MNK, UmmaMajorA, UmmaMajorB, SFVectorSize>();
|
||||
}
|
||||
}
|
||||
else {
|
||||
static_assert(cutlass::detail::dependent_false<ElementAMma>,
|
||||
"Unsupported policy for SM103 collective builder.");
|
||||
}
|
||||
}
|
||||
|
||||
template<
|
||||
class ElementAMma,
|
||||
class ElementBMma,
|
||||
class ElementAccumulator,
|
||||
class ElementSF,
|
||||
class TileShape_MNK,
|
||||
class ClusterShape_MNK,
|
||||
UMMA::Major UmmaMajorA,
|
||||
UMMA::Major UmmaMajorB,
|
||||
uint32_t SFVectorSize,
|
||||
class BuilderScheduleTag,
|
||||
bool Is2SM
|
||||
>
|
||||
struct Sm103TrivialBlockscaledMma {};
|
||||
|
||||
template<
|
||||
class ElementAMma,
|
||||
class ElementBMma,
|
||||
class ElementAccumulator,
|
||||
class ElementSF,
|
||||
class TileShape_MNK,
|
||||
class ClusterShape_MNK,
|
||||
UMMA::Major UmmaMajorA,
|
||||
UMMA::Major UmmaMajorB,
|
||||
uint32_t SFVectorSize,
|
||||
class BuilderScheduleTag
|
||||
>
|
||||
struct Sm103TrivialBlockscaledMma< ElementAMma,
|
||||
ElementBMma,
|
||||
ElementAccumulator,
|
||||
ElementSF,
|
||||
TileShape_MNK,
|
||||
ClusterShape_MNK,
|
||||
UmmaMajorA,
|
||||
UmmaMajorB,
|
||||
SFVectorSize,
|
||||
BuilderScheduleTag,
|
||||
true /*Is2SM*/> {
|
||||
using type = decltype(sm103_make_blockscaled_2sm_tiled_mma<ElementAMma, ElementBMma, ElementAccumulator, ElementSF,
|
||||
TileShape_MNK, ClusterShape_MNK, UmmaMajorA, UmmaMajorB, SFVectorSize>());
|
||||
};
|
||||
|
||||
template<
|
||||
class ElementAMma,
|
||||
class ElementBMma,
|
||||
class ElementAccumulator,
|
||||
class ElementSF,
|
||||
class TileShape_MNK,
|
||||
class ClusterShape_MNK,
|
||||
UMMA::Major UmmaMajorA,
|
||||
UMMA::Major UmmaMajorB,
|
||||
uint32_t SFVectorSize,
|
||||
class BuilderScheduleTag
|
||||
>
|
||||
struct Sm103TrivialBlockscaledMma< ElementAMma,
|
||||
ElementBMma,
|
||||
ElementAccumulator,
|
||||
ElementSF,
|
||||
TileShape_MNK,
|
||||
ClusterShape_MNK,
|
||||
UmmaMajorA,
|
||||
UmmaMajorB,
|
||||
SFVectorSize,
|
||||
BuilderScheduleTag,
|
||||
false /*Is2SM*/> {
|
||||
using type = decltype(sm103_make_blockscaled_1sm_tiled_mma<ElementAMma, ElementBMma, ElementAccumulator, ElementSF,
|
||||
TileShape_MNK, ClusterShape_MNK, UmmaMajorA, UmmaMajorB, SFVectorSize>());
|
||||
};
|
||||
|
||||
template <class ElementType>
|
||||
CUTLASS_HOST_DEVICE
|
||||
static constexpr bool
|
||||
is_sm103_block_scale_input() {
|
||||
// Allowed input element datatype for block-scaling GEMM
|
||||
return ( cute::is_same_v<ElementType, cutlass::type_erased_dynamic_float4_t> ||
|
||||
cute::is_same_v<ElementType, cutlass::float_e2m1_t>);
|
||||
}
|
||||
|
||||
template <class BlockScaleConfig, class MmaShapeA_MK>
|
||||
constexpr
|
||||
auto sm103_sfa_smem_atom_layout() {
|
||||
constexpr int SF_BUFFERS_PER_TILE_K = BlockScaleConfig::SFVecSize == 16 ? 4 : 2;
|
||||
auto mma_sfa_tiler = make_shape(get<0,0>(MmaShapeA_MK{})*get<1>(MmaShapeA_MK{}), get<0,1>(MmaShapeA_MK{}) * get<2>(MmaShapeA_MK{}) / Int<SF_BUFFERS_PER_TILE_K>{});
|
||||
return tiled_product(typename BlockScaleConfig::SfAtom{},
|
||||
make_layout(shape_div(mma_sfa_tiler, product_each(shape(typename BlockScaleConfig::SfAtom{})))));
|
||||
}
|
||||
|
||||
template <class BlockScaleConfig, class MmaShapeB_NK, int SFVecSize>
|
||||
constexpr
|
||||
auto sm103_sfb_smem_atom_layout() {
|
||||
auto sSFB = [&]() {
|
||||
constexpr int MMA_N = get<0>(MmaShapeB_NK{});
|
||||
constexpr int NonPow2N = 192;
|
||||
constexpr int NonPow2N_RoundUp = 256;
|
||||
// If MMA_N is 192, we need to operate at MMA_N = 256 granularity for UTCCP to work for ScaleFactorB.
|
||||
// Both TMA and UTCCP will transfer scale factor B as if we have 256 columns in B matrix.
|
||||
constexpr int MMA_N_SFB = (MMA_N == NonPow2N) ? NonPow2N_RoundUp : MMA_N;
|
||||
constexpr int SF_BUFFERS_PER_TILE_K = BlockScaleConfig::SFVecSize == 16 ? 4 : 2;
|
||||
auto mma_sfb_tiler = make_shape(Int<MMA_N_SFB>{}, get<1>(MmaShapeB_NK{}) / Int<SF_BUFFERS_PER_TILE_K>{});
|
||||
if constexpr(Int<MMA_N>{} == Int<128>{}) {
|
||||
return tiled_product(typename BlockScaleConfig::SfAtom{},
|
||||
make_layout(shape_div(mma_sfb_tiler,product_each(shape(typename BlockScaleConfig::SfAtom{})))));
|
||||
|
||||
}
|
||||
else {
|
||||
using SfKMajorAtom256 = Layout< Shape< Shape<_32,_4, _2>, Shape<Int<SFVecSize>, _4>>,
|
||||
Stride<Stride<_16,_4, Int<get<1>(mma_sfb_tiler)/SFVecSize/4*512>>, Stride< _0, _1>>>;
|
||||
return tiled_product(SfKMajorAtom256{},
|
||||
make_layout(shape_div(mma_sfb_tiler,product_each(shape(SfKMajorAtom256{})))));
|
||||
}
|
||||
}();
|
||||
return sSFB;
|
||||
}
|
||||
|
||||
} // namespace detail
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <
|
||||
class ElementPairA,
|
||||
class GmemLayoutATag,
|
||||
int AlignmentA,
|
||||
class ElementPairB,
|
||||
class GmemLayoutBTag,
|
||||
int AlignmentB,
|
||||
class ElementAccumulator,
|
||||
class TileShape_MNK, // (MmaAtomShapeM, MmaAtomShapeN, TileK)
|
||||
class ClusterShape_MNK, // Static cluster shape or dynamic (int, int, _1)
|
||||
class StageCountType,
|
||||
class BuilderScheduleTag
|
||||
>
|
||||
struct CollectiveBuilder<
|
||||
arch::Sm103,
|
||||
arch::OpClassBlockScaledTensorOp,
|
||||
ElementPairA,
|
||||
GmemLayoutATag,
|
||||
AlignmentA,
|
||||
ElementPairB,
|
||||
GmemLayoutBTag,
|
||||
AlignmentB,
|
||||
ElementAccumulator,
|
||||
TileShape_MNK,
|
||||
ClusterShape_MNK,
|
||||
StageCountType,
|
||||
BuilderScheduleTag,
|
||||
cute::enable_if_t<
|
||||
// Not paired input, Not Complex input
|
||||
(cute::is_tuple_v<ElementPairA> && cute::is_tuple_v<ElementPairB> &&
|
||||
not cute::is_complex_v<ElementPairA> && not cute::is_complex_v<ElementPairB>) &&
|
||||
// Blockscaled Gemm
|
||||
(cute::is_base_of_v<KernelScheduleSm103BlockScaledGemm, BuilderScheduleTag> ||
|
||||
cute::is_base_of_v<KernelSchedulePtrArraySm103BlockScaledGemm, BuilderScheduleTag> ||
|
||||
cute::is_same_v<KernelScheduleAuto, BuilderScheduleTag>) &&
|
||||
// Alignment check
|
||||
detail::sm1xx_blockscaled_gemm_is_aligned<remove_cvref_t<decltype(get<0>(ElementPairA{}))>,
|
||||
AlignmentA,
|
||||
remove_cvref_t<decltype(get<0>(ElementPairB{}))>,
|
||||
AlignmentB,
|
||||
BuilderScheduleTag>()>>
|
||||
{
|
||||
using ElementA = remove_cvref_t<decltype(get<0>(ElementPairA{}))>;
|
||||
using ElementB = remove_cvref_t<decltype(get<0>(ElementPairB{}))>;
|
||||
using ElementSF = remove_cvref_t<decltype(get<1>(ElementPairA{}))>;
|
||||
|
||||
static_assert(cute::is_tuple<ElementPairA>::value, "Expecting ElementPairA to be a tuple.");
|
||||
static_assert(cute::is_tuple<ElementPairB>::value, "Expecting ElementPairB to be a tuple.");
|
||||
|
||||
static_assert(cute::is_static_v<TileShape_MNK>, "TileShape has to be static");
|
||||
static_assert(cute::size<2>(TileShape_MNK{}) == _768{}, "TileShape_K should 768 for MMA kernels");
|
||||
|
||||
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>();
|
||||
static_assert(cutlass::gemm::detail::is_k_major_A<GmemLayoutATag>() && cutlass::gemm::detail::is_k_major_B<GmemLayoutBTag>(), "Only K major inputs are supported");
|
||||
|
||||
static_assert(cutlass::gemm::collective::detail::is_sm103_block_scale_input<ElementA>(), "Incorrect type for A matrix");
|
||||
static_assert(cutlass::gemm::collective::detail::is_sm103_block_scale_input<ElementB>(), "Incorrect type for B matrix");
|
||||
|
||||
static_assert(cute::is_same_v<ElementSF, cutlass::float_ue8m0_t> ||
|
||||
cute::is_same_v<ElementSF, cutlass::float_ue4m3_t>, "Incorrect scale factor type");
|
||||
|
||||
// Data type used by MMA instruction
|
||||
using ElementAMma = decltype(cutlass::gemm::collective::detail::sm1xx_kernel_input_element_to_mma_input_element<ElementA, false /*UseQmma=false*/>());
|
||||
using ElementBMma = decltype(cutlass::gemm::collective::detail::sm1xx_kernel_input_element_to_mma_input_element<ElementB, false /*UseQmma=false*/>());
|
||||
|
||||
static constexpr uint32_t SFVectorSize = detail::find_vector_size<BuilderScheduleTag>();
|
||||
|
||||
static constexpr bool is_2sm = cute::is_base_of_v<KernelSchedule2Sm, BuilderScheduleTag> ||
|
||||
(cute::is_same_v<BuilderScheduleTag, KernelScheduleAuto> &&
|
||||
(cute::is_static_v<ClusterShape_MNK> && cute::get<0>(ClusterShape_MNK{}) % 2 == 0));
|
||||
|
||||
using TiledMma = typename cutlass::gemm::collective::detail::Sm103TrivialBlockscaledMma<ElementAMma, ElementBMma, ElementAccumulator, ElementSF,
|
||||
TileShape_MNK, ClusterShape_MNK,
|
||||
UmmaMajorA, UmmaMajorB, SFVectorSize, BuilderScheduleTag, is_2sm>::type;
|
||||
|
||||
using AtomThrID = typename TiledMma::AtomThrID;
|
||||
using Sm1xxBlkScaledConfig = cutlass::detail::Sm103BlockScaledConfig<SFVectorSize>;
|
||||
|
||||
using ElementAMma_SmemAllocType = uint8_t;
|
||||
// ElementAMma;
|
||||
using ElementBMma_SmemAllocType = uint8_t;
|
||||
// ElementBMma;
|
||||
|
||||
// ((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 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 GmemTiledCopySFA = decltype(cutlass::gemm::collective::detail::sm100_cluster_shape_to_tma_atom_A(
|
||||
ClusterShape_MNK{}, AtomThrID{}));
|
||||
|
||||
using GmemTiledCopySFB = decltype(cutlass::gemm::collective::detail::sm100_cluster_shape_to_tma_atom_SFB(
|
||||
ClusterShape_MNK{}, AtomThrID{}));
|
||||
|
||||
using GmemTiledCopyPairA = decltype(cute::make_tuple(GmemTiledCopyA{}, GmemTiledCopySFA{}));
|
||||
using GmemTiledCopyPairB = decltype(cute::make_tuple(GmemTiledCopyB{}, GmemTiledCopySFB{}));
|
||||
|
||||
//
|
||||
// Construct SMEM layout (SmemLayoutAtom) for A and SFA
|
||||
//
|
||||
using SmemLayoutAtomA = UMMA::Layout_K_SW128_Atom<ElementAMma_SmemAllocType>;
|
||||
// A single indivisible block will hold 4 scale factors of 128 rows/columns (A/B matrix).
|
||||
// 4 is chosen to make consecutive 32bits of data to have scale factors for only a single row (col). 32bits corresponds to the TMEM word size
|
||||
static constexpr int MMA_M = cute::size<0>(TileShape_MNK{}) / cute::size(AtomThrID{});
|
||||
using SmemLayoutAtomSFA = decltype(detail::sm103_sfa_smem_atom_layout<Sm1xxBlkScaledConfig,MmaShapeA_MK>());
|
||||
using SmemLayoutAtomsA = decltype(cute::make_tuple(SmemLayoutAtomA{}, SmemLayoutAtomSFA{}));
|
||||
|
||||
//
|
||||
// Construct SMEM layout（SmemLayoutAtom）for B and SFB
|
||||
//
|
||||
|
||||
using SmemLayoutAtomB = UMMA::Layout_K_SW128_Atom<ElementBMma_SmemAllocType>;
|
||||
static constexpr int MMA_N = cute::size<1>(TileShape_MNK{});
|
||||
// If MMA_N is 192, we need to operate at MMA_N = 256 granularity for UTCCP to work for ScaleFactorB.
|
||||
// Both TMA and UTCCP will transfer scale factor B as if we have 256 columns in B matrix.
|
||||
using SmemLayoutAtomSFB = decltype(detail::sm103_sfb_smem_atom_layout<Sm1xxBlkScaledConfig,decltype(select<1,2>(TileShape_MNK{})),SFVectorSize>());
|
||||
using SmemLayoutAtomsB = decltype(cute::make_tuple(SmemLayoutAtomB{}, SmemLayoutAtomSFB{}));
|
||||
|
||||
//
|
||||
// Construct Strides for A, SFA, B, and SFB
|
||||
//
|
||||
using StrideA = cutlass::gemm::TagToStrideA_t<GmemLayoutATag>;
|
||||
using StrideB = cutlass::gemm::TagToStrideB_t<GmemLayoutBTag>;
|
||||
using InternalStrideA = cute::remove_pointer_t<StrideA>;
|
||||
using InternalStrideB = cute::remove_pointer_t<StrideB>;
|
||||
using InternalLayoutSFA = decltype(Sm1xxBlkScaledConfig::deduce_layoutSFA());
|
||||
using InternalLayoutSFB = decltype(Sm1xxBlkScaledConfig::deduce_layoutSFB());
|
||||
using LayoutSFA = cute::conditional_t<cute::is_same_v<InternalStrideA, StrideA>, InternalLayoutSFA, InternalLayoutSFA *>;
|
||||
using LayoutSFB = cute::conditional_t<cute::is_same_v<InternalStrideB, StrideB>, InternalLayoutSFB, InternalLayoutSFB *>;
|
||||
using StridePairA = decltype(cute::make_tuple(StrideA{}, LayoutSFA{}));
|
||||
using StridePairB = decltype(cute::make_tuple(StrideB{}, LayoutSFB{}));
|
||||
|
||||
//
|
||||
// Others
|
||||
//
|
||||
|
||||
static constexpr cutlass::sm103::detail::KernelPrefetchType PrefetchType = cute::is_base_of_v<KernelScheduleSm103BlockScaledMxNvf4UltraDisablePrefetch, BuilderScheduleTag>
|
||||
|| cute::is_base_of_v<KernelSchedulePtrArraySm103BlockScaledMxNvf4UltraDisablePrefetch, BuilderScheduleTag>
|
||||
? cutlass::sm103::detail::KernelPrefetchType::Disable :
|
||||
cutlass::sm103::detail::KernelPrefetchType::TmaPrefetch;
|
||||
|
||||
static constexpr uint32_t AccumulatorPipelineStageCount = (MMA_N == 256) ? 1 : 2;
|
||||
static constexpr uint32_t SchedulerPipelineStageCount = 3;
|
||||
|
||||
// AccumulatorPipeline = PipelineUmmaAsync
|
||||
static constexpr auto AccumulatorPipelineStorage = sizeof(typename cutlass::PipelineUmmaAsync<AccumulatorPipelineStageCount>::SharedStorage);
|
||||
// CLCPipeline = PipelineCLCFetchAsync
|
||||
static constexpr auto CLCPipelineStorage = sizeof(typename cutlass::PipelineCLCFetchAsync<SchedulerPipelineStageCount>::SharedStorage);
|
||||
// LoadOrderBarrier = OrderedSequenceBarrier<1,2>
|
||||
static constexpr auto LoadOrderBarrierStorage = sizeof(typename cutlass::OrderedSequenceBarrier<1,2>::SharedStorage);
|
||||
// CLC (scheduler) response
|
||||
static constexpr auto CLCResponseStorage = SchedulerPipelineStageCount * detail::CLCResponseSize;
|
||||
// CLC Throttle pipeline storage
|
||||
static constexpr auto CLCThrottlePipelineStorage = sizeof(typename cutlass::PipelineAsync<SchedulerPipelineStageCount>::SharedStorage);
|
||||
// Tmem dealloc
|
||||
static constexpr auto TmemDeallocStorage = sizeof(cutlass::arch::ClusterBarrier);
|
||||
// Tmem ptr storage
|
||||
static constexpr auto TmemBasePtrsStorage = AccumulatorPipelineStageCount * sizeof(uint32_t);
|
||||
// Tensormap Storage
|
||||
static constexpr bool IsArrayOfPointersGemm = cute::is_base_of_v<KernelSchedulePtrArraySm103BlockScaledGemm, BuilderScheduleTag>;
|
||||
static constexpr auto TensorMapStorage = IsArrayOfPointersGemm ? sizeof(cute::TmaDescriptor) * 4 /* for A, B, SFA and SFB */ : 0;
|
||||
// TMA Load Prefetch Storage
|
||||
static constexpr auto TmaPrefetchStorage = 0;
|
||||
// Smem usage that's not part of CollectiveEpilogue::SharedStorage & CollectiveMainloop::SharedStorage
|
||||
static constexpr auto KernelSmemCarveout = static_cast<int>( AccumulatorPipelineStorage +
|
||||
CLCPipelineStorage +
|
||||
LoadOrderBarrierStorage +
|
||||
CLCResponseStorage +
|
||||
CLCThrottlePipelineStorage +
|
||||
TmemDeallocStorage +
|
||||
TmemBasePtrsStorage +
|
||||
TensorMapStorage +
|
||||
TmaPrefetchStorage);
|
||||
// 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<Int<MMA_M>, Int<MMA_N/cute::size(AtomThrID{})>, _128>; // SmemAllocTypes are uint8_t. We always allocate 128bytes
|
||||
static constexpr auto PipelineStages = cutlass::gemm::collective::detail::sm103_compute_stage_count_or_override_blockscaled<
|
||||
Sm100ReducedSmemCapacityBytes, ElementAMma_SmemAllocType, ElementBMma_SmemAllocType, SmemTileShape, SmemLayoutAtomSFA, SmemLayoutAtomSFB>(StageCountType{});
|
||||
|
||||
using DispatchPolicy = typename cute::conditional_t<IsArrayOfPointersGemm,
|
||||
cutlass::gemm::MainloopSm103ArrayTmaUmmaWarpSpecializedBlockScaled<
|
||||
get<0>(PipelineStages),
|
||||
get<1>(PipelineStages),
|
||||
SchedulerPipelineStageCount,
|
||||
AccumulatorPipelineStageCount,
|
||||
ClusterShape_MNK,
|
||||
PrefetchType
|
||||
>,
|
||||
cutlass::gemm::MainloopSm103TmaUmmaWarpSpecializedBlockScaled<
|
||||
get<0>(PipelineStages),
|
||||
get<1>(PipelineStages),
|
||||
SchedulerPipelineStageCount,
|
||||
AccumulatorPipelineStageCount,
|
||||
ClusterShape_MNK,
|
||||
PrefetchType
|
||||
>
|
||||
>;
|
||||
|
||||
using CollectiveOp = cutlass::gemm::collective::CollectiveMma<
|
||||
DispatchPolicy,
|
||||
TileShape_MNK,
|
||||
ElementPairA,
|
||||
StridePairA,
|
||||
ElementPairB,
|
||||
StridePairB,
|
||||
TiledMma,
|
||||
GmemTiledCopyPairA,
|
||||
SmemLayoutAtomsA,
|
||||
void,
|
||||
cute::identity,
|
||||
GmemTiledCopyPairB,
|
||||
SmemLayoutAtomsB,
|
||||
void,
|
||||
cute::identity
|
||||
>;
|
||||
};
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace cutlass::gemm::collective
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
@@ -141,6 +141,18 @@ constexpr uint32_t find_vector_size() {
|
||||
cute::is_same_v<BuilderScheduleTag, KernelPtrArrayTmaWarpSpecialized2SmNvf4Sm100> ||
|
||||
cute::is_same_v<BuilderScheduleTag, KernelTmaWarpSpecializedNvf4Sm120> ||
|
||||
cute::is_same_v<BuilderScheduleTag, KernelTmaWarpSpecializedPingpongNvf4Sm120>
|
||||
|| cute::is_same_v<BuilderScheduleTag, KernelTmaWarpSpecialized1SmBlockScaledMxNvf4UltraVs16Sm103>
|
||||
|| cute::is_same_v<BuilderScheduleTag, KernelTmaWarpSpecialized2SmBlockScaledMxNvf4UltraVs16Sm103>
|
||||
|| cute::is_same_v<BuilderScheduleTag, KernelTmaWarpSpecialized1SmBlockScaledMxNvf4UltraVs16Sm103DisablePrefetch>
|
||||
|| cute::is_same_v<BuilderScheduleTag, KernelTmaWarpSpecialized2SmBlockScaledMxNvf4UltraVs16Sm103DisablePrefetch>
|
||||
|| cute::is_same_v<BuilderScheduleTag, KernelTmaWarpSpecialized1SmBlockScaledMxNvf4UltraVs16Sm103TmaPrefetch>
|
||||
|| cute::is_same_v<BuilderScheduleTag, KernelTmaWarpSpecialized2SmBlockScaledMxNvf4UltraVs16Sm103TmaPrefetch>
|
||||
|| cute::is_same_v<BuilderScheduleTag, KernelPtrArrayTmaWarpSpecialized1SmBlockScaledMxNvf4UltraVs16Sm103>
|
||||
|| cute::is_same_v<BuilderScheduleTag, KernelPtrArrayTmaWarpSpecialized2SmBlockScaledMxNvf4UltraVs16Sm103>
|
||||
|| cute::is_same_v<BuilderScheduleTag, KernelPtrArrayTmaWarpSpecialized1SmBlockScaledMxNvf4UltraVs16Sm103TmaPrefetch>
|
||||
|| cute::is_same_v<BuilderScheduleTag, KernelPtrArrayTmaWarpSpecialized2SmBlockScaledMxNvf4UltraVs16Sm103TmaPrefetch>
|
||||
|| cute::is_same_v<BuilderScheduleTag, KernelPtrArrayTmaWarpSpecialized1SmBlockScaledMxNvf4UltraVs16Sm103DisablePrefetch>
|
||||
|| cute::is_same_v<BuilderScheduleTag, KernelPtrArrayTmaWarpSpecialized2SmBlockScaledMxNvf4UltraVs16Sm103DisablePrefetch>
|
||||
) {
|
||||
return 16;
|
||||
}
|
||||
|
||||
@@ -47,6 +47,9 @@
|
||||
#include "cutlass/gemm/collective/builders/sm100_blockwise_umma_builder.inl"
|
||||
#include "cutlass/gemm/collective/builders/sm100_blockscaled_sparse_umma_builder.inl"
|
||||
#include "cutlass/gemm/collective/builders/sm100_simt_builder.inl"
|
||||
#include "cutlass/gemm/collective/builders/sm100_mixed_input_umma_builder.inl"
|
||||
#include "cutlass/gemm/collective/builders/sm100_cpasync_umma_builder.inl"
|
||||
#include "cutlass/gemm/collective/builders/sm103_blockscaled_umma_builder.inl"
|
||||
#include "cutlass/gemm/collective/builders/sm120_mma_builder.inl"
|
||||
#include "cutlass/gemm/collective/builders/sm120_blockscaled_mma_builder.inl"
|
||||
#include "cutlass/gemm/collective/builders/sm120_sparse_mma_builder.inl"
|
||||
|
||||
@@ -63,6 +63,10 @@
|
||||
#include "cutlass/gemm/collective/sm100_blockscaled_mma_array_warpspecialized.hpp"
|
||||
#include "cutlass/gemm/collective/sm100_mma_warpspecialized_blockwise_scaling.hpp"
|
||||
#include "cutlass/gemm/collective/sm100_mma_array_warpspecialized_blockwise_scaling.hpp"
|
||||
#include "cutlass/gemm/collective/sm100_mma_warpspecialized_mixed_input.hpp"
|
||||
#include "cutlass/gemm/collective/sm100_mma_cpasync_warpspecialized.hpp"
|
||||
#include "cutlass/gemm/collective/sm103_blockscaled_mma_warpspecialized.hpp"
|
||||
#include "cutlass/gemm/collective/sm103_blockscaled_mma_array_warpspecialized.hpp"
|
||||
#include "cutlass/gemm/collective/sm120_mma_tma.hpp"
|
||||
#include "cutlass/gemm/collective/sm120_blockscaled_mma_tma.hpp"
|
||||
#include "cutlass/gemm/collective/sm120_blockscaled_mma_array_tma.hpp"
|
||||
|
||||
@@ -0,0 +1,588 @@
|
||||
/***************************************************************************************************
|
||||
* Copyright (c) 2024 - 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/cutlass.h"
|
||||
#include "cutlass/detail/cluster.hpp"
|
||||
#include "cutlass/gemm/dispatch_policy.hpp"
|
||||
#include "cutlass/numeric_types.h"
|
||||
#include "cutlass/pipeline/pipeline.hpp"
|
||||
#include "cutlass/gemm/gemm.h"
|
||||
#include "cutlass/trace.h"
|
||||
#include "cutlass/kernel_hardware_info.hpp"
|
||||
#include "cutlass/arch/memory.h"
|
||||
|
||||
#include "cute/algorithm/functional.hpp"
|
||||
#include "cute/arch/cluster_sm90.hpp"
|
||||
#include "cute/atom/mma_atom.hpp"
|
||||
#include "cute/algorithm/gemm.hpp"
|
||||
#include "cute/numeric/arithmetic_tuple.hpp"
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace cutlass::gemm::collective {
|
||||
using namespace cute;
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
// WarpSpecialized Mainloop
|
||||
// Both DMA Load and MMA methods of this class must be run by a single thread that's picked by elect_one
|
||||
template <
|
||||
int Stages,
|
||||
int SchedulerPipelineStageCount,
|
||||
int AccumulatorPipelineStageCount,
|
||||
class ClusterShape, // Static cluster shape or dynamic (int, int, _1)
|
||||
class TileShape_, // (MmaAtomShapeM, MmaAtomShapeN, TileK)
|
||||
class ElementA_,
|
||||
class StrideA_,
|
||||
class ElementB_,
|
||||
class StrideB_,
|
||||
class TiledMma_,
|
||||
class GmemTiledCopyA_,
|
||||
class SmemLayoutAtomA_,
|
||||
class SmemCopyAtomA_,
|
||||
class TransformA_,
|
||||
class GmemTiledCopyB_,
|
||||
class SmemLayoutAtomB_,
|
||||
class SmemCopyAtomB_,
|
||||
class TransformB_>
|
||||
struct CollectiveMma<
|
||||
MainloopSm100UmmaCpAsyncWarpSpecialized<
|
||||
Stages,
|
||||
SchedulerPipelineStageCount,
|
||||
AccumulatorPipelineStageCount,
|
||||
ClusterShape>,
|
||||
TileShape_,
|
||||
ElementA_,
|
||||
StrideA_,
|
||||
ElementB_,
|
||||
StrideB_,
|
||||
TiledMma_,
|
||||
GmemTiledCopyA_,
|
||||
SmemLayoutAtomA_,
|
||||
SmemCopyAtomA_,
|
||||
TransformA_,
|
||||
GmemTiledCopyB_,
|
||||
SmemLayoutAtomB_,
|
||||
SmemCopyAtomB_,
|
||||
TransformB_>
|
||||
{
|
||||
using TiledMma = TiledMma_;
|
||||
using AtomThrShapeMNK = Shape<decltype(shape<0>(typename TiledMma::ThrLayoutVMNK{})), _1, _1>;
|
||||
|
||||
// Statically asserting to ensure only 1x1x1 cluster shape & 1sm setup is received
|
||||
static_assert(size(AtomThrShapeMNK{}) == 1, "Lower alignment SM100 GEMM only supports 1SM MMA");
|
||||
static_assert(size(ClusterShape{}) == 1, "CPASYNC does not support multicast so the cluster shape is restricted to 1, 1, 1");
|
||||
|
||||
using DispatchPolicy = MainloopSm100UmmaCpAsyncWarpSpecialized<
|
||||
Stages,
|
||||
SchedulerPipelineStageCount,
|
||||
AccumulatorPipelineStageCount,
|
||||
ClusterShape>;
|
||||
// TileShape refers to MmaTileShape to adapt for runtime cluster shape
|
||||
using TileShape = TileShape_;
|
||||
|
||||
CUTE_STATIC_ASSERT_V(evenly_divides(TileShape{}, tile_shape(TiledMma{})),
|
||||
"Static cluster shape used: TileShape should be evenly divided by TiledMma");
|
||||
|
||||
// Define A and B block shapes
|
||||
using MmaShapeA_MK = decltype(partition_shape_A(TiledMma{}, make_shape(size<0>(TileShape{}), size<2>(TileShape{}))));
|
||||
using MmaShapeB_NK = decltype(partition_shape_B(TiledMma{}, make_shape(size<1>(TileShape{}), size<2>(TileShape{}))));
|
||||
using LoadShapeA_MK = decltype(select<0,2>(TileShape{}));
|
||||
using LoadShapeB_NK = decltype(select<1,2>(TileShape{}));
|
||||
|
||||
// CtaShape_MNK is queried from collective in all kernel layers
|
||||
using CtaShape_MNK = TileShape;
|
||||
|
||||
using ElementA = ElementA_;
|
||||
using ElementAMma = typename TiledMma::ValTypeA;
|
||||
using StrideA = StrideA_;
|
||||
using ElementB = ElementB_;
|
||||
using ElementBMma = typename TiledMma::ValTypeB;
|
||||
using StrideB = StrideB_;
|
||||
|
||||
static constexpr bool IsRuntimeDataTypeA = cute::is_same_v<ElementA, cutlass::type_erased_dynamic_float8_t>;
|
||||
static constexpr bool IsRuntimeDataTypeB = cute::is_same_v<ElementB, cutlass::type_erased_dynamic_float8_t>;
|
||||
|
||||
static_assert(IsRuntimeDataTypeA == IsRuntimeDataTypeB,
|
||||
"ElementA and ElementB should be both runtime or both static.");
|
||||
|
||||
static constexpr bool IsRuntimeDataType = IsRuntimeDataTypeA && IsRuntimeDataTypeB;
|
||||
|
||||
|
||||
using ElementAccumulator = typename TiledMma::ValTypeC;
|
||||
using GmemTiledCopyA = GmemTiledCopyA_;
|
||||
using GmemTiledCopyB = GmemTiledCopyB_;
|
||||
using SmemLayoutAtomA = SmemLayoutAtomA_;
|
||||
using SmemLayoutAtomB = SmemLayoutAtomB_;
|
||||
using SmemCopyAtomA = SmemCopyAtomA_;
|
||||
using SmemCopyAtomB = SmemCopyAtomB_;
|
||||
using TransformA = TransformA_;
|
||||
using TransformB = TransformB_;
|
||||
using ArchTag = typename DispatchPolicy::ArchTag;
|
||||
|
||||
using MainloopPipeline = cutlass::PipelineUmmaConsumerAsync<DispatchPolicy::Stages, AtomThrShapeMNK>;
|
||||
using MainloopPipelineState = typename MainloopPipeline::PipelineState;
|
||||
|
||||
static_assert(size(GmemTiledCopyA{}) == size(GmemTiledCopyB{}), "A and B GmemTiledCopy should share the same thread count");
|
||||
static constexpr int NumLoadThreads = size(GmemTiledCopyA{});
|
||||
|
||||
static_assert(rank(SmemLayoutAtomA{}) == 2, "SmemLayoutAtomA must be rank 2 (M,K)");
|
||||
static_assert(((size<0,0>(MmaShapeA_MK{}) * size<1>(MmaShapeA_MK{})) % size<0>(SmemLayoutAtomA{})) == 0,
|
||||
"SmemLayoutAtom must evenly divide tile shape.");
|
||||
static_assert(((size<0,1>(MmaShapeA_MK{}) * size<2>(MmaShapeA_MK{})) % size<1>(SmemLayoutAtomA{})) == 0,
|
||||
"SmemLayoutAtom must evenly divide tile shape.");
|
||||
static_assert(cute::is_void_v<SmemCopyAtomA>,
|
||||
"SM100 UMMA cannot have a non-void copy atom for smem sourced instructions.");
|
||||
|
||||
static_assert(rank(SmemLayoutAtomB{}) == 2, "SmemLayoutAtomB must be rank 2 (N,K)");
|
||||
static_assert(((size<0,0>(MmaShapeB_NK{}) * size<1>(MmaShapeB_NK{})) % size<0>(SmemLayoutAtomB{})) == 0,
|
||||
"SmemLayoutAtom must evenly divide tile shape.");
|
||||
static_assert(((size<0,1>(MmaShapeB_NK{}) * size<2>(MmaShapeB_NK{})) % size<1>(SmemLayoutAtomB{})) == 0,
|
||||
"SmemLayoutAtom must evenly divide tile shape.");
|
||||
static_assert(cute::is_void_v<SmemCopyAtomB>,
|
||||
"SM100 UMMA cannot have a non-void copy atom for smem sourced instructions.");
|
||||
|
||||
// Tile along K mode first before tiling over MN. PIPE mode last as usual.
|
||||
// (MMA_TILE_M,MMA_TILE_K),MMA_M,MMA_K,PIPE)
|
||||
using MmaSmemLayoutA = decltype(UMMA::tile_to_mma_shape(
|
||||
SmemLayoutAtomA{},
|
||||
append(MmaShapeA_MK{}, Int<DispatchPolicy::Stages>{}),
|
||||
conditional_t< ::cutlass::gemm::detail::is_major<0,StrideA>(), Step<_2,_1,_3>, Step<_1,_2,_3>>{}));
|
||||
|
||||
using LoadSmemLayoutA = decltype(tile_to_shape(
|
||||
SmemLayoutAtomA{},
|
||||
append(LoadShapeA_MK{}, Int<DispatchPolicy::Stages>{}),
|
||||
conditional_t< ::cutlass::gemm::detail::is_major<0,StrideA>(), Step<_2,_1,_3>, Step<_1,_2,_3>>{}));
|
||||
|
||||
using MmaSmemLayoutB = decltype(UMMA::tile_to_mma_shape(
|
||||
SmemLayoutAtomB{},
|
||||
append(MmaShapeB_NK{}, Int<DispatchPolicy::Stages>{}),
|
||||
conditional_t< ::cutlass::gemm::detail::is_major<0,StrideB>(), Step<_2,_1,_3>, Step<_1,_2,_3>>{}));
|
||||
|
||||
using LoadSmemLayoutB = decltype(tile_to_shape(
|
||||
SmemLayoutAtomB{},
|
||||
append(LoadShapeB_NK{}, Int<DispatchPolicy::Stages>{}),
|
||||
conditional_t< ::cutlass::gemm::detail::is_major<0,StrideB>(), Step<_2,_1,_3>, Step<_1,_2,_3>>{}));
|
||||
|
||||
|
||||
static_assert(DispatchPolicy::Stages >= 2, "Specialization requires Stages set to value 1 or more.");
|
||||
static_assert(cute::is_base_of<cute::UMMA::DescriptorIterator, typename TiledMma::FrgTypeA>::value &&
|
||||
cute::is_base_of<cute::UMMA::DescriptorIterator, typename TiledMma::FrgTypeB>::value,
|
||||
"MMA atom must source both A and B operand from smem_desc for this mainloop.");
|
||||
|
||||
using SmemAllocTypeA = cute::conditional_t<cute::sizeof_bits_v<ElementAMma> < 8, uint8_t, ElementAMma>;
|
||||
using SmemAllocTypeB = cute::conditional_t<cute::sizeof_bits_v<ElementBMma> < 8, uint8_t, ElementBMma>;
|
||||
|
||||
using BitTypeElementA = cute::uint_bit_t<cute::sizeof_bits_v<ElementA>>;
|
||||
using BitTypeElementB = cute::uint_bit_t<cute::sizeof_bits_v<ElementB>>;
|
||||
|
||||
using ArrayElementA = cute::conditional_t<IsRuntimeDataTypeA, BitTypeElementA, ElementA>;
|
||||
using ArrayElementB = cute::conditional_t<IsRuntimeDataTypeB, BitTypeElementB, ElementB>;
|
||||
|
||||
using RuntimeDataTypeA = cute::conditional_t<IsRuntimeDataTypeA, cute::UMMA::MXF8F6F4Format, void*>;
|
||||
using RuntimeDataTypeB = cute::conditional_t<IsRuntimeDataTypeB, cute::UMMA::MXF8F6F4Format, void*>;
|
||||
|
||||
struct SharedStorage {
|
||||
struct TensorStorage : cute::aligned_struct<128, _0> {
|
||||
cute::array_aligned<SmemAllocTypeA, cute::cosize_v<LoadSmemLayoutA>> smem_A;
|
||||
cute::array_aligned<SmemAllocTypeB, cute::cosize_v<LoadSmemLayoutB>> smem_B;
|
||||
} tensors;
|
||||
|
||||
using PipelineStorage = typename MainloopPipeline::SharedStorage;
|
||||
PipelineStorage pipeline;
|
||||
};
|
||||
|
||||
// Expose shared storage for tensors/pipelines separately to allow kernel layer to reorder them.
|
||||
using TensorStorage = typename SharedStorage::TensorStorage;
|
||||
using PipelineStorage = typename SharedStorage::PipelineStorage;
|
||||
|
||||
// Host side kernel arguments
|
||||
struct Arguments {
|
||||
ArrayElementA const* ptr_A{nullptr};
|
||||
StrideA dA{};
|
||||
ArrayElementB const* ptr_B{nullptr};
|
||||
StrideB dB{};
|
||||
RuntimeDataTypeA runtime_data_type_a{};
|
||||
RuntimeDataTypeB runtime_data_type_b{};
|
||||
};
|
||||
|
||||
// Device side kernel params
|
||||
struct Params {
|
||||
ArrayElementA const* ptr_A{nullptr};
|
||||
StrideA dA{};
|
||||
ArrayElementB const* ptr_B{nullptr};
|
||||
StrideB dB{};
|
||||
RuntimeDataTypeA runtime_data_type_a;
|
||||
RuntimeDataTypeB runtime_data_type_b;
|
||||
};
|
||||
|
||||
template <class ProblemShape>
|
||||
static constexpr Params
|
||||
to_underlying_arguments(
|
||||
ProblemShape const& problem_shape,
|
||||
Arguments const& args,
|
||||
[[maybe_unused]] void* workspace,
|
||||
cutlass::KernelHardwareInfo const& hw_info = cutlass::KernelHardwareInfo{}) {
|
||||
// Optionally append 1s until problem shape is rank-4 (MNKL), in case it is only rank-3 (MNK)
|
||||
auto problem_shape_MNKL = append<4>(problem_shape, 1);
|
||||
auto [M,N,K,L] = problem_shape_MNKL;
|
||||
auto ptr_A = recast_ptr<ElementAMma>(args.ptr_A);
|
||||
auto ptr_B = recast_ptr<ElementBMma>(args.ptr_B);
|
||||
|
||||
return {
|
||||
args.ptr_A,
|
||||
args.dA,
|
||||
args.ptr_B,
|
||||
args.dB,
|
||||
args.runtime_data_type_a,
|
||||
args.runtime_data_type_b
|
||||
};
|
||||
}
|
||||
|
||||
template <class ProblemShape>
|
||||
static bool
|
||||
can_implement(
|
||||
ProblemShape const& problem_shape,
|
||||
[[maybe_unused]] Arguments const& args) {
|
||||
auto problem_shape_MNKL = append<4>(problem_shape, 1);
|
||||
auto [M,N,K,L] = problem_shape_MNKL;
|
||||
bool implementable = true;
|
||||
implementable = implementable && cutlass::detail::check_alignment<GmemTiledCopyA::NumValSrc>(cute::make_shape(M,K,L), StrideA{});
|
||||
implementable = implementable && cutlass::detail::check_alignment<GmemTiledCopyB::NumValSrc>(cute::make_shape(N,K,L), StrideB{});
|
||||
|
||||
if (!implementable) {
|
||||
CUTLASS_TRACE_HOST(" CAN IMPLEMENT: Problem Size doesn't meet the minimum alignment requirements for CpAsync.\n");
|
||||
}
|
||||
return implementable;
|
||||
}
|
||||
|
||||
/// Construct A Single Stage's Accumulator Shape
|
||||
CUTLASS_DEVICE auto
|
||||
partition_accumulator_shape() {
|
||||
auto acc_shape = partition_shape_C(TiledMma{}, take<0,2>(TileShape{})); // ((MMA_TILE_M,MMA_TILE_N),MMA_M,MMA_N)
|
||||
|
||||
return acc_shape;
|
||||
}
|
||||
|
||||
/// Set up the data needed by this collective for load.
|
||||
/// Return tuple element contain
|
||||
/// gA_mkl - The tiled tensor for input A
|
||||
/// gB_nkl - The tiled tensor for input B
|
||||
/// tAsA - partitioned smem tensor for A
|
||||
/// tBsB - partitioned smem tensor for B
|
||||
template <class ProblemShape_MNKL>
|
||||
CUTLASS_DEVICE auto
|
||||
load_init(
|
||||
ProblemShape_MNKL const& problem_shape_MNKL,
|
||||
Params const& params,
|
||||
TensorStorage& shared_tensors) const {
|
||||
using X = Underscore;
|
||||
// Separate out problem shape for convenience
|
||||
auto [M,N,K,L] = problem_shape_MNKL;
|
||||
|
||||
// Represent the full tensors
|
||||
Tensor mA_mkl = make_tensor(make_gmem_ptr(params.ptr_A), make_shape(M,K,L), params.dA); //(m,k,l)
|
||||
Tensor mB_nkl = make_tensor(make_gmem_ptr(params.ptr_B), make_shape(N,K,L), params.dB); //(n,k,l)
|
||||
// Partition for cpasync
|
||||
Tensor gA_mkl = local_tile(mA_mkl, TileShape{}, make_coord(_,_,_), Step<_1, X,_1>{}); // (BLK_M,BLK_K,m,k,l)
|
||||
Tensor gB_nkl = local_tile(mB_nkl, TileShape{}, make_coord(_,_,_), Step< X,_1,_1>{}); // (BLK_N,BLK_K,n,k,l)
|
||||
|
||||
// Build the coordinate tensors with the same shape as input matrices
|
||||
Tensor cA_mk = make_identity_tensor(make_shape(M,K));
|
||||
Tensor cB_nk = make_identity_tensor(make_shape(N,K));
|
||||
|
||||
// Slice the coordinate tensors in the same way as A/B tensor partitioning
|
||||
Tensor cgA_mk = local_tile(cA_mk, TileShape{}, make_coord(_,_,_), Step<_1, X,_1>{}); // (BLK_M,BLK_K,m,k)
|
||||
Tensor cgB_nk = local_tile(cB_nk, TileShape{}, make_coord(_,_,_), Step< X,_1,_1>{}); // (BLK_N,BLK_K,n,k)
|
||||
|
||||
Tensor sA = make_tensor(make_smem_ptr(shared_tensors.smem_A.data()), LoadSmemLayoutA{});
|
||||
Tensor sB = make_tensor(make_smem_ptr(shared_tensors.smem_B.data()), LoadSmemLayoutB{});
|
||||
|
||||
GmemTiledCopyA gmem_to_smem_a_tiled_copy;
|
||||
GmemTiledCopyB gmem_to_smem_b_tiled_copy;
|
||||
|
||||
int thread_idx = threadIdx.x % NumLoadThreads;
|
||||
auto thr_copy_a = gmem_to_smem_a_tiled_copy.get_slice(thread_idx);
|
||||
auto thr_copy_b = gmem_to_smem_b_tiled_copy.get_slice(thread_idx);
|
||||
|
||||
return cute::make_tuple(
|
||||
gA_mkl, gB_nkl, // gmem
|
||||
cgA_mk, cgB_nk, // crd
|
||||
sA, sB, // smem
|
||||
problem_shape_MNKL,
|
||||
gmem_to_smem_a_tiled_copy, gmem_to_smem_b_tiled_copy,
|
||||
thr_copy_a, thr_copy_b);
|
||||
}
|
||||
|
||||
/// Set up the data needed by this collective for mma compute.
|
||||
template <class FrgEngine, class FrgLayout>
|
||||
CUTLASS_DEVICE auto
|
||||
mma_init(
|
||||
Params const& params,
|
||||
[[maybe_unused]] cute::tuple<cute::Tensor<FrgEngine, FrgLayout>, cute::Tensor<FrgEngine, FrgLayout>> const& accumulators_pair,
|
||||
TensorStorage& shared_tensors) const {
|
||||
Tensor sA = make_tensor(make_smem_ptr(shared_tensors.smem_A.data()), MmaSmemLayoutA{}); // (BLK_M,BLK_K,PIPE)
|
||||
Tensor sB = make_tensor(make_smem_ptr(shared_tensors.smem_B.data()), MmaSmemLayoutB{}); // (BLK_N,BLK_K,PIPE)
|
||||
|
||||
// Allocate "fragments/descriptors" for A and B matrices
|
||||
Tensor tCrA = TiledMma::make_fragment_A(sA); // (MMA,MMA_M,MMA_K,PIPE)
|
||||
Tensor tCrB = TiledMma::make_fragment_B(sB); // (MMA,MMA_N,MMA_K,PIPE)
|
||||
|
||||
CUTE_STATIC_ASSERT_V(Int<DispatchPolicy::Stages>{} == size<3>(sA)); // PIPE
|
||||
CUTE_STATIC_ASSERT_V(Int<DispatchPolicy::Stages>{} == size<3>(sB));
|
||||
|
||||
TiledMma tiled_mma;
|
||||
|
||||
if constexpr (IsRuntimeDataType) {
|
||||
// Update instruction descriptor according to runtime argument.
|
||||
// Applying bitmask (0b111) to help compiler deduce that the conversion and assignment are safe.
|
||||
tiled_mma.idesc_.a_format_ = uint8_t(params.runtime_data_type_a) & 0b111;
|
||||
tiled_mma.idesc_.b_format_ = uint8_t(params.runtime_data_type_b) & 0b111;
|
||||
}
|
||||
|
||||
return cute::make_tuple(tiled_mma, tCrA, tCrB);
|
||||
}
|
||||
|
||||
/// Perform a collective-scoped matrix multiply-accumulate
|
||||
/// Producer Perspective
|
||||
template <
|
||||
class GTensorA, class GTensorB,
|
||||
class CTensorA, class CTensorB,
|
||||
class STensorA, class STensorB,
|
||||
class ProblemShape_MNKL,
|
||||
class TiledCopyA, class TiledCopyB,
|
||||
class ThreadCopyA, class ThreadCopyB,
|
||||
class TileCoordMNKL,
|
||||
class KTileIterator
|
||||
>
|
||||
CUTLASS_DEVICE auto
|
||||
load(
|
||||
Params const& params,
|
||||
MainloopPipeline mainloop_pipeline,
|
||||
MainloopPipelineState mainloop_pipe_producer_state,
|
||||
cute::tuple<GTensorA, GTensorB,
|
||||
CTensorA, CTensorB,
|
||||
STensorA, STensorB,
|
||||
ProblemShape_MNKL,
|
||||
TiledCopyA, TiledCopyB,
|
||||
ThreadCopyA, ThreadCopyB> const& load_inputs,
|
||||
TileCoordMNKL const& cta_coord_mnkl,
|
||||
KTileIterator k_tile_iter, int k_tile_count) {
|
||||
// Unpack from load_inputs
|
||||
GTensorA tAgA_mkl = get<0>(load_inputs);
|
||||
GTensorB tBgB_nkl = get<1>(load_inputs);
|
||||
CTensorA cgA_mk = get<2>(load_inputs);
|
||||
CTensorB cgB_nk = get<3>(load_inputs);
|
||||
STensorA sA = get<4>(load_inputs);
|
||||
STensorB sB = get<5>(load_inputs);
|
||||
ProblemShape_MNKL problem_shape_MNKL = get<6>(load_inputs);
|
||||
TiledCopyA gmem_to_smem_a_tiled_copy = get<7>(load_inputs);
|
||||
TiledCopyB gmem_to_smem_b_tiled_copy = get<8>(load_inputs);
|
||||
ThreadCopyA thr_copy_a = get<9>(load_inputs);
|
||||
ThreadCopyB thr_copy_b = get<10>(load_inputs);
|
||||
auto [M,N,K,L] = problem_shape_MNKL;
|
||||
|
||||
// Slice out the work coord from partitioned tensors
|
||||
Tensor gA_in = tAgA_mkl(_, _, get<0>(cta_coord_mnkl), _, get<3>(cta_coord_mnkl));
|
||||
Tensor gB_in = tBgB_nkl(_, _, get<1>(cta_coord_mnkl), _, get<3>(cta_coord_mnkl));
|
||||
|
||||
// Repeat slicing out coordinate tensor exactly the same as input tensor does
|
||||
Tensor cgA_mk_in = cgA_mk(_, _, get<0>(cta_coord_mnkl), _);
|
||||
Tensor cgB_nk_in = cgB_nk(_, _, get<1>(cta_coord_mnkl), _);
|
||||
|
||||
auto k_residue = K - size<1>(gB_in) * size<2>(gA_in);
|
||||
|
||||
// Shift tensor so residue_k is at origin (Can't read any k_coord < residue_k)
|
||||
// This aligns the tensor with BLK_K for all but the 0th k_tile
|
||||
Tensor gA = domain_offset(make_coord(0, k_residue, 0), gA_in);
|
||||
Tensor gB = domain_offset(make_coord(0, k_residue, 0), gB_in);
|
||||
|
||||
Tensor cA = domain_offset(make_coord(0, k_residue, 0), cgA_mk_in);
|
||||
Tensor cB = domain_offset(make_coord(0, k_residue, 0), cgB_nk_in);
|
||||
|
||||
auto tAgA = thr_copy_a.partition_S(gA);
|
||||
auto tAsA = thr_copy_a.partition_D(sA);
|
||||
|
||||
auto tBgB = thr_copy_b.partition_S(gB);
|
||||
auto tBsB = thr_copy_b.partition_D(sB);
|
||||
|
||||
// Allocate predicate tensors for m and n
|
||||
Tensor tApA = make_tensor<bool>(make_shape(size<1>(tAsA), size<2>(tAsA)), Stride<_1,_0>{});
|
||||
Tensor tBpB = make_tensor<bool>(make_shape(size<1>(tBsB), size<2>(tBsB)), Stride<_1,_0>{});
|
||||
|
||||
Tensor tAcA = thr_copy_a.partition_S(cA);
|
||||
Tensor tBcB = thr_copy_b.partition_S(cB);
|
||||
|
||||
// Copy gmem to smem for *k_tile_iter, predicating for k residue
|
||||
Tensor tAgAk = tAgA(_,_,_,*k_tile_iter);
|
||||
Tensor tBgBk = tBgB(_,_,_,*k_tile_iter);
|
||||
|
||||
// Repeating on predicators with the same operations on tAgA and tBgB
|
||||
Tensor tAcAk = tAcA(_,_,_,*k_tile_iter);
|
||||
Tensor tBcBk = tBcB(_,_,_,*k_tile_iter);
|
||||
|
||||
// Set predicates for m bounds
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int m = 0; m < size<0>(tApA); ++m) {
|
||||
tApA(m,0) = elem_less(get<0>(tAcAk(0,m,0)), M); // blk_m coord < M
|
||||
}
|
||||
// Set predicates for n bounds
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int n = 0; n < size<0>(tBpB); ++n) {
|
||||
tBpB(n,0) = elem_less(get<0>(tBcBk(0,n,0)), N); // blk_n coord < N
|
||||
}
|
||||
|
||||
// 0-th stage with predication on k to account for residue
|
||||
// For performance consideration,
|
||||
// this predicated block for K-tail is only activated when there is k-residue
|
||||
if (k_residue != 0 && k_tile_count > 0) {
|
||||
// LOCK mainloop_pipe_producer_state for _writing_
|
||||
mainloop_pipeline.producer_acquire(mainloop_pipe_producer_state);
|
||||
int write_stage = mainloop_pipe_producer_state.index();
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int k = 0; k < size<2>(tAsA); ++k) {
|
||||
if ( int(get<1>(tAcAk(0,0,k))) >= 0) { // blk_k coord < K
|
||||
copy_if(gmem_to_smem_a_tiled_copy, tApA(_,k), tAgAk(_,_,k), tAsA(_,_,k,write_stage));
|
||||
}
|
||||
else {
|
||||
clear(tAsA(_,_,k,write_stage));
|
||||
}
|
||||
}
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int k = 0; k < size<2>(tBsB); ++k) {
|
||||
if (int(get<1>(tBcBk(0,0,k))) >= 0) { // blk_k coord < K
|
||||
copy_if(gmem_to_smem_b_tiled_copy, tBpB(_,k), tBgBk(_,_,k), tBsB(_,_,k,write_stage));
|
||||
}
|
||||
else {
|
||||
clear(tBsB(_,_,k,write_stage));
|
||||
}
|
||||
}
|
||||
++k_tile_iter;
|
||||
--k_tile_count;
|
||||
|
||||
// UNLOCK mainloop_pipe_producer_state
|
||||
mainloop_pipeline.producer_commit(mainloop_pipe_producer_state, cutlass::arch::cpasync_barrier_arrive);
|
||||
|
||||
// Advance mainloop_pipe_producer_state
|
||||
++mainloop_pipe_producer_state;
|
||||
}
|
||||
|
||||
auto barrier_token = mainloop_pipeline.producer_try_acquire(mainloop_pipe_producer_state);
|
||||
|
||||
// Issue the Mainloop loads
|
||||
CUTLASS_PRAGMA_NO_UNROLL
|
||||
while (k_tile_count > 0) {
|
||||
auto mainloop_pipe_producer_state_curr = mainloop_pipe_producer_state;
|
||||
++mainloop_pipe_producer_state;
|
||||
mainloop_pipeline.producer_acquire(mainloop_pipe_producer_state_curr, barrier_token);
|
||||
barrier_token = mainloop_pipeline.producer_try_acquire(mainloop_pipe_producer_state);
|
||||
int write_stage = mainloop_pipe_producer_state_curr.index();
|
||||
|
||||
copy_if(gmem_to_smem_a_tiled_copy, tApA, tAgA(_,_,_,*k_tile_iter), tAsA(_,_,_,write_stage));
|
||||
copy_if(gmem_to_smem_b_tiled_copy, tBpB, tBgB(_,_,_,*k_tile_iter), tBsB(_,_,_,write_stage));
|
||||
|
||||
mainloop_pipeline.producer_commit(mainloop_pipe_producer_state_curr, cutlass::arch::cpasync_barrier_arrive);
|
||||
|
||||
--k_tile_count;
|
||||
++k_tile_iter;
|
||||
}
|
||||
|
||||
return cute::make_tuple(mainloop_pipe_producer_state, k_tile_iter);
|
||||
|
||||
}
|
||||
|
||||
/// Perform a Producer Epilogue to prevent early exit of ctas in a Cluster
|
||||
CUTLASS_DEVICE void
|
||||
load_tail(MainloopPipeline mainloop_pipeline, MainloopPipelineState mainloop_pipe_producer_state) {
|
||||
// Issue the epilogue waits
|
||||
// This helps avoid early exit of ctas in Cluster
|
||||
// Waits for all stages to either be released (all
|
||||
// Consumer UNLOCKs), or if the stage was never used
|
||||
// then would just be acquired since the phase was
|
||||
// still inverted from make_producer_start_state
|
||||
mainloop_pipeline.producer_tail(mainloop_pipe_producer_state);
|
||||
}
|
||||
|
||||
/// Perform a collective-scoped matrix multiply-accumulate
|
||||
/// Consumer Perspective
|
||||
template <
|
||||
class FrgEngine, class FrgLayout,
|
||||
class FragmentA, class FragmentB
|
||||
>
|
||||
CUTLASS_DEVICE auto
|
||||
mma(MainloopPipeline mainloop_pipeline,
|
||||
MainloopPipelineState mainloop_pipe_consumer_state,
|
||||
cute::tuple<cute::Tensor<FrgEngine, FrgLayout>, cute::Tensor<FrgEngine, FrgLayout>> const& accumulators_pair,
|
||||
cute::tuple<TiledMma, FragmentA, FragmentB> const& mma_inputs,
|
||||
int k_tile_count
|
||||
) {
|
||||
static_assert(is_tmem<FrgEngine>::value, "Accumulator must be tmem resident.");
|
||||
static_assert(rank(FrgLayout{}) == 3, "Accumulator must be MMA-partitioned: (MMA, MMA_M, MMA_N)");
|
||||
auto accumulators = get<0>(accumulators_pair);
|
||||
auto [tiled_mma, tCrA, tCrB] = mma_inputs;
|
||||
|
||||
//
|
||||
// PIPELINED MAIN LOOP
|
||||
//
|
||||
tiled_mma.accumulate_ = UMMA::ScaleOut::Zero;
|
||||
|
||||
CUTLASS_PRAGMA_NO_UNROLL
|
||||
while (k_tile_count > 0) {
|
||||
mainloop_pipeline.consumer_wait(mainloop_pipe_consumer_state);
|
||||
|
||||
int read_stage = mainloop_pipe_consumer_state.index();
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int k_block = 0; k_block < size<2>(tCrA); ++k_block) {
|
||||
// (V,M) x (V,N) => (V,M,N)
|
||||
cute::gemm(tiled_mma, tCrA(_,_,k_block,read_stage), tCrB(_,_,k_block,read_stage), accumulators);
|
||||
tiled_mma.accumulate_ = UMMA::ScaleOut::One;
|
||||
}
|
||||
|
||||
mainloop_pipeline.consumer_release(mainloop_pipe_consumer_state);
|
||||
--k_tile_count;
|
||||
++mainloop_pipe_consumer_state;
|
||||
}
|
||||
|
||||
return mainloop_pipe_consumer_state;
|
||||
}
|
||||
|
||||
};
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace cutlass::gemm::collective
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
@@ -347,20 +347,20 @@ struct CollectiveMma<
|
||||
|
||||
template<
|
||||
class KTileCount,
|
||||
class GTensorPartitionedScaleA, class GTensorPartitionedScaleB,
|
||||
class IdentTensorPartitionedScaleA, class IdentTensorPartitionedScaleB,
|
||||
class GTensorScaleA, class GTensorScaleB,
|
||||
class IdentTensorScaleA, class IdentTensorScaleB,
|
||||
class STensorScaleA, class STensorScaleB
|
||||
>
|
||||
struct LoadSFParams {
|
||||
// for scheduler
|
||||
KTileCount k_tiles;
|
||||
|
||||
GTensorPartitionedScaleA tSFAgSFA_mkl;
|
||||
GTensorPartitionedScaleB tSFBgSFB_nkl;
|
||||
IdentTensorPartitionedScaleA tSFAIdentSFA_mkl;
|
||||
IdentTensorPartitionedScaleB tSFBIdentSFB_nkl;
|
||||
STensorScaleA tSFAsSFA;
|
||||
STensorScaleB tSFBsSFB;
|
||||
GTensorScaleA gSFA_mkl;
|
||||
GTensorScaleB gSFB_nkl;
|
||||
IdentTensorScaleA identSFA_mkl;
|
||||
IdentTensorScaleB identSFB_nkl;
|
||||
STensorScaleA sSFA;
|
||||
STensorScaleB sSFB;
|
||||
|
||||
LayoutSFA layout_SFA;
|
||||
LayoutSFB layout_SFB;
|
||||
@@ -368,14 +368,14 @@ struct CollectiveMma<
|
||||
CUTLASS_DEVICE
|
||||
LoadSFParams (
|
||||
KTileCount k_tiles_,
|
||||
GTensorPartitionedScaleA tSFAgSFA_mkl_, GTensorPartitionedScaleB tSFBgSFB_nkl_,
|
||||
IdentTensorPartitionedScaleA tSFAIdentSFA_mkl_, IdentTensorPartitionedScaleB tSFBIdentSFB_nkl_,
|
||||
STensorScaleA tSFAsSFA_, STensorScaleB tSFBsSFB_,
|
||||
GTensorScaleA gSFA_mkl_, GTensorScaleB gSFB_nkl_,
|
||||
IdentTensorScaleA identSFA_mkl_, IdentTensorScaleB identSFB_nkl_,
|
||||
STensorScaleA sSFA_, STensorScaleB sSFB_,
|
||||
LayoutSFA layout_SFA_, LayoutSFB layout_SFB_)
|
||||
: k_tiles(k_tiles_)
|
||||
, tSFAgSFA_mkl(tSFAgSFA_mkl_), tSFBgSFB_nkl(tSFBgSFB_nkl_)
|
||||
, tSFAIdentSFA_mkl(tSFAIdentSFA_mkl_), tSFBIdentSFB_nkl(tSFBIdentSFB_nkl_)
|
||||
, tSFAsSFA(tSFAsSFA_), tSFBsSFB(tSFBsSFB_)
|
||||
, gSFA_mkl(gSFA_mkl_), gSFB_nkl(gSFB_nkl_)
|
||||
, identSFA_mkl(identSFA_mkl_), identSFB_nkl(identSFB_nkl_)
|
||||
, sSFA(sSFA_), sSFB(sSFB_)
|
||||
, layout_SFA(layout_SFA_), layout_SFB(layout_SFB_) {}
|
||||
};
|
||||
|
||||
@@ -732,35 +732,16 @@ struct CollectiveMma<
|
||||
static_assert(rank(decltype(gSFA_mkl){}) == 5);
|
||||
static_assert(rank(decltype(gSFB_nkl){}) == 5);
|
||||
|
||||
// 1 thread copies entire set of scalar
|
||||
GmemTiledCopySFA scale_copy_a{};
|
||||
GmemTiledCopySFB scale_copy_b{};
|
||||
|
||||
ThrCopy thr_scale_copy_a = scale_copy_a.get_slice(threadIdx.x % size(scale_copy_a));
|
||||
ThrCopy thr_scale_copy_b = scale_copy_b.get_slice(threadIdx.x % size(scale_copy_b));
|
||||
|
||||
Tensor sSFA = make_tensor(make_smem_ptr(shared_tensors.smem_SFA.begin()),
|
||||
SmemLayoutScaleA{}); // (CTA_M,CTA_K,P)
|
||||
Tensor sSFB = make_tensor(make_smem_ptr(shared_tensors.smem_SFB.begin()),
|
||||
SmemLayoutScaleB{}); // (CTA_M,CTA_K,P)
|
||||
|
||||
Tensor tSFAgSFA_mkl = thr_scale_copy_a.partition_S(gSFA_mkl); // (CPY, BLK_M, BLK_K, m, k, l)
|
||||
Tensor tSFAIdentSFA_mkl = thr_scale_copy_a.partition_S(identSFA_mkl); // (CPY, BLK_M, BLK_K, m, k, l)
|
||||
|
||||
Tensor tSFAsSFA = thr_scale_copy_a.partition_D(sSFA);
|
||||
|
||||
Tensor tSFBgSFB_nkl = thr_scale_copy_b.partition_S(gSFB_nkl); // (CPY, BLK_N, BLK_K, m, k, l)
|
||||
Tensor tSFBIdentSFB_nkl = thr_scale_copy_b.partition_S(identSFB_nkl); // (CPY, BLK_N, BLK_K, m, k, l)
|
||||
Tensor tSFBsSFB = thr_scale_copy_b.partition_D(sSFB);
|
||||
|
||||
static_assert(rank(decltype(tSFAgSFA_mkl){}) == 6);
|
||||
static_assert(rank(decltype(tSFBgSFB_nkl){}) == 6);
|
||||
|
||||
LoadSFParams load_params {
|
||||
size<3>(gSFA_mkl),
|
||||
tSFAgSFA_mkl, tSFBgSFB_nkl, // for input scale tensor values
|
||||
tSFAIdentSFA_mkl, tSFBIdentSFB_nkl, // for predicating scale tensor copies
|
||||
tSFAsSFA, tSFBsSFB, // for scale tensor values
|
||||
gSFA_mkl, gSFB_nkl, // for input scale tensor values
|
||||
identSFA_mkl, identSFB_nkl, // for predicating scale tensor copies
|
||||
sSFA, sSFB, // for scale tensor values
|
||||
mainloop_params.layout_SFA, // for predicating scale tensor copies
|
||||
mainloop_params.layout_SFB // for predicating scale tensor copies
|
||||
};
|
||||
@@ -922,24 +903,44 @@ struct CollectiveMma<
|
||||
KTileIterator k_tile_iter, int k_tile_count) {
|
||||
|
||||
auto [unused_k_tiles,
|
||||
tSFAgSFA_mkl, tSFBgSFB_nkl,
|
||||
tSFAIdentSFA_mkl, tSFBIdentSFB_nkl,
|
||||
tSFAsSFA, tSFBsSFB,
|
||||
gSFA_mkl, gSFB_nkl,
|
||||
identSFA_mkl, identSFB_nkl,
|
||||
sSFA, sSFB,
|
||||
layout_SFA, layout_SFB] = load_inputs;
|
||||
|
||||
// slice out the work coord from partitioned tensors
|
||||
GmemTiledCopySFA scale_copy_a{};
|
||||
GmemTiledCopySFB scale_copy_b{};
|
||||
|
||||
Tensor tSFAgSFA = tSFAgSFA_mkl(_, _, _, get<0>(cta_coord_mnkl), _, get<3>(cta_coord_mnkl));
|
||||
Tensor gSFA_k_compact = filter_zeros(
|
||||
gSFA_mkl(_, _, get<0>(cta_coord_mnkl), _, get<3>(cta_coord_mnkl))); // (BLK_M_CPT, BLK_K_CPT, k_cpt)
|
||||
Tensor gSFB_k_compact = filter_zeros(
|
||||
gSFB_nkl(_, _, get<1>(cta_coord_mnkl), _, get<3>(cta_coord_mnkl))); // (BLK_N_CPT, BLK_K_CPT, k_cpt)
|
||||
|
||||
Tensor tSFBgSFB = tSFBgSFB_nkl(_, _, _, get<1>(cta_coord_mnkl), _, get<3>(cta_coord_mnkl));
|
||||
Tensor identSFA_k_compact = filter_zeros(
|
||||
identSFA_mkl(_, _, get<0>(cta_coord_mnkl), _, get<3>(cta_coord_mnkl)),
|
||||
gSFA_k_compact.stride()); // (BLK_M_CPT, BLK_K_CPT, k_cpt)
|
||||
Tensor identSFB_k_compact = filter_zeros(
|
||||
identSFB_nkl(_, _, get<1>(cta_coord_mnkl), _, get<3>(cta_coord_mnkl)),
|
||||
gSFB_k_compact.stride()); // (BLK_N_CPT, BLK_K_CPT, k_cpt)
|
||||
|
||||
Tensor thr_tile_SFA_k = tSFAIdentSFA_mkl(_0{}, _, _, get<0>(cta_coord_mnkl), _, get<3>(cta_coord_mnkl));
|
||||
Tensor thr_tile_pSFA = make_tensor<bool>(shape(filter_zeros(thr_tile_SFA_k(_,_,_0{}), tSFAgSFA(_0{},_,_,_0{}).stride())));
|
||||
Tensor thr_tile_SFB_k = tSFBIdentSFB_nkl(_0{}, _, _, get<1>(cta_coord_mnkl), _, get<3>(cta_coord_mnkl));
|
||||
Tensor sSFA_compact = filter_zeros(sSFA); // (BLK_M_CPT, BLK_K_CPT, P)
|
||||
Tensor sSFB_compact = filter_zeros(sSFB); // (BLK_N_CPT, BLK_K_CPT, P)
|
||||
|
||||
Tensor thr_tile_pSFB = make_tensor<bool>(shape(filter_zeros(thr_tile_SFB_k(_,_,_0{}), tSFBgSFB(_0{},_,_,_0{}).stride())));
|
||||
ThrCopy thr_scale_copy_a = scale_copy_a.get_slice(threadIdx.x % size(scale_copy_a));
|
||||
ThrCopy thr_scale_copy_b = scale_copy_b.get_slice(threadIdx.x % size(scale_copy_b));
|
||||
|
||||
Tensor tSFAgSFA_k_compact = thr_scale_copy_a.partition_S(gSFA_k_compact); // (CPY, BLK_M, BLK_K, k)
|
||||
Tensor tSFAIdentSFA_k_compact = thr_scale_copy_a.partition_S(identSFA_k_compact); // (CPY, BLK_M, BLK_K, k)
|
||||
|
||||
Tensor tSFAsSFA_compact = thr_scale_copy_a.partition_D(sSFA_compact);
|
||||
|
||||
Tensor tSFBgSFB_k_compact = thr_scale_copy_b.partition_S(gSFB_k_compact); // (CPY, BLK_N, BLK_K, k)
|
||||
Tensor tSFBIdentSFB_k_compact = thr_scale_copy_b.partition_S(identSFB_k_compact); // (CPY, BLK_N, BLK_K, k)
|
||||
Tensor tSFBsSFB_compact = thr_scale_copy_b.partition_D(sSFB_compact);
|
||||
|
||||
Tensor thr_tile_pSFA = make_fragment_like<bool>(tSFAgSFA_k_compact(_0{},_,_,_0{}));
|
||||
Tensor thr_tile_pSFB = make_fragment_like<bool>(tSFBgSFB_k_compact(_0{},_,_,_0{}));
|
||||
|
||||
// Issue the loads
|
||||
CUTLASS_PRAGMA_NO_UNROLL
|
||||
@@ -949,18 +950,22 @@ struct CollectiveMma<
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int i = 0; i < size(thr_tile_pSFA); ++i) {
|
||||
Tensor thr_tile_SFA = filter_zeros(thr_tile_SFA_k(_,_,*k_tile_iter), tSFAgSFA(_0{},_,_,_0{}).stride());
|
||||
thr_tile_pSFA(i) = elem_less(thr_tile_SFA(i), shape(filter_zeros(layout_SFA))) && threadIdx.x % 32 < size(scale_copy_a);
|
||||
Tensor tSFAIdentSFA_compact = tSFAIdentSFA_k_compact(_0{},_,_,*k_tile_iter);
|
||||
thr_tile_pSFA(i) = elem_less(tSFAIdentSFA_compact(i),
|
||||
shape(filter_zeros(layout_SFA))) && threadIdx.x % 32 < size(scale_copy_a);
|
||||
}
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int i = 0; i < size(thr_tile_pSFB); ++i) {
|
||||
Tensor thr_tile_SFB = filter_zeros(thr_tile_SFB_k(_,_,*k_tile_iter), tSFBgSFB(_0{},_,_,_0{}).stride());
|
||||
thr_tile_pSFB(i) = elem_less(thr_tile_SFB(i), shape(filter_zeros(layout_SFB))) && threadIdx.x % 32 < size(scale_copy_b);
|
||||
Tensor tSFBIdentSFB_compact = tSFBIdentSFB_k_compact(_0{},_,_,*k_tile_iter);
|
||||
thr_tile_pSFB(i) = elem_less(tSFBIdentSFB_compact(i),
|
||||
shape(filter_zeros(layout_SFB))) && threadIdx.x % 32 < size(scale_copy_b);
|
||||
}
|
||||
|
||||
copy_if(scale_copy_a, thr_tile_pSFA, filter_zeros(tSFAgSFA(_,_,_,*k_tile_iter)), filter_zeros(tSFAsSFA(_,_,_,mainloop_sf_pipe_producer_state.index())));
|
||||
copy_if(scale_copy_b, thr_tile_pSFB, filter_zeros(tSFBgSFB(_,_,_,*k_tile_iter)), filter_zeros(tSFBsSFB(_,_,_,mainloop_sf_pipe_producer_state.index())));
|
||||
copy_if(scale_copy_a, thr_tile_pSFA, tSFAgSFA_k_compact(_,_,_,*k_tile_iter),
|
||||
tSFAsSFA_compact(_,_,_,mainloop_sf_pipe_producer_state.index()));
|
||||
copy_if(scale_copy_b, thr_tile_pSFB, tSFBgSFB_k_compact(_,_,_,*k_tile_iter),
|
||||
tSFBsSFB_compact(_,_,_,mainloop_sf_pipe_producer_state.index()));
|
||||
mainloop_sf_pipeline.producer_commit(mainloop_sf_pipe_producer_state, cutlass::arch::cpasync_barrier_arrive_noinc);
|
||||
|
||||
__syncwarp();
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
+9
-1
@@ -153,7 +153,15 @@ struct CollectiveMma<
|
||||
static_assert((size<0>(TileShape{}) % ScaleGranularityM) == 0, "FP8 scaling granularity must evenly divide tile shape along M.");
|
||||
static_assert((size<1>(TileShape{}) % ScaleGranularityN) == 0, "FP8 scaling granularity must evenly divide tile shape along N.");
|
||||
|
||||
using ScaleConfig = ::cutlass::detail::Sm90BlockwiseScaleConfig<ScaleGranularityM, ScaleGranularityN, ScaleGranularityK>;
|
||||
static constexpr bool MMajorSFA = size<0,1>(InternalLayoutSFA{}.stride()) == 1;
|
||||
static constexpr bool NMajorSFB = size<0,1>(InternalLayoutSFB{}.stride()) == 1;
|
||||
|
||||
using ScaleConfig = ::cutlass::detail::Sm90BlockwiseScaleConfig<
|
||||
ScaleGranularityM,
|
||||
ScaleGranularityN,
|
||||
ScaleGranularityK,
|
||||
MMajorSFA ? cute::GMMA::Major::MN : cute::GMMA::Major::K,
|
||||
NMajorSFB ? cute::GMMA::Major::MN : cute::GMMA::Major::K>;
|
||||
using SmemLayoutAtomSFA = decltype(ScaleConfig::smem_atom_layoutSFA(TileShape{}));
|
||||
using SmemLayoutAtomSFB = decltype(ScaleConfig::smem_atom_layoutSFB(TileShape{}));
|
||||
|
||||
|
||||
+15
-7
@@ -134,9 +134,12 @@ struct CollectiveMma<
|
||||
static constexpr int ScaleMsPerTile = size<0>(TileShape{}) / ScaleGranularityM;
|
||||
static constexpr int ScaleNsPerTile = size<1>(TileShape{}) / ScaleGranularityN;
|
||||
|
||||
static constexpr bool MMajorSFA = size<0,1>(LayoutSFA{}.stride()) == 1;
|
||||
static constexpr bool NMajorSFB = size<0,1>(LayoutSFB{}.stride()) == 1;
|
||||
|
||||
static constexpr int ScaleTmaThreshold = 32;
|
||||
static constexpr bool IsTmaLoadSFA = ScaleMsPerTile >= ScaleTmaThreshold && ScaleNsPerTile < ScaleTmaThreshold;
|
||||
static constexpr bool IsTmaLoadSFB = ScaleNsPerTile >= ScaleTmaThreshold && ScaleMsPerTile < ScaleTmaThreshold;
|
||||
static constexpr bool IsTmaLoadSFA = ScaleMsPerTile >= ScaleTmaThreshold && ScaleNsPerTile < ScaleTmaThreshold && MMajorSFA;
|
||||
static constexpr bool IsTmaLoadSFB = ScaleNsPerTile >= ScaleTmaThreshold && ScaleMsPerTile < ScaleTmaThreshold && NMajorSFB;
|
||||
// Two threads per CTA are producers (1 for operand tile `tma`, and 32 for scales `cp.async`)
|
||||
static constexpr int NumProducerThreadEvents = ((IsTmaLoadSFA && IsTmaLoadSFB)? 1 : 33);
|
||||
|
||||
@@ -151,7 +154,12 @@ struct CollectiveMma<
|
||||
static_assert((size<0>(TileShape{}) % ScaleGranularityM) == 0, "FP8 scaling granularity must evenly divide tile shape along M.");
|
||||
static_assert((size<1>(TileShape{}) % ScaleGranularityN) == 0, "FP8 scaling granularity must evenly divide tile shape along N.");
|
||||
|
||||
using ScaleConfig = ::cutlass::detail::Sm90BlockwiseScaleConfig<ScaleGranularityM, ScaleGranularityN, ScaleGranularityK>;
|
||||
using ScaleConfig = ::cutlass::detail::Sm90BlockwiseScaleConfig<
|
||||
ScaleGranularityM,
|
||||
ScaleGranularityN,
|
||||
ScaleGranularityK,
|
||||
MMajorSFA ? cute::GMMA::Major::MN : cute::GMMA::Major::K,
|
||||
NMajorSFB ? cute::GMMA::Major::MN : cute::GMMA::Major::K>;
|
||||
using SmemLayoutAtomSFA = decltype(ScaleConfig::smem_atom_layoutSFA(TileShape{}));
|
||||
using SmemLayoutAtomSFB = decltype(ScaleConfig::smem_atom_layoutSFB(TileShape{}));
|
||||
|
||||
@@ -170,8 +178,8 @@ struct CollectiveMma<
|
||||
using CopyAtomSFA = Copy_Atom<SM80_CP_ASYNC_CACHEALWAYS<ElementBlockScale>, ElementBlockScale>;
|
||||
using CopyAtomSFB = Copy_Atom<SM80_CP_ASYNC_CACHEALWAYS<ElementBlockScale>, ElementBlockScale>;
|
||||
|
||||
static constexpr int AlignmentSFA = 1;
|
||||
static constexpr int AlignmentSFB = 1;
|
||||
static constexpr int AlignmentSFA = IsTmaLoadSFA ? 128 / cutlass::sizeof_bits<ElementBlockScale>::value : 1;
|
||||
static constexpr int AlignmentSFB = IsTmaLoadSFB ? 128 / cutlass::sizeof_bits<ElementBlockScale>::value : 1;
|
||||
|
||||
// Block scaling smem layout
|
||||
using SmemLayoutSFA = decltype(make_layout(
|
||||
@@ -669,7 +677,7 @@ struct CollectiveMma<
|
||||
Tensor tSFAcSFA_compact = filter_zeros(tSFAcSFA);
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int i = 0; i < size(tSFApSFA); ++i) {
|
||||
tSFApSFA(i) = load_sfa && elem_less(get<0>(tSFAcSFA_compact(i)), get<0>(SFA_shape));
|
||||
tSFApSFA(i) = load_sfa && elem_less(tSFAcSFA_compact(i), SFA_shape);
|
||||
}
|
||||
|
||||
bool load_sfb = thread_idx < ScaleNsPerTile;
|
||||
@@ -677,7 +685,7 @@ struct CollectiveMma<
|
||||
Tensor tSFBcSFB_compact = filter_zeros(tSFBcSFB);
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int i = 0; i < size(tSFBpSFB); ++i) {
|
||||
tSFBpSFB(i) = load_sfb && elem_less(get<0>(tSFBcSFB_compact(i)), get<0>(SFB_shape));
|
||||
tSFBpSFB(i) = load_sfb && elem_less(tSFBcSFB_compact(i), SFB_shape);
|
||||
}
|
||||
int write_stage = smem_pipe_write.index();
|
||||
// Copy scale tensors from global memory to shared memory
|
||||
|
||||
@@ -393,6 +393,7 @@ public:
|
||||
[[maybe_unused]] dim3 fallback_cluster = dim3{0,0,0};
|
||||
if constexpr (GemmKernel::ArchTag::kMinComputeCapability == 100
|
||||
|| GemmKernel::ArchTag::kMinComputeCapability == 101
|
||||
|| GemmKernel::ArchTag::kMinComputeCapability == 103
|
||||
) {
|
||||
if constexpr (!cute::is_static_v<typename GemmKernel::DispatchPolicy::ClusterShape>) {
|
||||
fallback_cluster = params.hw_info.cluster_shape_fallback;
|
||||
@@ -473,6 +474,7 @@ public:
|
||||
if constexpr (GemmKernel::ArchTag::kMinComputeCapability == 100
|
||||
|| GemmKernel::ArchTag::kMinComputeCapability == 101
|
||||
|| GemmKernel::ArchTag::kMinComputeCapability == 120
|
||||
|| GemmKernel::ArchTag::kMinComputeCapability == 103
|
||||
) {
|
||||
if constexpr (is_static_1x1x1) {
|
||||
#if (CUTLASS_DEBUG_TRACE_LEVEL > 1)
|
||||
|
||||
@@ -34,14 +34,13 @@
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "cutlass/cutlass.h"
|
||||
#if defined(__CUDACC_RTC__)
|
||||
#include <cuda/std/limits>
|
||||
#include CUDA_STD_HEADER(limits)
|
||||
#else
|
||||
#include <limits>
|
||||
#endif
|
||||
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cutlass/numeric_types.h"
|
||||
#include "cutlass/arch/arch.h"
|
||||
#include "cutlass/device_kernel.h"
|
||||
|
||||
@@ -0,0 +1,183 @@
|
||||
/***************************************************************************************************
|
||||
* Copyright (c) 2024 - 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
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cutlass/numeric_types.h"
|
||||
#include "cutlass/arch/arch.h"
|
||||
#include "cutlass/device_kernel.h"
|
||||
|
||||
#include "cutlass/gemm/gemm.h"
|
||||
#include "cutlass/gemm/threadblock/threadblock_swizzle.h"
|
||||
#include "cutlass/gemm/kernel/gemm_universal.h"
|
||||
|
||||
#include "cutlass/gemm/kernel/default_gemm_universal.h"
|
||||
#include "cutlass/gemm/device/default_gemm_configuration.h"
|
||||
#include "cutlass/gemm/device/gemm_universal_base.h"
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace cutlass {
|
||||
namespace gemm {
|
||||
namespace device {
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <typename GemvKernel_>
|
||||
class GemvBlockScaled {
|
||||
public:
|
||||
|
||||
using GemvKernel = GemvKernel_;
|
||||
|
||||
|
||||
using ElementA = typename GemvKernel::ElementA;
|
||||
using LayoutA = typename GemvKernel::LayoutA;
|
||||
using ElementB = typename GemvKernel::ElementB;
|
||||
using ElementC = typename GemvKernel::ElementC;
|
||||
|
||||
using ElementSFA = typename GemvKernel::ElementSFA;
|
||||
using ElementSFB = typename GemvKernel::ElementSFB;
|
||||
|
||||
using ElementAccumulator = typename GemvKernel::ElementAccumulator;
|
||||
using EpilogueOutputOp = typename GemvKernel::EpilogueOutputOp;
|
||||
|
||||
static ComplexTransform const kTransformA = GemvKernel::kTransformA;
|
||||
static ComplexTransform const kTransformB = GemvKernel::kTransformB;
|
||||
|
||||
static int const kThreadCount = GemvKernel::kThreadCount;
|
||||
static int const kThreadsPerRow = GemvKernel::kThreadsPerRow;
|
||||
|
||||
using Arguments = typename GemvKernel::Arguments;
|
||||
using Params = typename GemvKernel::Params;
|
||||
|
||||
private:
|
||||
|
||||
Params params_;
|
||||
|
||||
public:
|
||||
|
||||
/// Constructs the GemvBlockScaled.
|
||||
GemvBlockScaled() = default;
|
||||
|
||||
/// Determines whether the GemvBlockScaled can execute the given problem.
|
||||
static Status can_implement(Arguments const &args) {
|
||||
|
||||
return GemvKernel::can_implement(args);
|
||||
}
|
||||
|
||||
/// Gets the workspace size
|
||||
static size_t get_workspace_size(Arguments const &args) {
|
||||
|
||||
return 0;
|
||||
}
|
||||
|
||||
/// Computes the grid shape
|
||||
static dim3 get_grid_shape(Arguments const &args, dim3 const &block) {
|
||||
if(platform::is_same<LayoutA, layout::ColumnMajor>::value) {
|
||||
return dim3((args.problem_size.row() + (block.x - 1)) / block.x, 1, args.batch_count % 65536);
|
||||
}
|
||||
else {
|
||||
return dim3((args.problem_size.row() + (block.y - 1)) / block.y, 1, args.batch_count % 65536);
|
||||
}
|
||||
}
|
||||
|
||||
/// Computes the block shape
|
||||
static dim3 get_block_shape() {
|
||||
if(platform::is_same<LayoutA, layout::ColumnMajor>::value) {
|
||||
return dim3(kThreadCount, 1, 1);
|
||||
}
|
||||
else {
|
||||
return dim3(kThreadsPerRow, kThreadCount / kThreadsPerRow, 1);
|
||||
}
|
||||
}
|
||||
|
||||
/// Initializes GemvBlockScaled state from arguments.
|
||||
Status initialize(Arguments const &args, void *workspace = nullptr, cudaStream_t stream = nullptr) {
|
||||
params_ = Params(args);
|
||||
return Status::kSuccess;
|
||||
}
|
||||
|
||||
/// Lightweight update given a subset of arguments
|
||||
Status update(Arguments const &args, void *workspace = nullptr) {
|
||||
return params_.update(args);
|
||||
}
|
||||
|
||||
/// Runs the kernel using initialized state.
|
||||
Status run(cudaStream_t stream = nullptr) {
|
||||
const dim3 block = get_block_shape();
|
||||
const dim3 grid = get_grid_shape(params_, block);
|
||||
|
||||
int smem_size = int(sizeof(typename GemvKernel::SharedStorage));
|
||||
|
||||
cutlass::arch::synclog_setup();
|
||||
cutlass::Kernel<GemvKernel><<<grid, block, smem_size, stream>>>(params_);
|
||||
|
||||
cudaError_t result = cudaGetLastError();
|
||||
if (result == cudaSuccess) {
|
||||
return Status::kSuccess;
|
||||
} else {
|
||||
return Status::kErrorInternal;
|
||||
}
|
||||
}
|
||||
|
||||
/// Runs the kernel using initialized state.
|
||||
Status operator()(cudaStream_t stream = nullptr) {
|
||||
return run(stream);
|
||||
}
|
||||
|
||||
/// Runs the kernel using initialized state.
|
||||
Status operator()(
|
||||
Arguments const &args,
|
||||
void *workspace = nullptr,
|
||||
cudaStream_t stream = nullptr) {
|
||||
|
||||
Status status = initialize(args, workspace, stream);
|
||||
|
||||
if (status == Status::kSuccess) {
|
||||
status = run(stream);
|
||||
}
|
||||
|
||||
return status;
|
||||
}
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace device
|
||||
} // namespace gemm
|
||||
} // namespace cutlass
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
@@ -36,6 +36,7 @@
|
||||
#include "cute/layout.hpp"
|
||||
#include "cute/numeric/integral_constant.hpp" // cute::false_type
|
||||
#include "cute/atom/copy_traits_sm100.hpp"
|
||||
#include "cutlass/detail/collective/sm103_kernel_type.hpp"
|
||||
//////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace cutlass::detail {
|
||||
@@ -423,6 +424,21 @@ struct MainloopSm90ArrayTmaGmmaWarpSpecializedBlockScaling
|
||||
"KernelSchedule must be one of the warp specialized FP8 block scale policies");
|
||||
};
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
//
|
||||
// Kernel Scheduler Tag
|
||||
//
|
||||
|
||||
// Dense GEMM: SM100 tensor op policy that applies to both 1SM and 2SM MMA atoms
|
||||
template<
|
||||
int SchedulerPipelineStageCount_,
|
||||
int AccumulatorPipelineStageCount_
|
||||
>
|
||||
struct KernelWarpSpecializedSm100 final {
|
||||
static constexpr int SchedulerPipelineStageCount = SchedulerPipelineStageCount_;
|
||||
static constexpr int AccumulatorPipelineStageCount = AccumulatorPipelineStageCount_;
|
||||
};
|
||||
|
||||
template<
|
||||
int SchedulerPipelineStageCount_,
|
||||
@@ -461,6 +477,24 @@ struct KernelPtrArrayTmaWarpSpecializedMmaTransformSm100 final {
|
||||
static constexpr int AccumulatorPipelineStageCount = AccumulatorPipelineStageCount_;
|
||||
};
|
||||
|
||||
template<
|
||||
int SchedulerPipelineStageCount_,
|
||||
int AccumulatorPipelineStageCount_
|
||||
>
|
||||
struct KernelTmaWarpSpecializedBlockScaledSm103 final {
|
||||
static constexpr int SchedulerPipelineStageCount = SchedulerPipelineStageCount_;
|
||||
static constexpr int AccumulatorPipelineStageCount = AccumulatorPipelineStageCount_;
|
||||
};
|
||||
|
||||
template<
|
||||
int SchedulerPipelineStageCount_,
|
||||
int AccumulatorPipelineStageCount_
|
||||
>
|
||||
struct KernelPtrArrayTmaWarpSpecializedBlockScaledSm103 final {
|
||||
static constexpr int SchedulerPipelineStageCount = SchedulerPipelineStageCount_;
|
||||
static constexpr int AccumulatorPipelineStageCount = AccumulatorPipelineStageCount_;
|
||||
};
|
||||
|
||||
// Sparse Gemm
|
||||
template<
|
||||
int SchedulerPipelineStageCount_,
|
||||
@@ -665,6 +699,8 @@ struct KernelScheduleSm100DenseGemm : KernelScheduleSm100 {}; // Base policy
|
||||
// Dense GEMM: Specialize for 1SM vs 2SM
|
||||
struct KernelTmaWarpSpecialized1SmSm100 final : KernelSchedule1Sm, KernelScheduleSm100DenseGemm {}; // Use for 1SM Dense GEMM Kernels for Collective Mainloop Builder
|
||||
struct KernelTmaWarpSpecialized2SmSm100 final : KernelSchedule2Sm, KernelScheduleSm100DenseGemm {}; // Use for 2SM Dense GEMM Kernels for Collective Mainloop Builder
|
||||
struct KernelWarpSpecialized1SmSm100 final : KernelSchedule1Sm, KernelScheduleSm100DenseGemm {}; // Use for 1SM Dense GEMM Kernels for Collective Mainloop Builder Without TMA
|
||||
|
||||
///////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
// SM100 Ptr-Array Dense GEMM Dispatch Policies
|
||||
///////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
@@ -721,6 +757,8 @@ struct KernelScheduleSm100MixedInputGemm : KernelScheduleSm100 {};
|
||||
struct KernelTmaWarpSpecializedMixedInputSmemSm100 : KernelScheduleSm100MixedInputGemm { };
|
||||
struct KernelTmaWarpSpecialized1SmMixedInputSm100 final : KernelSchedule1Sm, KernelScheduleSm100MixedInputGemm { };
|
||||
struct KernelTmaWarpSpecialized1SmMixedInputSmemSm100 final : KernelSchedule1Sm, KernelTmaWarpSpecializedMixedInputSmemSm100 { };
|
||||
struct KernelTmaWarpSpecialized2SmMixedInputSm100 final : KernelSchedule2Sm, KernelScheduleSm100MixedInputGemm { };
|
||||
struct KernelTmaWarpSpecialized2SmMixedInputSmemSm100 final : KernelSchedule2Sm, KernelTmaWarpSpecializedMixedInputSmemSm100 { };
|
||||
|
||||
///////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
// SM100 Ptr-Array FastF32 (9xBF16) GEMM Dispatch Policies
|
||||
@@ -789,6 +827,54 @@ struct KernelSparseTmaWarpSpecialized2SmNvf4Sm100 final : KernelSchedule2
|
||||
struct KernelSparseTmaWarpSpecialized1SmMxf4Sm100 final : KernelSchedule1Sm, KernelScheduleSparseMxNvf4Sm100 { };
|
||||
struct KernelSparseTmaWarpSpecialized2SmMxf4Sm100 final : KernelSchedule2Sm, KernelScheduleSparseMxNvf4Sm100 { };
|
||||
|
||||
///////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
//
|
||||
// SM103 Dispatch Policies
|
||||
//
|
||||
///////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
struct KernelScheduleSm103 {};
|
||||
struct KernelScheduleSm103BlockScaledGemm : KernelScheduleSm103 {};
|
||||
struct KernelScheduleSm103BlockScaledMxNvf4UltraTmaPrefetch : KernelScheduleSm103BlockScaledGemm {};
|
||||
struct KernelScheduleSm103BlockScaledMxNvf4UltraDisablePrefetch : KernelScheduleSm103BlockScaledGemm {};
|
||||
|
||||
// Blockscaled Gemm: Specialized for instruction type, scale factor vector size, and 1SM vs. 2SM
|
||||
// These are the public dispatch policy name
|
||||
struct KernelTmaWarpSpecialized1SmBlockScaledMxNvf4UltraVs16Sm103TmaPrefetch final : KernelSchedule1Sm, KernelScheduleSm103BlockScaledMxNvf4UltraTmaPrefetch { };
|
||||
struct KernelTmaWarpSpecialized2SmBlockScaledMxNvf4UltraVs16Sm103TmaPrefetch final : KernelSchedule2Sm, KernelScheduleSm103BlockScaledMxNvf4UltraTmaPrefetch { };
|
||||
struct KernelTmaWarpSpecialized1SmBlockScaledMxNvf4UltraVs32Sm103TmaPrefetch final : KernelSchedule1Sm, KernelScheduleSm103BlockScaledMxNvf4UltraTmaPrefetch { };
|
||||
struct KernelTmaWarpSpecialized2SmBlockScaledMxNvf4UltraVs32Sm103TmaPrefetch final : KernelSchedule2Sm, KernelScheduleSm103BlockScaledMxNvf4UltraTmaPrefetch { };
|
||||
|
||||
struct KernelTmaWarpSpecialized1SmBlockScaledMxNvf4UltraVs16Sm103DisablePrefetch final : KernelSchedule1Sm, KernelScheduleSm103BlockScaledMxNvf4UltraDisablePrefetch { };
|
||||
struct KernelTmaWarpSpecialized2SmBlockScaledMxNvf4UltraVs16Sm103DisablePrefetch final : KernelSchedule2Sm, KernelScheduleSm103BlockScaledMxNvf4UltraDisablePrefetch { };
|
||||
struct KernelTmaWarpSpecialized1SmBlockScaledMxNvf4UltraVs32Sm103DisablePrefetch final : KernelSchedule1Sm, KernelScheduleSm103BlockScaledMxNvf4UltraDisablePrefetch { };
|
||||
struct KernelTmaWarpSpecialized2SmBlockScaledMxNvf4UltraVs32Sm103DisablePrefetch final : KernelSchedule2Sm, KernelScheduleSm103BlockScaledMxNvf4UltraDisablePrefetch { };
|
||||
|
||||
using KernelTmaWarpSpecialized1SmBlockScaledMxNvf4UltraVs16Sm103 = KernelTmaWarpSpecialized1SmBlockScaledMxNvf4UltraVs16Sm103TmaPrefetch;
|
||||
using KernelTmaWarpSpecialized2SmBlockScaledMxNvf4UltraVs16Sm103 = KernelTmaWarpSpecialized2SmBlockScaledMxNvf4UltraVs16Sm103TmaPrefetch;
|
||||
using KernelTmaWarpSpecialized1SmBlockScaledMxNvf4UltraVs32Sm103 = KernelTmaWarpSpecialized1SmBlockScaledMxNvf4UltraVs32Sm103TmaPrefetch;
|
||||
using KernelTmaWarpSpecialized2SmBlockScaledMxNvf4UltraVs32Sm103 = KernelTmaWarpSpecialized2SmBlockScaledMxNvf4UltraVs32Sm103TmaPrefetch;
|
||||
|
||||
|
||||
struct KernelSchedulePtrArraySm103BlockScaledGemm : KernelScheduleSm103 {};
|
||||
struct KernelSchedulePtrArraySm103BlockScaledMxNvf4UltraTmaPrefetch : KernelSchedulePtrArraySm103BlockScaledGemm {};
|
||||
struct KernelSchedulePtrArraySm103BlockScaledMxNvf4UltraDisablePrefetch : KernelSchedulePtrArraySm103BlockScaledGemm {};
|
||||
|
||||
struct KernelPtrArrayTmaWarpSpecialized1SmBlockScaledMxNvf4UltraVs16Sm103TmaPrefetch final : KernelSchedule1Sm, KernelSchedulePtrArraySm103BlockScaledMxNvf4UltraTmaPrefetch { };
|
||||
struct KernelPtrArrayTmaWarpSpecialized2SmBlockScaledMxNvf4UltraVs16Sm103TmaPrefetch final : KernelSchedule2Sm, KernelSchedulePtrArraySm103BlockScaledMxNvf4UltraTmaPrefetch { };
|
||||
struct KernelPtrArrayTmaWarpSpecialized1SmBlockScaledMxNvf4UltraVs32Sm103TmaPrefetch final : KernelSchedule1Sm, KernelSchedulePtrArraySm103BlockScaledMxNvf4UltraTmaPrefetch { };
|
||||
struct KernelPtrArrayTmaWarpSpecialized2SmBlockScaledMxNvf4UltraVs32Sm103TmaPrefetch final : KernelSchedule2Sm, KernelSchedulePtrArraySm103BlockScaledMxNvf4UltraTmaPrefetch { };
|
||||
|
||||
struct KernelPtrArrayTmaWarpSpecialized1SmBlockScaledMxNvf4UltraVs16Sm103DisablePrefetch final : KernelSchedule1Sm, KernelSchedulePtrArraySm103BlockScaledMxNvf4UltraDisablePrefetch { };
|
||||
struct KernelPtrArrayTmaWarpSpecialized2SmBlockScaledMxNvf4UltraVs16Sm103DisablePrefetch final : KernelSchedule2Sm, KernelSchedulePtrArraySm103BlockScaledMxNvf4UltraDisablePrefetch { };
|
||||
struct KernelPtrArrayTmaWarpSpecialized1SmBlockScaledMxNvf4UltraVs32Sm103DisablePrefetch final : KernelSchedule1Sm, KernelSchedulePtrArraySm103BlockScaledMxNvf4UltraDisablePrefetch { };
|
||||
struct KernelPtrArrayTmaWarpSpecialized2SmBlockScaledMxNvf4UltraVs32Sm103DisablePrefetch final : KernelSchedule2Sm, KernelSchedulePtrArraySm103BlockScaledMxNvf4UltraDisablePrefetch { };
|
||||
|
||||
using KernelPtrArrayTmaWarpSpecialized1SmBlockScaledMxNvf4UltraVs16Sm103 = KernelPtrArrayTmaWarpSpecialized1SmBlockScaledMxNvf4UltraVs16Sm103DisablePrefetch;
|
||||
using KernelPtrArrayTmaWarpSpecialized2SmBlockScaledMxNvf4UltraVs16Sm103 = KernelPtrArrayTmaWarpSpecialized2SmBlockScaledMxNvf4UltraVs16Sm103DisablePrefetch;
|
||||
using KernelPtrArrayTmaWarpSpecialized1SmBlockScaledMxNvf4UltraVs32Sm103 = KernelPtrArrayTmaWarpSpecialized1SmBlockScaledMxNvf4UltraVs32Sm103DisablePrefetch;
|
||||
using KernelPtrArrayTmaWarpSpecialized2SmBlockScaledMxNvf4UltraVs32Sm103 = KernelPtrArrayTmaWarpSpecialized2SmBlockScaledMxNvf4UltraVs32Sm103DisablePrefetch;
|
||||
|
||||
///////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
//
|
||||
// SM120 Dispatch Policies
|
||||
@@ -844,6 +930,25 @@ struct KernelSparseTmaWarpSpecializedMxf4Sm120 final : KernelScheduleS
|
||||
struct KernelSparseTmaWarpSpecializedMxf8f6f4Sm120 final : KernelScheduleSparseMxf8f6f4Sm120 { };
|
||||
struct KernelSparseTmaWarpSpecializedMxf8f6f4Acc2x4Sm120 final : KernelScheduleSparseMxf8f6f4Sm120, KernelScheduleAcc2x4Sm120 { };
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
//
|
||||
// Collective Mainloop Dispatch Policies
|
||||
//
|
||||
|
||||
// n-buffer in smem, pipelined with Blackwell UMMA and CPASYNC, Warp specialized dynamic schedule
|
||||
template<
|
||||
int Stages_,
|
||||
int SchedulerPipelineStageCount_,
|
||||
int AccumulatorPipelineStageCount_,
|
||||
class ClusterShape_ = Shape<_1,_1,_1>
|
||||
>
|
||||
struct MainloopSm100UmmaCpAsyncWarpSpecialized {
|
||||
constexpr static int Stages = Stages_;
|
||||
using ClusterShape = ClusterShape_;
|
||||
using ArchTag = arch::Sm100;
|
||||
using Schedule = KernelWarpSpecializedSm100<SchedulerPipelineStageCount_, AccumulatorPipelineStageCount_>;
|
||||
};
|
||||
|
||||
// n-buffer in smem, pipelined with Blackwell UMMA and TMA, Warp specialized dynamic schedule
|
||||
template<
|
||||
@@ -983,6 +1088,34 @@ struct MainloopSm100TmaUmmaWarpSpecializedFastF32 {
|
||||
};
|
||||
|
||||
|
||||
// n-buffer in smem, pipelined with Blackwell Mixed Input kernel with UMMA (HwScaled) and TMA,
|
||||
template<
|
||||
// Number of Pipeline stages for
|
||||
// MainloopLoad <-> Conversion <-> MainLoad
|
||||
int Load2TransformPipelineStageCount_,
|
||||
// Number of Pipeline stages for
|
||||
// MainloopLoad <-> Conversion <-> MainLoad
|
||||
int Transform2MmaPipelineStageCount_,
|
||||
// TileScheduler pipeline depth
|
||||
int SchedulerPipelineStageCount_,
|
||||
// Accmulator pipeline depth
|
||||
int AccumulatorPipelineStageCount_,
|
||||
// ClusterShape for the kernel
|
||||
class ClusterShape_ = Shape<_1,_1,_1>
|
||||
>
|
||||
struct MainloopSm100TmaUmmaWarpSpecializedMixedInput {
|
||||
constexpr static int Load2TransformPipelineStageCount = Load2TransformPipelineStageCount_;
|
||||
constexpr static int Load2MmaPipelineStageCount = Load2TransformPipelineStageCount_;
|
||||
constexpr static int Transform2MmaPipelineStageCount = Transform2MmaPipelineStageCount_;
|
||||
constexpr static detail::KernelInputTransformType InputTransformType = detail::KernelInputTransformType::MixedInput;
|
||||
using ClusterShape = ClusterShape_;
|
||||
using ArchTag = arch::Sm100;
|
||||
using Schedule = KernelTmaWarpSpecializedMixedInputTransformSm100<SchedulerPipelineStageCount_, AccumulatorPipelineStageCount_>;
|
||||
|
||||
// For backwards compatibility with GemmUniversalAdapter.
|
||||
constexpr static int Stages = Load2TransformPipelineStageCount;
|
||||
};
|
||||
|
||||
|
||||
// n-buffer in smem, pipelined with Blackwell UMMA and TMA, Warp specialized dynamic schedule
|
||||
template<
|
||||
@@ -1064,9 +1197,49 @@ struct MainloopSm100ArrayTmaUmmaWarpSpecializedFastF32 {
|
||||
};
|
||||
|
||||
|
||||
// n-buffer in smem, pipelined with Blackwell UMMA and TMA, Warp specialized dynamic schedule
|
||||
template<
|
||||
int LoadABPipelineStageCount_,
|
||||
int LoadSFPipelineStageCount_,
|
||||
int SchedulerPipelineStageCount_,
|
||||
int AccumulatorPipelineStageCount_,
|
||||
class ClusterShape_ = Shape<_1,_1,_1>,
|
||||
cutlass::sm103::detail::KernelPrefetchType PrefetchType_ = cutlass::sm103::detail::KernelPrefetchType::TmaPrefetch
|
||||
>
|
||||
struct MainloopSm103TmaUmmaWarpSpecializedBlockScaled {
|
||||
constexpr static int LoadABPipelineStageCount = LoadABPipelineStageCount_;
|
||||
constexpr static int LoadSFPipelineStageCount = LoadSFPipelineStageCount_;
|
||||
using ClusterShape = ClusterShape_;
|
||||
using ArchTag = arch::Sm103;
|
||||
constexpr static bool IsOverlappingAccum = AccumulatorPipelineStageCount_ == 1;
|
||||
using Schedule = KernelTmaWarpSpecializedBlockScaledSm103<SchedulerPipelineStageCount_, AccumulatorPipelineStageCount_>;
|
||||
// For backwards compatibility with GemmUniversalAdapter.
|
||||
constexpr static int Stages = LoadABPipelineStageCount;
|
||||
constexpr static cutlass::sm103::detail::KernelPrefetchType PrefetchType = PrefetchType_;
|
||||
};
|
||||
|
||||
// Mainloop schedule for array-based TMA
|
||||
|
||||
template<
|
||||
int LoadABPipelineStageCount_,
|
||||
int LoadSFPipelineStageCount_,
|
||||
int SchedulerPipelineStageCount_,
|
||||
int AccumulatorPipelineStageCount_,
|
||||
class ClusterShape_ = Shape<_1,_1,_1>,
|
||||
cutlass::sm103::detail::KernelPrefetchType PrefetchType_ = cutlass::sm103::detail::KernelPrefetchType::TmaPrefetch
|
||||
>
|
||||
struct MainloopSm103ArrayTmaUmmaWarpSpecializedBlockScaled {
|
||||
constexpr static int LoadABPipelineStageCount = LoadABPipelineStageCount_;
|
||||
constexpr static int LoadSFPipelineStageCount = LoadSFPipelineStageCount_;
|
||||
using ClusterShape = ClusterShape_;
|
||||
using ArchTag = arch::Sm103;
|
||||
constexpr static bool IsOverlappingAccum = AccumulatorPipelineStageCount_ == 1;
|
||||
using Schedule = KernelPtrArrayTmaWarpSpecializedBlockScaledSm103<SchedulerPipelineStageCount_, AccumulatorPipelineStageCount_>;
|
||||
// For backwards compatibility with GemmUniversalAdapter.
|
||||
constexpr static int Stages = LoadABPipelineStageCount;
|
||||
constexpr static cutlass::sm103::detail::KernelPrefetchType PrefetchType = PrefetchType_;
|
||||
};
|
||||
|
||||
template<
|
||||
int Stages_,
|
||||
int SchedulerPipelineStageCount_,
|
||||
|
||||
@@ -68,8 +68,13 @@ struct IsCutlass3ArrayKernel<ProblemShape, cute::void_t<typename ProblemShape::U
|
||||
#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_tma_warpspecialized_mixed_input_transform.hpp"
|
||||
#include "cutlass/gemm/kernel/sm100_gemm_array_tma_warpspecialized_input_transform.hpp"
|
||||
#include "cutlass/gemm/kernel/sm100_gemm_array_tma_warpspecialized_mma_transform.hpp"
|
||||
#include "cutlass/gemm/kernel/sm100_sparse_gemm_tma_warpspecialized.hpp"
|
||||
#include "cutlass/gemm/kernel/sm100_gemm_cpasync_warpspecialized.hpp"
|
||||
#include "cutlass/gemm/kernel/sm103_blockscaled_gemm_tma_warpspecialized.hpp"
|
||||
#include "cutlass/gemm/kernel/sm103_blockscaled_gemm_array_tma_warpspecialized.hpp"
|
||||
#include "cutlass/gemm/kernel/sm120_gemm_tma_warpspecialized_cooperative_asymmetric_dma.hpp"
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
@@ -0,0 +1,885 @@
|
||||
/***************************************************************************************************
|
||||
* Copyright (c) 2024 - 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
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "cutlass/arch/cache_operation.h" /// cutlass::arch::CacheOperation
|
||||
#include "cutlass/arch/memory.h" // cutlass::arch::global_load
|
||||
#include "cutlass/arch/memory_sm80.h" // cp.async helpers, ldsm, cp_async_wait
|
||||
#include "cutlass/complex.h" // cutlass::ComplexTransform:
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cutlass/fast_math.h" // cutlass::fast_max
|
||||
#include "cutlass/layout/matrix.h" // cutlass::layout::RowMajor
|
||||
#include "cutlass/matrix_coord.h" // cutlass::MatrixCoord
|
||||
#include "cutlass/numeric_conversion.h" // cutlass::FloatRoundStyle, cutlass::NumericConverter
|
||||
#include "cutlass/numeric_types.h" // cutlass::float_e4m3_t
|
||||
#include "cutlass/platform/platform.h" // cutlass::is_same_v
|
||||
#include "cutlass/tensor_ref.h" // cutlass::TensorRef
|
||||
#include "cutlass/semaphore.h" // split-k
|
||||
|
||||
#include "cute/algorithm/functional.hpp" // cute::for_each
|
||||
#include "cute/numeric/arithmetic_tuple.hpp" // cute::make_int_sequence
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace cutlass {
|
||||
namespace gemm {
|
||||
namespace kernel {
|
||||
|
||||
using namespace cute;
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <
|
||||
typename ElementA_,
|
||||
typename LayoutA_,
|
||||
typename ElementB_,
|
||||
typename ElementC_,
|
||||
typename ElementAccumulator_,
|
||||
typename EpilogueOutputOp_,
|
||||
int kElementsPerAccess_ = 1, ///< Number of elements involved in a global access.
|
||||
int kThreadCount_ = 0, ///< Number of threads in the thread block.
|
||||
/// It will be calculated automatically if set to 0.
|
||||
int kThreadsPerRow_ = 0, ///< Number of threads in the k dimension.
|
||||
/// It will be calculated automatically if set to 0.
|
||||
typename ElementSFA_ = cutlass::float_e4m3_t,
|
||||
typename ElementSFB_ = cutlass::float_e4m3_t,
|
||||
int kSFVecSize_ = 16
|
||||
>
|
||||
struct GemvBlockScaled;
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
//
|
||||
// Specializations
|
||||
//
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
// GEMV for row-major A matrix
|
||||
template <typename ElementA_,
|
||||
typename ElementB_,
|
||||
typename ElementC_,
|
||||
typename ElementAccumulator_,
|
||||
typename EpilogueOutputOp_,
|
||||
int kElementsPerAccess_,
|
||||
int kThreadCount_,
|
||||
int kThreadsPerRow_,
|
||||
typename ElementSFA_,
|
||||
typename ElementSFB_,
|
||||
int kSFVecSize_>
|
||||
struct GemvBlockScaled<ElementA_,
|
||||
cutlass::layout::RowMajor,
|
||||
ElementB_,
|
||||
ElementC_,
|
||||
ElementAccumulator_,
|
||||
EpilogueOutputOp_,
|
||||
kElementsPerAccess_,
|
||||
kThreadCount_,
|
||||
kThreadsPerRow_,
|
||||
ElementSFA_,
|
||||
ElementSFB_,
|
||||
kSFVecSize_>
|
||||
{
|
||||
public:
|
||||
using ElementA = ElementA_;
|
||||
using ElementSFA = ElementSFA_;
|
||||
using LayoutA = cutlass::layout::RowMajor;
|
||||
using TensorRefA = cutlass::TensorRef<ElementA, LayoutA>;
|
||||
static_assert(cutlass::sizeof_bits<ElementSFA>::value == 8, "ElementSFA should be FP8 type");
|
||||
|
||||
using ElementB = ElementB_;
|
||||
using ElementSFB = ElementSFB_;
|
||||
using LayoutB = cutlass::layout::ColumnMajor;
|
||||
static_assert(cutlass::sizeof_bits<ElementSFB>::value == 8, "ElementSFB should be FP8 type");
|
||||
|
||||
using ElementC = ElementC_;
|
||||
using LayoutC = cutlass::layout::ColumnMajor;
|
||||
|
||||
using ElementAccumulator = ElementAccumulator_;
|
||||
|
||||
static constexpr cutlass::ComplexTransform kTransformA = cutlass::ComplexTransform::kNone;
|
||||
static constexpr cutlass::ComplexTransform kTransformB = cutlass::ComplexTransform::kNone;
|
||||
|
||||
static constexpr FloatRoundStyle Round = cutlass::FloatRoundStyle::round_to_nearest;
|
||||
|
||||
// number of return elements in a global access
|
||||
static constexpr int kElementsPerAccess = kElementsPerAccess_;
|
||||
static constexpr int kSFVecSize = kSFVecSize_;
|
||||
static constexpr int kSFPerAccess = cutlass::const_max(1, kElementsPerAccess / kSFVecSize);
|
||||
|
||||
static_assert(kSFVecSize == 16, "Only SFVecSize = 16 is supported");
|
||||
// Hardcode some check for easier debug
|
||||
static_assert(kElementsPerAccess == 32, "for fp4 kernel, 32 elt per access");
|
||||
static_assert(kSFPerAccess == 2, "fpr fp4 kernel, 2 sf read per thread");
|
||||
|
||||
static constexpr bool kDequantizeA = cutlass::sizeof_bits<ElementA>::value == 4;
|
||||
static constexpr bool kDequantizeB = cutlass::sizeof_bits<ElementB>::value == 4;
|
||||
static constexpr int kPackedElementsA = cutlass::sizeof_bits<ElementA>::value == 4 ? 2 : 1;
|
||||
static constexpr int kPackedElementsB = cutlass::sizeof_bits<ElementB>::value == 4 ? 2 : 1;
|
||||
static constexpr int kPackedElements = cutlass::const_max(kPackedElementsA, kPackedElementsB);
|
||||
|
||||
static_assert(kDequantizeA == true, "kDequantizeA should be true");
|
||||
static_assert(kDequantizeB == true, "kDequantizeB should be true");
|
||||
|
||||
using FragmentA = cutlass::Array<ElementA, kElementsPerAccess>;
|
||||
using FragmentB = cutlass::Array<ElementB, kElementsPerAccess>;
|
||||
using FragmentCompute = cutlass::Array<ElementAccumulator, kElementsPerAccess>;
|
||||
using FragmentSFA = cutlass::Array<ElementSFA, kSFPerAccess>;
|
||||
using FragmentSFB = cutlass::Array<ElementSFB, kSFPerAccess>;
|
||||
using FragmentPackedA = cutlass::Array<ElementA, kPackedElements>;
|
||||
using FragmentPackedB = cutlass::Array<ElementB, kPackedElements>;
|
||||
|
||||
static_assert(sizeof_bits<FragmentA>::value == 128, "FragmentA should be 128 bits");
|
||||
static_assert(sizeof_bits<FragmentB>::value == 128, "FragmentB should be 128 bits");
|
||||
|
||||
// // thread block shape (kThreadsPerRow, kThreadCount / kThreadsPerRow, 1)
|
||||
static constexpr int kThreadCount = (kThreadCount_ <= 0) ? 128 : kThreadCount_;
|
||||
static constexpr int kThreadsPerRow = (kThreadsPerRow_ <= 0) ?
|
||||
cutlass::const_min(static_cast<int>(kThreadCount / cutlass::bits_to_bytes(kElementsPerAccess * cutlass::sizeof_bits<ElementA>::value)), 16) :
|
||||
kThreadsPerRow_;
|
||||
static constexpr int kThreadsPerCol = kThreadCount / kThreadsPerRow;
|
||||
|
||||
static constexpr int kStageCount = 4;
|
||||
static constexpr int kBufferCount = 2;
|
||||
|
||||
// Number of elements stored in shared memory per stage for operands A and B.
|
||||
// Each thread contributes `kElementsPerAccess / kPackedElements{A,B}` packed
|
||||
// values.
|
||||
static constexpr int kSmemPerStageA = kThreadCount * kElementsPerAccess / kPackedElementsA;
|
||||
// B is uniform across all threads in the same k-column, so only store it once per k-thread
|
||||
static constexpr int kSmemPerStageB = kThreadsPerRow * kElementsPerAccess / kPackedElementsB;
|
||||
|
||||
using EpilogueOutputOp = EpilogueOutputOp_;
|
||||
|
||||
// Ensure epilogue and mainloop have same thread layout
|
||||
static_assert(kThreadCount == EpilogueOutputOp::kThreadCount, "mainloop, epilogue thread count mismatch");
|
||||
static_assert(kThreadsPerRow == EpilogueOutputOp::kThreadsPerRow, "mainloop, epilogue thread per row mismatch");
|
||||
static_assert(kThreadsPerCol == EpilogueOutputOp::kThreadsPerCol, "mainloop, epilogue thread per col mismatch");
|
||||
|
||||
//
|
||||
// Structures
|
||||
//
|
||||
|
||||
/// Argument structure
|
||||
struct Arguments
|
||||
{
|
||||
MatrixCoord problem_size;
|
||||
int32_t batch_count{0};
|
||||
typename EpilogueOutputOp::Params epilogue;
|
||||
|
||||
TensorRefA ref_A;
|
||||
|
||||
ElementB const *ptr_B{nullptr};
|
||||
ElementC const *ptr_C{nullptr};
|
||||
ElementC *ptr_D{nullptr};
|
||||
|
||||
ElementSFA const *ptr_SFA{nullptr};
|
||||
ElementSFB const *ptr_SFB{nullptr};
|
||||
|
||||
int64_t stride_A{0};
|
||||
int64_t batch_stride_A{0};
|
||||
int64_t batch_stride_B{0};
|
||||
int64_t batch_stride_C{0};
|
||||
int64_t batch_stride_D{0};
|
||||
|
||||
int64_t batch_stride_SFA{0};
|
||||
int64_t batch_stride_SFB{0};
|
||||
int64_t batch_stride_SFD{0};
|
||||
};
|
||||
|
||||
using Params = Arguments;
|
||||
|
||||
/// Shared memory storage structure
|
||||
struct SharedStorage
|
||||
{
|
||||
using EpilogueStorage = typename EpilogueOutputOp::SharedStorage;
|
||||
EpilogueStorage epilogue;
|
||||
|
||||
alignas(16) ElementA smem_A[kBufferCount][kStageCount][kSmemPerStageA];
|
||||
alignas(16) ElementB smem_B[kBufferCount][kStageCount][kSmemPerStageB];
|
||||
alignas(16) ElementSFA smem_SFA[kBufferCount][kStageCount][kThreadCount * kSFPerAccess];
|
||||
alignas(16) ElementSFB smem_SFB[kBufferCount][kStageCount][kThreadsPerRow * kSFPerAccess];
|
||||
};
|
||||
|
||||
public:
|
||||
//
|
||||
// Methods
|
||||
//
|
||||
/// Determines whether kernel satisfies alignment
|
||||
static Status can_implement(cutlass::MatrixCoord const &problem_size)
|
||||
{
|
||||
if (problem_size.column() % kElementsPerAccess != 0) {
|
||||
return Status::kErrorMisalignedOperand;
|
||||
}
|
||||
return Status::kSuccess;
|
||||
}
|
||||
|
||||
static Status can_implement(Arguments const &args)
|
||||
{
|
||||
return can_implement(args.problem_size);
|
||||
}
|
||||
|
||||
/// Executes one GEMV
|
||||
CUTLASS_DEVICE
|
||||
void operator()(Params const ¶ms, SharedStorage &shared_storage)
|
||||
{
|
||||
EpilogueOutputOp epilogue(params.epilogue, shared_storage.epilogue);
|
||||
|
||||
// Converters only needed for regular GEMV fallback case
|
||||
NumericConverter<ElementAccumulator, ElementA, Round> A_converter;
|
||||
NumericConverter<ElementAccumulator, ElementB, Round> B_converter;
|
||||
NumericConverter<ElementAccumulator, ElementSFA, Round> SFA_converter;
|
||||
NumericConverter<ElementAccumulator, ElementSFB, Round> SFB_converter;
|
||||
|
||||
const int32_t gemm_m = params.problem_size.row();
|
||||
[[maybe_unused]] static constexpr int32_t gemm_n = 1;
|
||||
const int32_t gemm_k = params.problem_size.column();
|
||||
const int32_t gemm_batch = params.batch_count;
|
||||
|
||||
// Loop over batch indices
|
||||
for (int batch_idx = blockIdx.z; batch_idx < gemm_batch; batch_idx += gridDim.z) {
|
||||
|
||||
int idx_col_k = threadIdx.x;
|
||||
int idx_row_m = blockIdx.x * blockDim.y + threadIdx.y;
|
||||
|
||||
if (idx_row_m < gemm_m) {
|
||||
// problem_size (row = m, column = k)
|
||||
// matrix A (batch, m, k)
|
||||
// vector B (batch, k, 1)
|
||||
// vector C (batch, m, 1)
|
||||
// vector D (batch, m, 1)
|
||||
// move in the batch dimension
|
||||
ElementA const *ptr_A = params.ref_A.data() + batch_idx * params.batch_stride_A / kPackedElementsA;
|
||||
ElementB const *ptr_B = params.ptr_B + batch_idx * params.batch_stride_B / kPackedElementsB;
|
||||
ElementC const *ptr_C = params.ptr_C + batch_idx * params.batch_stride_C;
|
||||
ElementC *ptr_D = params.ptr_D + batch_idx * params.batch_stride_D;
|
||||
|
||||
// move in the k dimension
|
||||
ptr_A += idx_col_k * kElementsPerAccess / kPackedElementsA;
|
||||
ptr_B += idx_col_k * kElementsPerAccess / kPackedElementsB;
|
||||
|
||||
// move in the m dimension
|
||||
ptr_A += idx_row_m * params.stride_A / kPackedElementsA;
|
||||
ptr_C += idx_row_m;
|
||||
ptr_D += idx_row_m;
|
||||
|
||||
ElementSFA const *ptr_SF_A{nullptr};
|
||||
ElementSFB const *ptr_SF_B{nullptr};
|
||||
int global_k{0};
|
||||
|
||||
int SF_blocks_by_M = (gemm_m + 127) >> 7;
|
||||
int SF_blocks_by_K = (gemm_k / kSFVecSize + 3) >> 2;
|
||||
|
||||
// move in the batch dimension
|
||||
ptr_SF_A = params.ptr_SFA + batch_idx * SF_blocks_by_M * SF_blocks_by_K * 512;
|
||||
ptr_SF_B = params.ptr_SFB + batch_idx * SF_blocks_by_K * 512;
|
||||
|
||||
// move in the m dimension
|
||||
ptr_SF_A += (((idx_row_m >> 7) * SF_blocks_by_K) << 9) + ((idx_row_m & 0x1f) << 4) + ((idx_row_m & 0x7f) >> 5 << 2);
|
||||
|
||||
global_k = idx_col_k * kElementsPerAccess;
|
||||
|
||||
ElementAccumulator accum = ElementAccumulator(0);
|
||||
|
||||
// Local aliases
|
||||
const int tileA_k_local = kThreadsPerRow * kElementsPerAccess;
|
||||
const int total_tiles = gemm_k / tileA_k_local;
|
||||
|
||||
int unroll_col_k = 0; // total K elements consumed so far by this thread
|
||||
const int thread_id = threadIdx.y * kThreadsPerRow + threadIdx.x;
|
||||
const bool is_even_thread = (threadIdx.x % 2 == 0);
|
||||
const bool load_b = (threadIdx.y == 0);
|
||||
const int smem_sf_write_offset = (thread_id / 2) * 4; // 4 FP8 per even thread
|
||||
const int smem_sf_offset = thread_id * kSFPerAccess;
|
||||
|
||||
// Fast path: if the problem fits entirely in the tail path, skip SMEM
|
||||
if (total_tiles == 0) {
|
||||
accum += process_tail_elements(0, idx_col_k, gemm_k,
|
||||
ptr_A, ptr_B,
|
||||
ptr_SF_A, ptr_SF_B,
|
||||
A_converter, B_converter,
|
||||
SFA_converter, SFB_converter);
|
||||
} else {
|
||||
|
||||
// Scaling factors are now loaded from shared memory, no register pipeline needed
|
||||
|
||||
// Thread-local SMEM line offset
|
||||
const int thread_linear = threadIdx.y * kThreadsPerRow + threadIdx.x;
|
||||
const int smem_offset_A = thread_linear * (kElementsPerAccess / kPackedElementsA);
|
||||
// Only one row of threads (threadIdx.y == 0) loads B
|
||||
const int smem_offset_B = threadIdx.x * (kElementsPerAccess / kPackedElementsB);
|
||||
|
||||
// PROLOGUE – prime first kStageCount-1 stages into buffer 0
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int b = 0; b < kBufferCount - 1; ++b) {
|
||||
// Load all stages using the helper function
|
||||
load_stages_gmem_to_smem(
|
||||
b, // buffer_idx
|
||||
kStageCount, // num_stages
|
||||
unroll_col_k, // passed by reference
|
||||
global_k, // passed by reference
|
||||
tileA_k_local,
|
||||
smem_offset_A,
|
||||
smem_offset_B,
|
||||
smem_sf_write_offset,
|
||||
is_even_thread,
|
||||
load_b,
|
||||
true, // valid_tile = true for prologue
|
||||
ptr_A,
|
||||
ptr_B,
|
||||
ptr_SF_A,
|
||||
ptr_SF_B,
|
||||
shared_storage);
|
||||
}
|
||||
cutlass::arch::cp_async_fence();
|
||||
|
||||
// Ensure first stage committed
|
||||
cutlass::arch::cp_async_wait<kBufferCount - 2>();
|
||||
__syncthreads();
|
||||
|
||||
// Register double buffering for A/B fragments and SFA/SFB like SM80
|
||||
FragmentA fragA_reg[2];
|
||||
FragmentB fragB_reg[2];
|
||||
FragmentSFA fragSFA_reg[2];
|
||||
FragmentSFB fragSFB_reg[2];
|
||||
|
||||
// Current pipe index in smem to read from
|
||||
int smem_pipe_read = 0;
|
||||
// Current pipe index in smem to write to
|
||||
int smem_pipe_write = kBufferCount - 1;
|
||||
|
||||
// PREFETCH register pipeline - load first kblock (stage 0) into register bank 0
|
||||
if constexpr (kStageCount > 1)
|
||||
{
|
||||
int frag_idx = 0;
|
||||
|
||||
// Load fragments using the helper function
|
||||
load_smem_fragments(
|
||||
fragA_reg[frag_idx],
|
||||
fragB_reg[frag_idx],
|
||||
fragSFA_reg[frag_idx],
|
||||
fragSFB_reg[frag_idx],
|
||||
smem_pipe_read,
|
||||
0, // k_block = 0
|
||||
smem_offset_A,
|
||||
smem_offset_B,
|
||||
smem_sf_offset,
|
||||
shared_storage);
|
||||
|
||||
}
|
||||
|
||||
// Mainloop
|
||||
int tile_idx = 0;
|
||||
while (tile_idx < total_tiles) {
|
||||
int smem_pipe_read_curr = smem_pipe_read;
|
||||
|
||||
for_each(make_int_sequence<kStageCount>{}, [&] (auto k_block)
|
||||
{
|
||||
if (k_block == kStageCount - 1)
|
||||
{
|
||||
cutlass::arch::cp_async_wait<kBufferCount - 2>();
|
||||
__syncthreads();
|
||||
|
||||
smem_pipe_read_curr = smem_pipe_read;
|
||||
}
|
||||
|
||||
// Load A/B/SFA/SFB smem->regs for k_block_next
|
||||
auto k_block_next = (k_block + Int<1>{}) % kStageCount;
|
||||
int frag_idx_next = (k_block + 1) & 1;
|
||||
|
||||
// Prefetch next kblock data using saved pipe index
|
||||
load_smem_fragments(
|
||||
fragA_reg[frag_idx_next],
|
||||
fragB_reg[frag_idx_next],
|
||||
fragSFA_reg[frag_idx_next],
|
||||
fragSFB_reg[frag_idx_next],
|
||||
smem_pipe_read_curr,
|
||||
k_block_next,
|
||||
smem_offset_A,
|
||||
smem_offset_B,
|
||||
smem_sf_offset,
|
||||
shared_storage);
|
||||
// Copy gmem to smem before computing gemm on each k-pipe
|
||||
if (k_block == 0)
|
||||
{
|
||||
// Use predicate instead of branch for cp_async
|
||||
bool valid_tile = (global_k < gemm_k);
|
||||
|
||||
// Load all stages using the helper function
|
||||
load_stages_gmem_to_smem(
|
||||
smem_pipe_write, // buffer_idx
|
||||
kStageCount, // num_stages
|
||||
unroll_col_k, // passed by reference
|
||||
global_k, // passed by reference
|
||||
tileA_k_local,
|
||||
smem_offset_A,
|
||||
smem_offset_B,
|
||||
smem_sf_write_offset,
|
||||
is_even_thread,
|
||||
load_b,
|
||||
valid_tile,
|
||||
ptr_A,
|
||||
ptr_B,
|
||||
ptr_SF_A,
|
||||
ptr_SF_B,
|
||||
shared_storage);
|
||||
|
||||
cutlass::arch::cp_async_fence();
|
||||
|
||||
// Advance the pipe indices
|
||||
smem_pipe_write = smem_pipe_read;
|
||||
++smem_pipe_read;
|
||||
smem_pipe_read = (smem_pipe_read == kBufferCount) ? 0 : smem_pipe_read;
|
||||
}
|
||||
|
||||
{
|
||||
int frag_idx = k_block & 1;
|
||||
|
||||
// Compute using current fragments
|
||||
accum += blockscaled_multiply_add(
|
||||
fragA_reg[frag_idx], fragB_reg[frag_idx],
|
||||
fragSFA_reg[frag_idx],
|
||||
fragSFB_reg[frag_idx]);
|
||||
}
|
||||
});
|
||||
|
||||
tile_idx += kStageCount;
|
||||
}
|
||||
|
||||
// Drain outstanding async copies
|
||||
cutlass::arch::cp_async_wait<0>();
|
||||
__syncthreads();
|
||||
|
||||
// Tail elements that don't fill a full tile
|
||||
if (unroll_col_k + idx_col_k * kPackedElementsA < gemm_k) {
|
||||
accum += process_tail_elements(unroll_col_k, idx_col_k, gemm_k,
|
||||
ptr_A, ptr_B,
|
||||
ptr_SF_A, ptr_SF_B,
|
||||
A_converter, B_converter,
|
||||
SFA_converter, SFB_converter);
|
||||
}
|
||||
}
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int mask = (kThreadsPerRow >> 1); mask > 0; mask >>= 1) {
|
||||
accum += ElementAccumulator(__shfl_xor_sync(0xFFFFFFFF, static_cast<float>(accum), mask, 32));
|
||||
}
|
||||
|
||||
auto frag_acc = static_cast<typename EpilogueOutputOp::ElementAccumulator>(accum);
|
||||
auto frag_c = static_cast<typename EpilogueOutputOp::ElementC>(*(ptr_C));
|
||||
|
||||
// Applying blockscaled epilogue
|
||||
epilogue(frag_acc, frag_c, batch_idx);
|
||||
}
|
||||
}
|
||||
} //end of operator()
|
||||
|
||||
private:
|
||||
// Load multiple stages from global to shared memory
|
||||
CUTLASS_DEVICE
|
||||
void load_stages_gmem_to_smem(
|
||||
int buffer_idx,
|
||||
int num_stages,
|
||||
int& unroll_col_k,
|
||||
int& global_k,
|
||||
int tileA_k_local,
|
||||
int smem_offset_A,
|
||||
int smem_offset_B,
|
||||
int smem_sf_write_offset,
|
||||
bool is_even_thread,
|
||||
bool load_b,
|
||||
bool valid_tile,
|
||||
ElementA const* ptr_A,
|
||||
ElementB const* ptr_B,
|
||||
ElementSFA const* ptr_SF_A,
|
||||
ElementSFB const* ptr_SF_B,
|
||||
SharedStorage& shared_storage) {
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int s = 0; s < num_stages; ++s) {
|
||||
// Load scaling factors using cp.async - only even threads participate
|
||||
// Calculate SF indices for this thread
|
||||
int SF_idx = global_k / kSFVecSize;
|
||||
int SF_offset_by_k = ((SF_idx >> 2) << 9) + (SF_idx & 0x3);
|
||||
|
||||
void *smem_ptr_SFA = &shared_storage.smem_SFA[buffer_idx][s][smem_sf_write_offset];
|
||||
const void *gmem_ptr_SFA = ptr_SF_A + SF_offset_by_k;
|
||||
// Load 4 FP8 values (32 bits) - for this thread and next thread
|
||||
cutlass::arch::cp_async<sizeof(uint32_t)>(smem_ptr_SFA, gmem_ptr_SFA, valid_tile && is_even_thread);
|
||||
|
||||
void *smem_ptr_SFB = &shared_storage.smem_SFB[buffer_idx][s][(threadIdx.x / 2) * 4];
|
||||
const void *gmem_ptr_SFB = ptr_SF_B + SF_offset_by_k;
|
||||
// Load 4 FP8 values (32 bits) - for this thread and next thread, only if threadIdx.y == 0
|
||||
cutlass::arch::cp_async<sizeof(uint32_t)>(smem_ptr_SFB, gmem_ptr_SFB, valid_tile && load_b && is_even_thread);
|
||||
|
||||
void *smem_ptr_A = &shared_storage.smem_A[buffer_idx][s][smem_offset_A];
|
||||
const void *gmem_ptr_A = ptr_A + unroll_col_k / kPackedElementsA;
|
||||
cutlass::arch::cp_async<sizeof(FragmentA)>(smem_ptr_A, gmem_ptr_A, valid_tile);
|
||||
|
||||
void *smem_ptr_B = &shared_storage.smem_B[buffer_idx][s][smem_offset_B];
|
||||
const void *gmem_ptr_B = ptr_B + unroll_col_k / kPackedElementsB;
|
||||
cutlass::arch::cp_async<sizeof(FragmentB)>(smem_ptr_B, gmem_ptr_B, valid_tile && load_b);
|
||||
|
||||
unroll_col_k += tileA_k_local;
|
||||
global_k += tileA_k_local;
|
||||
}
|
||||
}
|
||||
|
||||
/// Fused blockscaled GEMV computation using PTX
|
||||
CUTLASS_DEVICE
|
||||
ElementAccumulator blockscaled_multiply_add(
|
||||
FragmentA const& fragA,
|
||||
FragmentB const& fragB,
|
||||
FragmentSFA const& fragSFA,
|
||||
FragmentSFB const& fragSFB) {
|
||||
|
||||
#if defined(CUDA_PTX_FP4FP6_CVT_ENABLED)
|
||||
uint16_t const& src_fragSFA_packed = reinterpret_cast<uint16_t const&>(fragSFA);
|
||||
uint16_t const& src_fragSFB_packed = reinterpret_cast<uint16_t const&>(fragSFB);
|
||||
|
||||
uint32_t const* src_fragA_packed = reinterpret_cast<uint32_t const*>(&fragA);
|
||||
uint32_t const* src_fragB_packed = reinterpret_cast<uint32_t const*>(&fragB);
|
||||
|
||||
ElementAccumulator out;
|
||||
uint16_t* out_fp16 = reinterpret_cast<uint16_t*>(&out);
|
||||
|
||||
asm volatile( \
|
||||
"{\n" \
|
||||
// declare registers for A / B tensors
|
||||
".reg .b8 byte0_0, byte0_1, byte0_2, byte0_3;\n" \
|
||||
".reg .b8 byte0_4, byte0_5, byte0_6, byte0_7;\n" \
|
||||
".reg .b8 byte1_0, byte1_1, byte1_2, byte1_3;\n" \
|
||||
".reg .b8 byte1_4, byte1_5, byte1_6, byte1_7;\n" \
|
||||
".reg .b8 byte2_0, byte2_1, byte2_2, byte2_3;\n" \
|
||||
".reg .b8 byte2_4, byte2_5, byte2_6, byte2_7;\n" \
|
||||
".reg .b8 byte3_0, byte3_1, byte3_2, byte3_3;\n" \
|
||||
".reg .b8 byte3_4, byte3_5, byte3_6, byte3_7;\n" \
|
||||
|
||||
// declare registers for accumulators
|
||||
".reg .f16x2 accum_0_0, accum_0_1, accum_0_2, accum_0_3;\n" \
|
||||
".reg .f16x2 accum_1_0, accum_1_1, accum_1_2, accum_1_3;\n" \
|
||||
".reg .f16x2 accum_2_0, accum_2_1, accum_2_2, accum_2_3;\n" \
|
||||
".reg .f16x2 accum_3_0, accum_3_1, accum_3_2, accum_3_3;\n" \
|
||||
|
||||
// declare registers for scaling factors
|
||||
".reg .f16x2 sfa_f16x2;\n" \
|
||||
".reg .f16x2 sfb_f16x2;\n" \
|
||||
".reg .f16x2 sf_f16x2;\n" \
|
||||
|
||||
// declare registers for conversion
|
||||
".reg .f16x2 cvt_0_0, cvt_0_1, cvt_0_2, cvt_0_3;\n" \
|
||||
".reg .f16x2 cvt_0_4, cvt_0_5, cvt_0_6, cvt_0_7;\n" \
|
||||
".reg .f16x2 cvt_1_0, cvt_1_1, cvt_1_2, cvt_1_3;\n" \
|
||||
".reg .f16x2 cvt_1_4, cvt_1_5, cvt_1_6, cvt_1_7;\n" \
|
||||
".reg .f16x2 cvt_2_0, cvt_2_1, cvt_2_2, cvt_2_3;\n" \
|
||||
".reg .f16x2 cvt_2_4, cvt_2_5, cvt_2_6, cvt_2_7;\n" \
|
||||
".reg .f16x2 cvt_3_0, cvt_3_1, cvt_3_2, cvt_3_3;\n" \
|
||||
".reg .f16x2 cvt_3_4, cvt_3_5, cvt_3_6, cvt_3_7;\n" \
|
||||
".reg .f16 result_f16, lane0, lane1;\n" \
|
||||
".reg .f16x2 mul_f16x2_0, mul_f16x2_1;\n" \
|
||||
|
||||
// convert scaling factors from fp8 to f16x2
|
||||
"cvt.rn.f16x2.e4m3x2 sfa_f16x2, %1;\n" \
|
||||
"cvt.rn.f16x2.e4m3x2 sfb_f16x2, %2;\n" \
|
||||
|
||||
// clear accumulators
|
||||
"mov.b32 accum_0_0, 0;\n" \
|
||||
"mov.b32 accum_0_1, 0;\n" \
|
||||
"mov.b32 accum_0_2, 0;\n" \
|
||||
"mov.b32 accum_0_3, 0;\n" \
|
||||
"mov.b32 accum_1_0, 0;\n" \
|
||||
"mov.b32 accum_1_1, 0;\n" \
|
||||
"mov.b32 accum_1_2, 0;\n" \
|
||||
"mov.b32 accum_1_3, 0;\n" \
|
||||
"mov.b32 accum_2_0, 0;\n" \
|
||||
"mov.b32 accum_2_1, 0;\n" \
|
||||
"mov.b32 accum_2_2, 0;\n" \
|
||||
"mov.b32 accum_2_3, 0;\n" \
|
||||
"mov.b32 accum_3_0, 0;\n" \
|
||||
"mov.b32 accum_3_1, 0;\n" \
|
||||
"mov.b32 accum_3_2, 0;\n" \
|
||||
"mov.b32 accum_3_3, 0;\n" \
|
||||
|
||||
// multiply, unpacking and permuting scale factors
|
||||
"mul.rn.f16x2 sf_f16x2, sfa_f16x2, sfb_f16x2;\n" \
|
||||
"mov.b32 {lane0, lane1}, sf_f16x2;\n" \
|
||||
"mov.b32 mul_f16x2_0, {lane0, lane0};\n" \
|
||||
"mov.b32 mul_f16x2_1, {lane1, lane1};\n" \
|
||||
|
||||
// unpacking A and B tensors
|
||||
"mov.b32 {byte0_0, byte0_1, byte0_2, byte0_3}, %3;\n" \
|
||||
"mov.b32 {byte0_4, byte0_5, byte0_6, byte0_7}, %4;\n" \
|
||||
"mov.b32 {byte1_0, byte1_1, byte1_2, byte1_3}, %5;\n" \
|
||||
"mov.b32 {byte1_4, byte1_5, byte1_6, byte1_7}, %6;\n" \
|
||||
"mov.b32 {byte2_0, byte2_1, byte2_2, byte2_3}, %7;\n" \
|
||||
"mov.b32 {byte2_4, byte2_5, byte2_6, byte2_7}, %8;\n" \
|
||||
"mov.b32 {byte3_0, byte3_1, byte3_2, byte3_3}, %9;\n" \
|
||||
"mov.b32 {byte3_4, byte3_5, byte3_6, byte3_7}, %10;\n" \
|
||||
|
||||
// convert A and B tensors from fp4 to f16x2
|
||||
|
||||
// A[0 - 7] and B[0 - 7]
|
||||
"cvt.rn.f16x2.e2m1x2 cvt_0_0, byte0_0;\n" \
|
||||
"cvt.rn.f16x2.e2m1x2 cvt_0_1, byte0_1;\n" \
|
||||
"cvt.rn.f16x2.e2m1x2 cvt_0_2, byte0_2;\n" \
|
||||
"cvt.rn.f16x2.e2m1x2 cvt_0_3, byte0_3;\n" \
|
||||
"cvt.rn.f16x2.e2m1x2 cvt_0_4, byte0_4;\n" \
|
||||
"cvt.rn.f16x2.e2m1x2 cvt_0_5, byte0_5;\n" \
|
||||
"cvt.rn.f16x2.e2m1x2 cvt_0_6, byte0_6;\n" \
|
||||
"cvt.rn.f16x2.e2m1x2 cvt_0_7, byte0_7;\n" \
|
||||
|
||||
// A[8 - 15] and B[8 - 15]
|
||||
"cvt.rn.f16x2.e2m1x2 cvt_1_0, byte1_0;\n" \
|
||||
"cvt.rn.f16x2.e2m1x2 cvt_1_1, byte1_1;\n" \
|
||||
"cvt.rn.f16x2.e2m1x2 cvt_1_2, byte1_2;\n" \
|
||||
"cvt.rn.f16x2.e2m1x2 cvt_1_3, byte1_3;\n" \
|
||||
"cvt.rn.f16x2.e2m1x2 cvt_1_4, byte1_4;\n" \
|
||||
"cvt.rn.f16x2.e2m1x2 cvt_1_5, byte1_5;\n" \
|
||||
"cvt.rn.f16x2.e2m1x2 cvt_1_6, byte1_6;\n" \
|
||||
"cvt.rn.f16x2.e2m1x2 cvt_1_7, byte1_7;\n" \
|
||||
|
||||
// A[16 - 23] and B[16 - 23]
|
||||
"cvt.rn.f16x2.e2m1x2 cvt_2_0, byte2_0;\n" \
|
||||
"cvt.rn.f16x2.e2m1x2 cvt_2_1, byte2_1;\n" \
|
||||
"cvt.rn.f16x2.e2m1x2 cvt_2_2, byte2_2;\n" \
|
||||
"cvt.rn.f16x2.e2m1x2 cvt_2_3, byte2_3;\n" \
|
||||
"cvt.rn.f16x2.e2m1x2 cvt_2_4, byte2_4;\n" \
|
||||
"cvt.rn.f16x2.e2m1x2 cvt_2_5, byte2_5;\n" \
|
||||
"cvt.rn.f16x2.e2m1x2 cvt_2_6, byte2_6;\n" \
|
||||
"cvt.rn.f16x2.e2m1x2 cvt_2_7, byte2_7;\n" \
|
||||
|
||||
// A[24 - 31] and B[24 - 31]
|
||||
"cvt.rn.f16x2.e2m1x2 cvt_3_0, byte3_0;\n" \
|
||||
"cvt.rn.f16x2.e2m1x2 cvt_3_1, byte3_1;\n" \
|
||||
"cvt.rn.f16x2.e2m1x2 cvt_3_2, byte3_2;\n" \
|
||||
"cvt.rn.f16x2.e2m1x2 cvt_3_3, byte3_3;\n" \
|
||||
"cvt.rn.f16x2.e2m1x2 cvt_3_4, byte3_4;\n" \
|
||||
"cvt.rn.f16x2.e2m1x2 cvt_3_5, byte3_5;\n" \
|
||||
"cvt.rn.f16x2.e2m1x2 cvt_3_6, byte3_6;\n" \
|
||||
"cvt.rn.f16x2.e2m1x2 cvt_3_7, byte3_7;\n" \
|
||||
|
||||
// fma for A[0 - 7] and B[0 - 7]
|
||||
"fma.rn.f16x2 accum_0_0, cvt_0_0, cvt_0_4, accum_0_0;\n" \
|
||||
"fma.rn.f16x2 accum_0_1, cvt_0_1, cvt_0_5, accum_0_1;\n" \
|
||||
"fma.rn.f16x2 accum_0_2, cvt_0_2, cvt_0_6, accum_0_2;\n" \
|
||||
"fma.rn.f16x2 accum_0_3, cvt_0_3, cvt_0_7, accum_0_3;\n" \
|
||||
|
||||
// fma for A[8 - 15] and B[8 - 15]
|
||||
"fma.rn.f16x2 accum_1_0, cvt_1_0, cvt_1_4, accum_1_0;\n" \
|
||||
"fma.rn.f16x2 accum_1_1, cvt_1_1, cvt_1_5, accum_1_1;\n" \
|
||||
"fma.rn.f16x2 accum_1_2, cvt_1_2, cvt_1_6, accum_1_2;\n" \
|
||||
"fma.rn.f16x2 accum_1_3, cvt_1_3, cvt_1_7, accum_1_3;\n" \
|
||||
|
||||
// fma for A[16 - 23] and B[16 - 23]
|
||||
"fma.rn.f16x2 accum_2_0, cvt_2_0, cvt_2_4, accum_2_0;\n" \
|
||||
"fma.rn.f16x2 accum_2_1, cvt_2_1, cvt_2_5, accum_2_1;\n" \
|
||||
"fma.rn.f16x2 accum_2_2, cvt_2_2, cvt_2_6, accum_2_2;\n" \
|
||||
"fma.rn.f16x2 accum_2_3, cvt_2_3, cvt_2_7, accum_2_3;\n" \
|
||||
|
||||
// fma for A[24 - 31] and B[24 - 31]
|
||||
"fma.rn.f16x2 accum_3_0, cvt_3_0, cvt_3_4, accum_3_0;\n" \
|
||||
"fma.rn.f16x2 accum_3_1, cvt_3_1, cvt_3_5, accum_3_1;\n" \
|
||||
"fma.rn.f16x2 accum_3_2, cvt_3_2, cvt_3_6, accum_3_2;\n" \
|
||||
"fma.rn.f16x2 accum_3_3, cvt_3_3, cvt_3_7, accum_3_3;\n" \
|
||||
|
||||
// tree reduction for accumulators
|
||||
"add.rn.f16x2 accum_0_0, accum_0_0, accum_0_1;\n" \
|
||||
"add.rn.f16x2 accum_0_2, accum_0_2, accum_0_3;\n" \
|
||||
"add.rn.f16x2 accum_1_0, accum_1_0, accum_1_1;\n" \
|
||||
"add.rn.f16x2 accum_1_2, accum_1_2, accum_1_3;\n" \
|
||||
"add.rn.f16x2 accum_2_0, accum_2_0, accum_2_1;\n" \
|
||||
"add.rn.f16x2 accum_2_2, accum_2_2, accum_2_3;\n" \
|
||||
"add.rn.f16x2 accum_3_0, accum_3_0, accum_3_1;\n" \
|
||||
"add.rn.f16x2 accum_3_2, accum_3_2, accum_3_3;\n" \
|
||||
|
||||
"add.rn.f16x2 accum_0_0, accum_0_0, accum_0_2;\n" \
|
||||
"add.rn.f16x2 accum_1_0, accum_1_0, accum_1_2;\n" \
|
||||
"add.rn.f16x2 accum_2_0, accum_2_0, accum_2_2;\n" \
|
||||
"add.rn.f16x2 accum_3_0, accum_3_0, accum_3_2;\n" \
|
||||
|
||||
"add.rn.f16x2 accum_0_0, accum_0_0, accum_1_0;\n" \
|
||||
"add.rn.f16x2 accum_2_0, accum_2_0, accum_3_0;\n" \
|
||||
|
||||
// apply scaling factors and final reduction
|
||||
"mul.rn.f16x2 accum_0_0, mul_f16x2_0, accum_0_0;\n" \
|
||||
"mul.rn.f16x2 accum_2_0, mul_f16x2_1, accum_2_0;\n" \
|
||||
|
||||
"add.rn.f16x2 accum_0_0, accum_0_0, accum_2_0;\n" \
|
||||
|
||||
"mov.b32 {lane0, lane1}, accum_0_0;\n" \
|
||||
"add.rn.f16 result_f16, lane0, lane1;\n" \
|
||||
|
||||
"mov.b16 %0, result_f16;\n" \
|
||||
|
||||
"}\n"
|
||||
: "=h"(out_fp16[0]) // 0
|
||||
: "h"(src_fragSFA_packed), "h"(src_fragSFB_packed), // 1, 2
|
||||
"r"(src_fragA_packed[0]), "r"(src_fragB_packed[0]), // 3, 4
|
||||
"r"(src_fragA_packed[1]), "r"(src_fragB_packed[1]), // 5, 6
|
||||
"r"(src_fragA_packed[2]), "r"(src_fragB_packed[2]), // 7, 8
|
||||
"r"(src_fragA_packed[3]), "r"(src_fragB_packed[3]) // 9, 10
|
||||
: "memory"
|
||||
);
|
||||
|
||||
return out;
|
||||
|
||||
#else
|
||||
NumericArrayConverter<ElementAccumulator, ElementA, kElementsPerAccess, Round> srcA_converter;
|
||||
NumericArrayConverter<ElementAccumulator, ElementB, kElementsPerAccess, Round> srcB_converter;
|
||||
NumericConverter<ElementAccumulator, ElementSFA, Round> SFA_converter;
|
||||
NumericConverter<ElementAccumulator, ElementSFB, Round> SFB_converter;
|
||||
|
||||
FragmentCompute fragA_Compute = srcA_converter(fragA);
|
||||
FragmentCompute fragB_Compute = srcB_converter(fragB);
|
||||
ElementAccumulator accum = ElementAccumulator(0);
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int i = 0; i < kSFPerAccess; i++) {
|
||||
ElementAccumulator accum_SF_block = ElementAccumulator(0);
|
||||
|
||||
int local_k_offset = i * kSFVecSize;
|
||||
ElementAccumulator multiplier{1};
|
||||
|
||||
multiplier = SFA_converter(fragSFA.at(i)) * SFB_converter(fragSFB.at(i));
|
||||
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int e = 0; e < kSFVecSize; e++) {
|
||||
accum_SF_block += fragA_Compute.at(e + local_k_offset) * fragB_Compute.at(e + local_k_offset);
|
||||
}
|
||||
|
||||
accum_SF_block *= multiplier;
|
||||
accum += accum_SF_block;
|
||||
}
|
||||
|
||||
return accum;
|
||||
|
||||
#endif
|
||||
}
|
||||
|
||||
CUTLASS_DEVICE
|
||||
ElementAccumulator process_tail_elements(
|
||||
int unroll_col_k,
|
||||
int idx_col_k,
|
||||
int gemm_k,
|
||||
ElementA const *ptr_A,
|
||||
ElementB const *ptr_B,
|
||||
ElementSFA const *ptr_SF_A,
|
||||
ElementSFB const *ptr_SF_B,
|
||||
NumericConverter<ElementAccumulator, ElementA, Round> const &A_converter,
|
||||
NumericConverter<ElementAccumulator, ElementB, Round> const &B_converter,
|
||||
NumericConverter<ElementAccumulator, ElementSFA, Round> const &SFA_converter,
|
||||
NumericConverter<ElementAccumulator, ElementSFB, Round> const &SFB_converter) {
|
||||
|
||||
ElementAccumulator accum = ElementAccumulator(0);
|
||||
|
||||
// calculate the rest of K elements
|
||||
// each thread fetch 1 element each time
|
||||
for (int k = unroll_col_k + idx_col_k * kPackedElementsA; k < gemm_k; k += kThreadsPerRow * kPackedElementsA) {
|
||||
// blockscaled GEMV
|
||||
int SF_idx = k / kSFVecSize;
|
||||
int SF_offset_by_k = ((SF_idx >> 2) << 9) + (SF_idx & 0x3);
|
||||
|
||||
ElementSFA sfa = *(ptr_SF_A + SF_offset_by_k);
|
||||
ElementSFB sfb = *(ptr_SF_B + SF_offset_by_k);
|
||||
|
||||
FragmentPackedA fragA;
|
||||
FragmentPackedB fragB;
|
||||
|
||||
// fetch from matrix A
|
||||
arch::global_load<FragmentPackedA, sizeof(FragmentPackedA), arch::CacheOperation::Always>(
|
||||
fragA,
|
||||
ptr_A - (idx_col_k * kElementsPerAccess - k) / kPackedElementsA,
|
||||
true);
|
||||
|
||||
// fetch from vector B
|
||||
arch::global_load<FragmentPackedB, sizeof(FragmentPackedB), arch::CacheOperation::Always>(
|
||||
fragB,
|
||||
ptr_B - (idx_col_k * kElementsPerAccess - k) / kPackedElementsB,
|
||||
true);
|
||||
|
||||
ElementAccumulator accum_SF_packed = ElementAccumulator(0);
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int e = 0; e < kPackedElements; e++) {
|
||||
accum_SF_packed += A_converter(fragA.at(e)) * B_converter(fragB.at(e));
|
||||
}
|
||||
|
||||
accum_SF_packed *= SFA_converter(sfa) * SFB_converter(sfb);
|
||||
|
||||
accum += accum_SF_packed;
|
||||
|
||||
}
|
||||
|
||||
return accum;
|
||||
}
|
||||
|
||||
// Load fragments from shared memory
|
||||
template<typename FragmentA, typename FragmentB, typename FragmentSFA, typename FragmentSFB>
|
||||
CUTLASS_DEVICE
|
||||
void load_smem_fragments(
|
||||
FragmentA& fragA,
|
||||
FragmentB& fragB,
|
||||
FragmentSFA& fragSFA,
|
||||
FragmentSFB& fragSFB,
|
||||
int smem_pipe_idx,
|
||||
int k_block,
|
||||
int smem_offset_A,
|
||||
int smem_offset_B,
|
||||
int smem_sf_offset,
|
||||
SharedStorage& shared_storage) const {
|
||||
|
||||
// Load A/B fragments
|
||||
arch::shared_load(fragA, &shared_storage.smem_A[smem_pipe_idx][k_block][smem_offset_A]);
|
||||
arch::shared_load(fragB, &shared_storage.smem_B[smem_pipe_idx][k_block][smem_offset_B]);
|
||||
|
||||
// Load SF fragments
|
||||
uint32_t smem_ptr = cutlass::arch::cutlass_get_smem_pointer(&shared_storage.smem_SFA[smem_pipe_idx][k_block][smem_sf_offset]);
|
||||
arch::shared_load<2>(&fragSFA, smem_ptr);
|
||||
smem_ptr = cutlass::arch::cutlass_get_smem_pointer(&shared_storage.smem_SFB[smem_pipe_idx][k_block][threadIdx.x * kSFPerAccess]);
|
||||
arch::shared_load<2>(&fragSFB, smem_ptr);
|
||||
|
||||
}
|
||||
|
||||
};
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace kernel
|
||||
} // namespace gemm
|
||||
} // namespace cutlass
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
@@ -0,0 +1,793 @@
|
||||
/***************************************************************************************************
|
||||
* Copyright (c) 2024 - 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/cutlass.h"
|
||||
#include "cutlass/workspace.h"
|
||||
#include "cutlass/kernel_hardware_info.hpp"
|
||||
#include "cutlass/detail/cluster.hpp"
|
||||
#include "cutlass/fast_math.h"
|
||||
#include "cute/arch/cluster_sm90.hpp"
|
||||
#include "cutlass/arch/arch.h"
|
||||
#include "cutlass/arch/barrier.h"
|
||||
#include "cutlass/arch/reg_reconfig.h"
|
||||
#include "cutlass/gemm/gemm.h"
|
||||
#include "cutlass/gemm/dispatch_policy.hpp"
|
||||
#include "cutlass/gemm/kernel/sm100_tile_scheduler.hpp"
|
||||
#include "cutlass/pipeline/pipeline.hpp"
|
||||
|
||||
#include "cute/tensor.hpp"
|
||||
#include "cute/arch/tmem_allocator_sm100.hpp"
|
||||
#include "cute/atom/mma_atom.hpp"
|
||||
|
||||
///////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace cutlass::gemm::kernel {
|
||||
|
||||
///////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <
|
||||
class ProblemShape_,
|
||||
class CollectiveMainloop_,
|
||||
class CollectiveEpilogue_,
|
||||
class TileSchedulerTag_
|
||||
>
|
||||
class GemmUniversal<
|
||||
ProblemShape_,
|
||||
CollectiveMainloop_,
|
||||
CollectiveEpilogue_,
|
||||
TileSchedulerTag_,
|
||||
cute::enable_if_t<
|
||||
cutlass::detail::is_kernel_tag_of_v<typename CollectiveMainloop_::DispatchPolicy::Schedule,
|
||||
KernelWarpSpecializedSm100>>>
|
||||
{
|
||||
public:
|
||||
using ProblemShape = ProblemShape_;
|
||||
static_assert(rank(ProblemShape{}) == 3 or rank(ProblemShape{}) == 4,
|
||||
"ProblemShape{} should be <M,N,K> or <M,N,K,L>");
|
||||
static constexpr bool IsGdcEnabled = false;
|
||||
// Mainloop derived types
|
||||
using CollectiveMainloop = CollectiveMainloop_;
|
||||
using TileShape = typename CollectiveMainloop::TileShape;
|
||||
using TiledMma = typename CollectiveMainloop::TiledMma;
|
||||
using ArchTag = typename CollectiveMainloop::ArchTag;
|
||||
using ElementA = typename CollectiveMainloop::ElementA;
|
||||
using StrideA = typename CollectiveMainloop::StrideA;
|
||||
using ElementB = typename CollectiveMainloop::ElementB;
|
||||
using StrideB = typename CollectiveMainloop::StrideB;
|
||||
using DispatchPolicy = typename CollectiveMainloop::DispatchPolicy;
|
||||
using ElementAccumulator = typename CollectiveMainloop::ElementAccumulator;
|
||||
using ClusterShape = typename DispatchPolicy::ClusterShape;
|
||||
using MainloopArguments = typename CollectiveMainloop::Arguments;
|
||||
using MainloopParams = typename CollectiveMainloop::Params;
|
||||
static_assert(ArchTag::kMinComputeCapability >= 100);
|
||||
|
||||
// Epilogue derived types
|
||||
using CollectiveEpilogue = CollectiveEpilogue_;
|
||||
using ElementC = typename CollectiveEpilogue::ElementC;
|
||||
using StrideC = typename CollectiveEpilogue::StrideC;
|
||||
using ElementD = typename CollectiveEpilogue::ElementD;
|
||||
using StrideD = typename CollectiveEpilogue::StrideD;
|
||||
using EpilogueArguments = typename CollectiveEpilogue::Arguments;
|
||||
using EpilogueParams = typename CollectiveEpilogue::Params;
|
||||
static constexpr bool IsComplex = CollectiveEpilogue::NumAccumulatorMtxs == 2;
|
||||
|
||||
// CLC pipeline depth
|
||||
// determines how many waves (stages-1) a warp can race ahead
|
||||
static constexpr uint32_t SchedulerPipelineStageCount = DispatchPolicy::Schedule::SchedulerPipelineStageCount;
|
||||
|
||||
// TileID scheduler
|
||||
// Get Blk and Scheduling tile shapes
|
||||
using CtaShape_MNK = typename CollectiveMainloop::CtaShape_MNK;
|
||||
using AtomThrShapeMNK = typename CollectiveMainloop::AtomThrShapeMNK;
|
||||
|
||||
static_assert(size(AtomThrShapeMNK{}) == 1, "Lower alignment kernel only supports 1x1x1 cluster shape.");
|
||||
using TileSchedulerTag = TileSchedulerTag_;
|
||||
using TileScheduler = typename detail::TileSchedulerSelector<
|
||||
TileSchedulerTag, ArchTag, CtaShape_MNK, ClusterShape, SchedulerPipelineStageCount>::Scheduler;
|
||||
using TileSchedulerArguments = typename TileScheduler::Arguments;
|
||||
using TileSchedulerParams = typename TileScheduler::Params;
|
||||
|
||||
// Warp specialization thread count per threadblock
|
||||
static constexpr uint32_t NumSchedThreads = NumThreadsPerWarp; // 1 warp
|
||||
static constexpr uint32_t NumMMAThreads = NumThreadsPerWarp; // 1 warp
|
||||
static constexpr uint32_t NumEmptyThreads = NumThreadsPerWarp; // 1 warp
|
||||
static constexpr uint32_t NumMainloopLoadThreads = CollectiveMainloop::NumLoadThreads; // 4 warps
|
||||
static constexpr uint32_t NumEpilogueLoadThreads = NumThreadsPerWarp; // 1 warp
|
||||
static constexpr uint32_t NumEpilogueThreads = CollectiveEpilogue::ThreadCount;
|
||||
static constexpr uint32_t NumEpilogueWarps = NumEpilogueThreads / NumThreadsPerWarp;
|
||||
|
||||
static constexpr uint32_t MaxThreadsPerBlock = NumSchedThreads +
|
||||
NumMainloopLoadThreads + NumMMAThreads +
|
||||
NumEpilogueLoadThreads + NumEpilogueThreads + NumEmptyThreads;
|
||||
static constexpr uint32_t MinBlocksPerMultiprocessor = 1;
|
||||
static constexpr uint32_t NumFixupBarriers = 1;
|
||||
static constexpr uint32_t CLCResponseSize = sizeof(typename TileScheduler::CLCResponse);
|
||||
|
||||
static constexpr bool IsSchedDynamicPersistent = TileScheduler::IsDynamicPersistent;
|
||||
|
||||
// Pipelines and pipeline states
|
||||
static constexpr uint32_t AccumulatorPipelineStageCount = DispatchPolicy::Schedule::AccumulatorPipelineStageCount;
|
||||
|
||||
// Pipeline and pipeline state types
|
||||
using MainloopPipeline = typename CollectiveMainloop::MainloopPipeline;
|
||||
using MainloopPipelineState = typename CollectiveMainloop::MainloopPipelineState;
|
||||
|
||||
using EpiLoadPipeline = typename CollectiveEpilogue::LoadPipeline;
|
||||
using EpiLoadPipelineState = typename CollectiveEpilogue::LoadPipelineState;
|
||||
|
||||
using EpiStorePipeline = typename CollectiveEpilogue::StorePipeline;
|
||||
using EpiStorePipelineState = typename CollectiveEpilogue::StorePipelineState;
|
||||
|
||||
using AccumulatorPipeline = cutlass::PipelineUmmaAsync<AccumulatorPipelineStageCount, AtomThrShapeMNK>;
|
||||
using AccumulatorPipelineState = typename AccumulatorPipeline::PipelineState;
|
||||
|
||||
using CLCPipeline = cutlass::PipelineCLCFetchAsync<SchedulerPipelineStageCount, ClusterShape>;
|
||||
using CLCPipelineState = typename CLCPipeline::PipelineState;
|
||||
|
||||
using TmemAllocator = cute::TMEM::Allocator1Sm;
|
||||
|
||||
static constexpr int EpilogueWarpRegs = 248;
|
||||
static constexpr int NonEpilogueWarpRegs = 128;
|
||||
|
||||
// Kernel level shared memory storage
|
||||
struct SharedStorage {
|
||||
struct PipelineStorage : cute::aligned_struct<16, _1> {
|
||||
using MainloopPipelineStorage = typename CollectiveMainloop::PipelineStorage;
|
||||
using EpiLoadPipelineStorage = typename CollectiveEpilogue::PipelineStorage;
|
||||
using CLCPipelineStorage = typename CLCPipeline::SharedStorage;
|
||||
using AccumulatorPipelineStorage = typename AccumulatorPipeline::SharedStorage;
|
||||
|
||||
alignas(16) MainloopPipelineStorage mainloop;
|
||||
alignas(16) EpiLoadPipelineStorage epi_load;
|
||||
alignas(16) CLCPipelineStorage clc;
|
||||
alignas(16) AccumulatorPipelineStorage accumulator;
|
||||
alignas(16) arch::ClusterBarrier tmem_dealloc;
|
||||
} pipelines;
|
||||
|
||||
alignas(16) typename TileScheduler::CLCResponse clc_response[SchedulerPipelineStageCount];
|
||||
uint32_t tmem_base_ptr;
|
||||
|
||||
struct TensorStorage : cute::aligned_struct<128, _1> {
|
||||
using MainloopTensorStorage = typename CollectiveMainloop::TensorStorage;
|
||||
using EpilogueTensorStorage = typename CollectiveEpilogue::TensorStorage;
|
||||
|
||||
MainloopTensorStorage mainloop;
|
||||
EpilogueTensorStorage epilogue;
|
||||
} tensors;
|
||||
|
||||
};
|
||||
|
||||
static constexpr int SharedStorageSize = sizeof(SharedStorage);
|
||||
static_assert(SharedStorageSize <= cutlass::arch::sm100_smem_capacity_bytes, "SMEM usage exceeded capacity.");
|
||||
|
||||
// Host facing host arguments
|
||||
struct Arguments {
|
||||
GemmUniversalMode mode{};
|
||||
ProblemShape problem_shape{};
|
||||
MainloopArguments mainloop{};
|
||||
EpilogueArguments epilogue{};
|
||||
KernelHardwareInfo hw_info{};
|
||||
TileSchedulerArguments scheduler{};
|
||||
};
|
||||
|
||||
// Kernel device entry point API
|
||||
struct Params {
|
||||
GemmUniversalMode mode{};
|
||||
ProblemShape problem_shape{};
|
||||
MainloopParams mainloop{};
|
||||
EpilogueParams epilogue{};
|
||||
KernelHardwareInfo hw_info{};
|
||||
TileSchedulerParams scheduler{};
|
||||
};
|
||||
|
||||
enum class WarpCategory : int32_t {
|
||||
MMA = 0,
|
||||
Sched = 1,
|
||||
EpilogueLoad = 3,
|
||||
Epilogue = 4,
|
||||
MainloopLoad = 8
|
||||
};
|
||||
|
||||
struct IsParticipant {
|
||||
uint32_t mma = false;
|
||||
uint32_t sched = false;
|
||||
uint32_t epi_load = false;
|
||||
uint32_t epilogue = false;
|
||||
uint32_t main_load = false;
|
||||
};
|
||||
|
||||
// Convert to underlying arguments.
|
||||
static
|
||||
Params
|
||||
to_underlying_arguments(Arguments const& args, void* workspace) {
|
||||
(void) workspace;
|
||||
auto problem_shape = args.problem_shape;
|
||||
auto problem_shape_MNKL = append<4>(problem_shape, 1);
|
||||
static constexpr uint32_t NumEpilogueSubTiles = 1;
|
||||
|
||||
// Get SM count if needed, otherwise use user supplied SM count
|
||||
int sm_count = args.hw_info.sm_count;
|
||||
if (sm_count != 0) {
|
||||
CUTLASS_TRACE_HOST(" WARNING: SM100 tile scheduler does not allow for user specified SM counts.\n"
|
||||
" To restrict a kernel's resource usage, consider using CUDA driver APIs instead (green contexts).");
|
||||
sm_count = KernelHardwareInfo::query_device_multiprocessor_count(args.hw_info.device_id);
|
||||
}
|
||||
|
||||
CUTLASS_TRACE_HOST("to_underlying_arguments(): Setting persistent grid SM count to " << sm_count);
|
||||
|
||||
KernelHardwareInfo hw_info{args.hw_info.device_id, sm_count};
|
||||
|
||||
// Calculate workspace pointers
|
||||
uint8_t* workspace_ptr = reinterpret_cast<uint8_t*>(workspace);
|
||||
size_t workspace_offset = 0;
|
||||
|
||||
// Epilogue
|
||||
void* epilogue_workspace = workspace_ptr + workspace_offset;
|
||||
workspace_offset += CollectiveEpilogue::get_workspace_size(args.problem_shape, args.epilogue);
|
||||
workspace_offset = round_nearest(workspace_offset, MinWorkspaceAlignment);
|
||||
|
||||
void* mainloop_workspace = nullptr;
|
||||
|
||||
// Tile scheduler
|
||||
void* scheduler_workspace = workspace_ptr + workspace_offset;
|
||||
workspace_offset += TileScheduler::template get_workspace_size<ProblemShape, ElementAccumulator>(
|
||||
args.scheduler, args.problem_shape, args.hw_info, NumFixupBarriers, NumEpilogueSubTiles, CollectiveEpilogue::NumAccumulatorMtxs);
|
||||
workspace_offset = round_nearest(workspace_offset, MinWorkspaceAlignment);
|
||||
|
||||
return {
|
||||
args.mode,
|
||||
args.problem_shape,
|
||||
CollectiveMainloop::to_underlying_arguments(args.problem_shape, args.mainloop, mainloop_workspace),
|
||||
CollectiveEpilogue::to_underlying_arguments(args.problem_shape, args.epilogue, epilogue_workspace),
|
||||
hw_info,
|
||||
TileScheduler::to_underlying_arguments(
|
||||
problem_shape_MNKL, TileShape{}, AtomThrShapeMNK{}, ClusterShape{},
|
||||
args.hw_info, args.scheduler, scheduler_workspace
|
||||
)
|
||||
};
|
||||
}
|
||||
|
||||
static bool
|
||||
can_implement(Arguments const& args) {
|
||||
bool implementable = (args.mode == GemmUniversalMode::kGemm) or
|
||||
(args.mode == GemmUniversalMode::kBatched && rank(ProblemShape{}) == 4);
|
||||
if (!implementable) {
|
||||
CUTLASS_TRACE_HOST(" CAN IMPLEMENT: Arguments or Problem Shape don't meet the requirements.\n");
|
||||
return implementable;
|
||||
}
|
||||
implementable &= CollectiveMainloop::can_implement(args.problem_shape, args.mainloop);
|
||||
implementable &= CollectiveEpilogue::can_implement(args.problem_shape, args.epilogue);
|
||||
implementable &= TileScheduler::can_implement(args.scheduler);
|
||||
|
||||
static constexpr int MaxClusterSize = 16;
|
||||
implementable &= size(ClusterShape{}) <= MaxClusterSize;
|
||||
|
||||
return implementable;
|
||||
}
|
||||
|
||||
static size_t
|
||||
get_workspace_size(Arguments const& args) {
|
||||
static constexpr uint32_t NumEpilogueSubTiles = 1;
|
||||
size_t workspace_size = 0;
|
||||
|
||||
// Epilogue
|
||||
workspace_size += CollectiveEpilogue::get_workspace_size(args.problem_shape, args.epilogue);
|
||||
workspace_size = round_nearest(workspace_size, MinWorkspaceAlignment);
|
||||
|
||||
// Tile scheduler
|
||||
workspace_size += TileScheduler::template get_workspace_size<ProblemShape, ElementAccumulator>(
|
||||
args.scheduler, args.problem_shape, args.hw_info, NumFixupBarriers, NumEpilogueSubTiles, CollectiveEpilogue::NumAccumulatorMtxs);
|
||||
workspace_size = round_nearest(workspace_size, MinWorkspaceAlignment);
|
||||
|
||||
return workspace_size;
|
||||
}
|
||||
|
||||
static cutlass::Status
|
||||
initialize_workspace(Arguments const& args, void* workspace = nullptr, cudaStream_t stream = nullptr,
|
||||
CudaHostAdapter* cuda_adapter = nullptr) {
|
||||
Status status = Status::kSuccess;
|
||||
uint8_t* workspace_ptr = reinterpret_cast<uint8_t*>(workspace);
|
||||
size_t workspace_offset = 0;
|
||||
static constexpr uint32_t NumEpilogueSubTiles = 1;
|
||||
|
||||
// Epilogue
|
||||
status = CollectiveEpilogue::initialize_workspace(args.problem_shape, args.epilogue, workspace_ptr + workspace_offset, stream, cuda_adapter);
|
||||
workspace_offset += CollectiveEpilogue::get_workspace_size(args.problem_shape, args.epilogue);
|
||||
status = cutlass::Status::kSuccess;
|
||||
workspace_offset = round_nearest(workspace_offset, MinWorkspaceAlignment);
|
||||
if (status != Status::kSuccess) {
|
||||
return status;
|
||||
}
|
||||
|
||||
// Tile scheduler
|
||||
status = TileScheduler::template initialize_workspace<ProblemShape, ElementAccumulator>(
|
||||
args.scheduler, workspace_ptr + workspace_offset, stream, args.problem_shape, args.hw_info, NumFixupBarriers, NumEpilogueSubTiles, CollectiveEpilogue::NumAccumulatorMtxs, cuda_adapter);
|
||||
workspace_offset += TileScheduler::template get_workspace_size<ProblemShape, ElementAccumulator>(
|
||||
args.scheduler, args.problem_shape, args.hw_info, NumFixupBarriers);
|
||||
workspace_offset = round_nearest(workspace_offset, MinWorkspaceAlignment);
|
||||
if (status != Status::kSuccess) {
|
||||
return status;
|
||||
}
|
||||
|
||||
return status;
|
||||
}
|
||||
|
||||
static dim3
|
||||
get_grid_shape(Params const& params) {
|
||||
auto cluster_shape = ClusterShape{};
|
||||
auto blk_shape = CtaShape_MNK{};
|
||||
auto problem_shape_MNKL = append<4>(params.problem_shape, Int<1>{});
|
||||
return TileScheduler::get_grid_shape(
|
||||
params.scheduler,
|
||||
problem_shape_MNKL,
|
||||
TileShape{},
|
||||
AtomThrShapeMNK{},
|
||||
cluster_shape,
|
||||
params.hw_info
|
||||
);
|
||||
|
||||
}
|
||||
|
||||
static dim3
|
||||
get_block_shape() {
|
||||
return dim3(MaxThreadsPerBlock, 1, 1);
|
||||
}
|
||||
|
||||
public:
|
||||
|
||||
CUTLASS_DEVICE
|
||||
void
|
||||
operator()(Params const& params, char* smem_buf) {
|
||||
|
||||
using namespace cute;
|
||||
using X = Underscore;
|
||||
// Separate out problem shape for convenience
|
||||
// Optionally append 1s until problem shape is rank-4 in case its is only rank-3 (MNK)
|
||||
auto problem_shape_MNKL = append<4>(params.problem_shape, Int<1>{});
|
||||
auto M = get<0>(problem_shape_MNKL);
|
||||
auto N = get<1>(problem_shape_MNKL);
|
||||
auto K = get<2>(problem_shape_MNKL);
|
||||
auto L = get<3>(problem_shape_MNKL);
|
||||
|
||||
// Account for more than one epilogue warp
|
||||
int warp_idx = canonical_warp_idx_sync();
|
||||
WarpCategory warp_category = warp_idx < static_cast<int>(WarpCategory::Epilogue) ? WarpCategory(warp_idx)
|
||||
: warp_idx < static_cast<int>(WarpCategory::MainloopLoad) ? WarpCategory::Epilogue
|
||||
: WarpCategory::MainloopLoad;
|
||||
uint32_t lane_predicate = cute::elect_one_sync();
|
||||
auto tile_shape = TileShape{};
|
||||
auto cluster_shape = ClusterShape{};
|
||||
constexpr int cluster_size = size(ClusterShape{});
|
||||
int cta_rank_in_cluster = cute::block_rank_in_cluster();
|
||||
bool is_first_cta_in_cluster = cta_rank_in_cluster == 0;
|
||||
int cta_coord_v = cta_rank_in_cluster % size<0>(typename TiledMma::AtomThrID{});
|
||||
bool is_mma_leader_cta = cta_coord_v == 0;
|
||||
int mma_leader_ctas = size(shape_div(cluster_shape, AtomThrShapeMNK{}));
|
||||
[[maybe_unused]] uint32_t mma_peer_cta_rank = cta_rank_in_cluster;
|
||||
|
||||
// Kernel level shared memory storage
|
||||
SharedStorage& shared_storage = *reinterpret_cast<SharedStorage*>(smem_buf);
|
||||
|
||||
// In a warp specialized kernel, collectives expose data movement and compute operations separately
|
||||
CollectiveMainloop collective_mainloop;
|
||||
CollectiveEpilogue collective_epilogue(params.epilogue, shared_storage.tensors.epilogue);
|
||||
|
||||
// Do we load source tensor C or other aux inputs
|
||||
bool is_epi_load_needed = collective_epilogue.is_producer_load_needed();
|
||||
|
||||
IsParticipant is_participant = {
|
||||
(warp_category == WarpCategory::MMA) && is_mma_leader_cta, // mma
|
||||
(warp_category == WarpCategory::Sched) && is_first_cta_in_cluster, // sched
|
||||
(warp_category == WarpCategory::EpilogueLoad) && is_epi_load_needed, // epi_load
|
||||
(warp_category == WarpCategory::Epilogue), // epilogue
|
||||
(warp_category == WarpCategory::MainloopLoad) // main_load
|
||||
};
|
||||
|
||||
// Mainloop Load pipeline
|
||||
typename MainloopPipeline::Params mainloop_pipeline_params;
|
||||
if (WarpCategory::MainloopLoad == warp_category) {
|
||||
mainloop_pipeline_params.role = MainloopPipeline::ThreadCategory::Producer;
|
||||
}
|
||||
if (WarpCategory::MMA == warp_category) {
|
||||
mainloop_pipeline_params.role = MainloopPipeline::ThreadCategory::Consumer;
|
||||
}
|
||||
|
||||
mainloop_pipeline_params.producer_arv_count = NumMainloopLoadThreads;
|
||||
mainloop_pipeline_params.consumer_arv_count = 1; // Only UMMA consumes the A and B buffers
|
||||
mainloop_pipeline_params.dst_blockid = cta_rank_in_cluster;
|
||||
mainloop_pipeline_params.initializing_warp = 0;
|
||||
MainloopPipeline mainloop_pipeline(shared_storage.pipelines.mainloop, mainloop_pipeline_params, cluster_shape);
|
||||
|
||||
// Epilogue Load pipeline
|
||||
typename EpiLoadPipeline::Params epi_load_pipeline_params;
|
||||
if (WarpCategory::EpilogueLoad == warp_category) {
|
||||
epi_load_pipeline_params.role = EpiLoadPipeline::ThreadCategory::Producer;
|
||||
}
|
||||
if (WarpCategory::Epilogue == warp_category) {
|
||||
epi_load_pipeline_params.role = EpiLoadPipeline::ThreadCategory::Consumer;
|
||||
}
|
||||
epi_load_pipeline_params.dst_blockid = cta_rank_in_cluster;
|
||||
epi_load_pipeline_params.producer_arv_count = NumEpilogueLoadThreads;
|
||||
epi_load_pipeline_params.consumer_arv_count = NumEpilogueThreads;
|
||||
epi_load_pipeline_params.transaction_bytes = CollectiveEpilogue::TmaTransactionBytes;
|
||||
epi_load_pipeline_params.initializing_warp = 3;
|
||||
EpiLoadPipeline epi_load_pipeline(shared_storage.pipelines.epi_load, epi_load_pipeline_params);
|
||||
|
||||
// Epilogue Store pipeline
|
||||
typename EpiStorePipeline::Params epi_store_pipeline_params;
|
||||
epi_store_pipeline_params.always_wait = true;
|
||||
EpiStorePipeline epi_store_pipeline(epi_store_pipeline_params);
|
||||
|
||||
// CLC pipeline
|
||||
typename CLCPipeline::Params clc_pipeline_params;
|
||||
if (WarpCategory::Sched == warp_category) {
|
||||
clc_pipeline_params.role = CLCPipeline::ThreadCategory::ProducerConsumer;
|
||||
}
|
||||
else {
|
||||
clc_pipeline_params.role = CLCPipeline::ThreadCategory::Consumer;
|
||||
}
|
||||
clc_pipeline_params.producer_blockid = 0;
|
||||
clc_pipeline_params.producer_arv_count = 1;
|
||||
clc_pipeline_params.consumer_arv_count = NumSchedThreads + cluster_size *
|
||||
(NumMainloopLoadThreads + NumEpilogueThreads + NumMMAThreads);
|
||||
|
||||
clc_pipeline_params.transaction_bytes = CLCResponseSize;
|
||||
clc_pipeline_params.initializing_warp = 1;
|
||||
CLCPipeline clc_pipeline(shared_storage.pipelines.clc, clc_pipeline_params, cluster_shape);
|
||||
|
||||
// Mainloop-Epilogue pipeline
|
||||
typename AccumulatorPipeline::Params accumulator_pipeline_params;
|
||||
if (WarpCategory::MMA == warp_category) {
|
||||
accumulator_pipeline_params.role = AccumulatorPipeline::ThreadCategory::Producer;
|
||||
}
|
||||
if (WarpCategory::Epilogue == warp_category) {
|
||||
accumulator_pipeline_params.role = AccumulatorPipeline::ThreadCategory::Consumer;
|
||||
}
|
||||
// Only one producer thread arrives on this barrier.
|
||||
accumulator_pipeline_params.producer_arv_count = 1;
|
||||
accumulator_pipeline_params.consumer_arv_count = size(AtomThrShapeMNK{}) * NumEpilogueThreads;
|
||||
accumulator_pipeline_params.initializing_warp = 2;
|
||||
AccumulatorPipeline accumulator_pipeline(shared_storage.pipelines.accumulator, accumulator_pipeline_params, cluster_shape);
|
||||
|
||||
// Tmem allocator
|
||||
TmemAllocator tmem_allocator{};
|
||||
|
||||
// Sync allocation status between MMA and epilogue warps within CTA
|
||||
arch::NamedBarrier tmem_allocation_result_barrier(NumMMAThreads + NumEpilogueThreads, cutlass::arch::ReservedNamedBarriers::TmemAllocBarrier);
|
||||
// Sync deallocation status between MMA warps of peer CTAs
|
||||
arch::ClusterBarrier& tmem_deallocation_result_barrier = shared_storage.pipelines.tmem_dealloc;
|
||||
[[maybe_unused]] uint32_t dealloc_barrier_phase = 0;
|
||||
|
||||
MainloopPipelineState mainloop_pipe_consumer_state;
|
||||
MainloopPipelineState mainloop_pipe_producer_state = cutlass::make_producer_start_state<MainloopPipeline>();
|
||||
|
||||
EpiLoadPipelineState epi_load_pipe_consumer_state;
|
||||
EpiLoadPipelineState epi_load_pipe_producer_state = cutlass::make_producer_start_state<EpiLoadPipeline>();
|
||||
|
||||
// epilogue store pipe is producer-only (consumer is TMA unit, waits via scoreboarding)
|
||||
EpiStorePipelineState epi_store_pipe_producer_state = cutlass::make_producer_start_state<EpiStorePipeline>();
|
||||
|
||||
CLCPipelineState clc_pipe_consumer_state;
|
||||
CLCPipelineState clc_pipe_producer_state = cutlass::make_producer_start_state<CLCPipeline>();
|
||||
|
||||
AccumulatorPipelineState accumulator_pipe_consumer_state;
|
||||
AccumulatorPipelineState accumulator_pipe_producer_state = cutlass::make_producer_start_state<AccumulatorPipeline>();
|
||||
|
||||
// We need this to guarantee that the Pipeline init is visible
|
||||
// To all producers and consumer threadblocks in the cluster
|
||||
pipeline_init_arrive_relaxed(cluster_size);
|
||||
|
||||
dim3 block_id_in_cluster = cute::block_id_in_cluster();
|
||||
// TileID scheduler
|
||||
TileScheduler scheduler(&shared_storage.clc_response[0], params.scheduler, block_id_in_cluster);
|
||||
typename TileScheduler::WorkTileInfo work_tile_info = scheduler.initial_work_tile_info(cluster_shape);
|
||||
auto cta_coord_mnkl = scheduler.work_tile_to_cta_coord(work_tile_info);
|
||||
|
||||
//
|
||||
// TMEM "Allocation"
|
||||
//
|
||||
auto acc_shape = collective_mainloop.partition_accumulator_shape();
|
||||
auto bulk_tmem = TiledMma::make_fragment_C(append(acc_shape,
|
||||
Int<AccumulatorPipelineStageCount>{}));
|
||||
|
||||
//
|
||||
// END PROLOGUE
|
||||
//
|
||||
|
||||
// Synchronization call. Blocks until barriers are initialized in shared memory.
|
||||
pipeline_init_wait(cluster_size);
|
||||
|
||||
if (is_participant.main_load) {
|
||||
cutlass::arch::warpgroup_reg_dealloc<NonEpilogueWarpRegs>();
|
||||
|
||||
auto load_inputs = collective_mainloop.load_init(
|
||||
problem_shape_MNKL, params.mainloop, shared_storage.tensors.mainloop);
|
||||
Tensor gA_mkl = get<0>(load_inputs);
|
||||
|
||||
do {
|
||||
// Get current work tile and fetch next work tile
|
||||
cta_coord_mnkl = scheduler.work_tile_to_cta_coord(work_tile_info);
|
||||
|
||||
auto [next_work_tile_info, increment_pipe] = scheduler.fetch_next_work(
|
||||
work_tile_info,
|
||||
clc_pipeline,
|
||||
clc_pipe_consumer_state
|
||||
);
|
||||
|
||||
// Get the number of K tiles to compute for this work as well as the starting K tile offset of the work.
|
||||
auto k_tile_iter = scheduler.get_k_tile_iterator(work_tile_info, problem_shape_MNKL, CtaShape_MNK{}, shape<3>(gA_mkl));
|
||||
auto k_tile_count = TileScheduler::get_work_k_tile_count(work_tile_info, problem_shape_MNKL, CtaShape_MNK{});
|
||||
|
||||
auto [mainloop_producer_state_next, unused_] = collective_mainloop.load(
|
||||
params.mainloop,
|
||||
mainloop_pipeline,
|
||||
mainloop_pipe_producer_state,
|
||||
load_inputs,
|
||||
cta_coord_mnkl,
|
||||
k_tile_iter, k_tile_count
|
||||
);
|
||||
mainloop_pipe_producer_state = mainloop_producer_state_next;
|
||||
|
||||
// Sync warp to prevent non-participating threads entering next wave early
|
||||
__syncwarp();
|
||||
work_tile_info = next_work_tile_info;
|
||||
|
||||
if (increment_pipe) {
|
||||
++clc_pipe_consumer_state;
|
||||
}
|
||||
} while (work_tile_info.is_valid());
|
||||
|
||||
collective_mainloop.load_tail(mainloop_pipeline, mainloop_pipe_producer_state);
|
||||
|
||||
}
|
||||
|
||||
else if (is_participant.sched) {
|
||||
cutlass::arch::warpgroup_reg_dealloc<NonEpilogueWarpRegs>();
|
||||
|
||||
if constexpr (IsSchedDynamicPersistent) {
|
||||
// Whether a new CLC query must be performed.
|
||||
// See comment below where this variable is updated for a description of
|
||||
// why this variable is needed.
|
||||
bool requires_clc_query = true;
|
||||
|
||||
cutlass::arch::wait_on_dependent_grids();
|
||||
|
||||
do {
|
||||
if (requires_clc_query) {
|
||||
// Query next clcID and update producer state
|
||||
clc_pipe_producer_state = scheduler.advance_to_next_work(clc_pipeline, clc_pipe_producer_state);
|
||||
}
|
||||
|
||||
// Fetch next work tile
|
||||
auto [next_work_tile_info, increment_pipe] = scheduler.fetch_next_work(
|
||||
work_tile_info,
|
||||
clc_pipeline,
|
||||
clc_pipe_consumer_state
|
||||
);
|
||||
|
||||
// Only perform a new CLC query if we consumed a new CLC query result in
|
||||
// `fetch_next_work`. An example of a case in which CLC `fetch_next_work` does
|
||||
// not consume a new CLC query response is when processing stream-K units.
|
||||
// The current stream-K scheduler uses single WorkTileInfo to track multiple
|
||||
// (potentially-partial) tiles to be computed via stream-K. In this case,
|
||||
// `fetch_next_work` simply performs in-place updates on the existing WorkTileInfo,
|
||||
// rather than consuming a CLC query response.
|
||||
requires_clc_query = increment_pipe;
|
||||
if (increment_pipe) {
|
||||
++clc_pipe_consumer_state;
|
||||
}
|
||||
|
||||
work_tile_info = next_work_tile_info;
|
||||
} while (work_tile_info.is_valid());
|
||||
clc_pipeline.producer_tail(clc_pipe_producer_state);
|
||||
}
|
||||
}
|
||||
|
||||
else if (is_participant.mma) {
|
||||
cutlass::arch::warpgroup_reg_dealloc<NonEpilogueWarpRegs>();
|
||||
|
||||
// Tmem allocation sequence
|
||||
tmem_allocator.allocate(TmemAllocator::Sm100TmemCapacityColumns, &shared_storage.tmem_base_ptr);
|
||||
__syncwarp();
|
||||
tmem_allocation_result_barrier.arrive();
|
||||
uint32_t tmem_base_ptr = shared_storage.tmem_base_ptr;
|
||||
bulk_tmem.data() = tmem_base_ptr;
|
||||
|
||||
// Pass the acc with tuple type since the bgrad kernel change the mma_init API
|
||||
auto mma_inputs = collective_mainloop.mma_init(params.mainloop, cute::make_tuple(bulk_tmem, bulk_tmem), shared_storage.tensors.mainloop);
|
||||
do {
|
||||
auto k_tile_count = TileScheduler::get_work_k_tile_count(work_tile_info, problem_shape_MNKL, CtaShape_MNK{});
|
||||
|
||||
// Fetch next work tile
|
||||
auto [next_work_tile_info, increment_pipe] = scheduler.fetch_next_work(
|
||||
work_tile_info,
|
||||
clc_pipeline,
|
||||
clc_pipe_consumer_state
|
||||
);
|
||||
|
||||
if (increment_pipe) {
|
||||
++clc_pipe_consumer_state;
|
||||
}
|
||||
|
||||
// Wait for tmem accumulator buffer to become empty with a flipped phase
|
||||
accumulator_pipeline.producer_acquire(accumulator_pipe_producer_state);
|
||||
|
||||
int acc_stage = accumulator_pipe_producer_state.index();
|
||||
Tensor accumulators = bulk_tmem(_,_,_,acc_stage);
|
||||
mainloop_pipe_consumer_state = collective_mainloop.mma(
|
||||
mainloop_pipeline,
|
||||
mainloop_pipe_consumer_state,
|
||||
// Pass the acc with tuple type since the bgrad kernel change the mma API
|
||||
cute::make_tuple(accumulators, accumulators),
|
||||
mma_inputs,
|
||||
k_tile_count
|
||||
);
|
||||
|
||||
accumulator_pipeline.producer_commit(accumulator_pipe_producer_state);
|
||||
|
||||
++accumulator_pipe_producer_state;
|
||||
work_tile_info = next_work_tile_info;
|
||||
} while (work_tile_info.is_valid());
|
||||
// Release the right to allocate before deallocations so that the next CTA can rasterize
|
||||
tmem_allocator.release_allocation_lock();
|
||||
|
||||
accumulator_pipeline.producer_tail(accumulator_pipe_producer_state);
|
||||
|
||||
// Free entire tmem allocation
|
||||
tmem_allocator.free(tmem_base_ptr, TmemAllocator::Sm100TmemCapacityColumns);
|
||||
}
|
||||
|
||||
else if (is_participant.epi_load) {
|
||||
cutlass::arch::warpgroup_reg_dealloc<NonEpilogueWarpRegs>();
|
||||
|
||||
bool do_tail_load = false;
|
||||
do {
|
||||
bool compute_epilogue = TileScheduler::compute_epilogue(work_tile_info, params.scheduler);
|
||||
|
||||
// Get current work tile and fetch next work tile
|
||||
auto [next_work_tile_info, increment_pipe] = scheduler.fetch_next_work(
|
||||
work_tile_info,
|
||||
clc_pipeline,
|
||||
clc_pipe_consumer_state
|
||||
);
|
||||
work_tile_info = next_work_tile_info;
|
||||
|
||||
if (increment_pipe) {
|
||||
++clc_pipe_consumer_state;
|
||||
}
|
||||
|
||||
if (compute_epilogue) {
|
||||
|
||||
epi_load_pipe_producer_state = collective_epilogue.load(
|
||||
epi_load_pipeline,
|
||||
epi_load_pipe_producer_state,
|
||||
problem_shape_MNKL,
|
||||
CtaShape_MNK{},
|
||||
cta_coord_mnkl,
|
||||
TileShape{},
|
||||
TiledMma{},
|
||||
shared_storage.tensors.epilogue
|
||||
);
|
||||
|
||||
do_tail_load = true;
|
||||
}
|
||||
|
||||
// Calculate the cta coordinates of the next work tile
|
||||
cta_coord_mnkl = scheduler.work_tile_to_cta_coord(work_tile_info);
|
||||
} while (work_tile_info.is_valid());
|
||||
if (do_tail_load) {
|
||||
collective_epilogue.load_tail(
|
||||
epi_load_pipeline, epi_load_pipe_producer_state,
|
||||
epi_store_pipeline, epi_store_pipe_producer_state);
|
||||
}
|
||||
}
|
||||
|
||||
else if (is_participant.epilogue) {
|
||||
cutlass::arch::warpgroup_reg_alloc<EpilogueWarpRegs>();
|
||||
|
||||
// Wait for tmem allocate here
|
||||
tmem_allocation_result_barrier.arrive_and_wait();
|
||||
uint32_t tmem_base_ptr = shared_storage.tmem_base_ptr;
|
||||
bulk_tmem.data() = tmem_base_ptr;
|
||||
|
||||
bool do_tail_store = false;
|
||||
do {
|
||||
// Fetch next work tile
|
||||
auto [next_work_tile_info, increment_pipe] = scheduler.fetch_next_work(
|
||||
work_tile_info,
|
||||
clc_pipeline,
|
||||
clc_pipe_consumer_state
|
||||
);
|
||||
|
||||
if (increment_pipe) {
|
||||
++clc_pipe_consumer_state;
|
||||
}
|
||||
// Accumulator stage slice
|
||||
int acc_stage = accumulator_pipe_consumer_state.index();
|
||||
Tensor accumulators = bulk_tmem(_,_,_,acc_stage);
|
||||
|
||||
accumulator_pipe_consumer_state = scheduler.template fixup<IsComplex>(
|
||||
TiledMma{},
|
||||
work_tile_info,
|
||||
accumulators,
|
||||
accumulator_pipeline,
|
||||
accumulator_pipe_consumer_state,
|
||||
typename CollectiveEpilogue::CopyOpT2R{}
|
||||
);
|
||||
|
||||
//
|
||||
// Epilogue and write to gD
|
||||
//
|
||||
if (scheduler.compute_epilogue(work_tile_info)) {
|
||||
auto [load_state_next, store_state_next, acc_state_next] = collective_epilogue.store(
|
||||
epi_load_pipeline,
|
||||
epi_load_pipe_consumer_state,
|
||||
epi_store_pipeline,
|
||||
epi_store_pipe_producer_state,
|
||||
accumulator_pipeline,
|
||||
accumulator_pipe_consumer_state,
|
||||
problem_shape_MNKL,
|
||||
CtaShape_MNK{},
|
||||
cta_coord_mnkl,
|
||||
TileShape{},
|
||||
TiledMma{},
|
||||
accumulators,
|
||||
shared_storage.tensors.epilogue
|
||||
);
|
||||
epi_load_pipe_consumer_state = load_state_next;
|
||||
epi_store_pipe_producer_state = store_state_next;
|
||||
accumulator_pipe_consumer_state = acc_state_next;
|
||||
do_tail_store = true;
|
||||
}
|
||||
|
||||
work_tile_info = next_work_tile_info;
|
||||
cta_coord_mnkl = scheduler.work_tile_to_cta_coord(work_tile_info);
|
||||
|
||||
} while (work_tile_info.is_valid());
|
||||
if (do_tail_store) {
|
||||
collective_epilogue.store_tail(
|
||||
epi_load_pipeline, epi_load_pipe_consumer_state,
|
||||
epi_store_pipeline, epi_store_pipe_producer_state,
|
||||
CtaShape_MNK{});
|
||||
}
|
||||
}
|
||||
|
||||
else {
|
||||
cutlass::arch::warpgroup_reg_dealloc<NonEpilogueWarpRegs>();
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
///////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace cutlass::gemm::kernel
|
||||
+1090
File diff suppressed because it is too large
Load Diff
@@ -88,9 +88,7 @@ public:
|
||||
static_assert(cute::is_static<TileShape>::value);
|
||||
|
||||
auto selected_cluster_shape = cutlass::detail::select_cluster_shape(cluster_shape_mnk, hw_info.cluster_shape);
|
||||
auto cta_shape = cute::conditional_return<not cute::is_static_v<ClusterShape>>(
|
||||
shape_div(tile_shape_mnk, atom_thr_shape_mnk), // Dynamic Cluster: For 2SM kernels, use CTA tile shape for the underlying scheduler
|
||||
shape_div(tile_shape_mnk, selected_cluster_shape)); // Static Cluster: Blackwell builders expects TileShape to be Cluster's Tile Shape, Hopper doesn't
|
||||
auto cta_shape = shape_div(tile_shape_mnk, atom_thr_shape_mnk); // For 2SM kernels, use CTA tile shape for the underlying scheduler
|
||||
|
||||
dim3 problem_blocks = get_tiled_cta_shape_mnl(
|
||||
problem_shapes,
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -411,9 +411,11 @@ public:
|
||||
using namespace cute;
|
||||
using X = Underscore;
|
||||
|
||||
#if (defined(__CUDA_ARCH_FEAT_SM90_ALL) || defined(__CUDA_ARCH_FEAT_SM120_ALL) || CUDA_ARCH_CONDITIONAL_OR_FAMILY(1200))
|
||||
# define ENABLE_SM90_KERNEL_LEVEL 1
|
||||
#endif
|
||||
# if (defined(__CUDA_ARCH_FEAT_SM90_ALL) || defined(__CUDA_ARCH_FEAT_SM120_ALL) || defined(__CUDA_ARCH_FEAT_SM121_ALL) ||\
|
||||
CUDA_ARCH_CONDITIONAL_OR_FAMILY(1200) || CUDA_ARCH_CONDITIONAL_OR_FAMILY(1210))
|
||||
# define ENABLE_SM90_KERNEL_LEVEL 1
|
||||
# endif
|
||||
|
||||
// Any Tensor Op MMA Atom in the ISA is arch conditional.
|
||||
#if ! defined(ENABLE_SM90_KERNEL_LEVEL)
|
||||
printf("ERROR : Arch conditional MMA instruction used without targeting appropriate compute capability. Aborting.\n");
|
||||
|
||||
@@ -423,9 +423,11 @@ public:
|
||||
using namespace cute;
|
||||
using X = Underscore;
|
||||
|
||||
#if (defined(__CUDA_ARCH_FEAT_SM90_ALL) || defined(__CUDA_ARCH_FEAT_SM120_ALL) || CUDA_ARCH_CONDITIONAL_OR_FAMILY(1200))
|
||||
# define ENABLE_SM90_KERNEL_LEVEL 1
|
||||
#endif
|
||||
# if (defined(__CUDA_ARCH_FEAT_SM90_ALL) || defined(__CUDA_ARCH_FEAT_SM120_ALL) || defined(__CUDA_ARCH_FEAT_SM121_ALL) ||\
|
||||
CUDA_ARCH_CONDITIONAL_OR_FAMILY(1200) || CUDA_ARCH_CONDITIONAL_OR_FAMILY(1210))
|
||||
# define ENABLE_SM90_KERNEL_LEVEL 1
|
||||
# endif
|
||||
|
||||
// Any Tensor Op MMA Atom in the ISA is arch conditional.
|
||||
#if ! defined(ENABLE_SM90_KERNEL_LEVEL)
|
||||
printf("ERROR : Arch conditional MMA instruction used without targeting appropriate compute capability. Aborting.\n");
|
||||
|
||||
@@ -270,9 +270,11 @@ public:
|
||||
using namespace cute;
|
||||
using X = Underscore;
|
||||
|
||||
#if (defined(__CUDA_ARCH_FEAT_SM90_ALL) || defined(__CUDA_ARCH_FEAT_SM120_ALL) || CUDA_ARCH_CONDITIONAL_OR_FAMILY(1200))
|
||||
# define ENABLE_SM90_KERNEL_LEVEL 1
|
||||
#endif
|
||||
# if (defined(__CUDA_ARCH_FEAT_SM90_ALL) || defined(__CUDA_ARCH_FEAT_SM120_ALL) || defined(__CUDA_ARCH_FEAT_SM121_ALL) ||\
|
||||
CUDA_ARCH_CONDITIONAL_OR_FAMILY(1200) || CUDA_ARCH_CONDITIONAL_OR_FAMILY(1210))
|
||||
# define ENABLE_SM90_KERNEL_LEVEL 1
|
||||
# endif
|
||||
|
||||
// Any Tensor Op MMA Atom in the WGMMA ISA is arch conditional to sm90a.
|
||||
#if ! defined(ENABLE_SM90_KERNEL_LEVEL)
|
||||
printf("ERROR : Arch conditional MMA instruction used without targeting sm90a compute capability. Aborting.\n");
|
||||
|
||||
@@ -342,9 +342,11 @@ public:
|
||||
using namespace cute;
|
||||
using X = Underscore;
|
||||
|
||||
#if (defined(__CUDA_ARCH_FEAT_SM90_ALL) || defined(__CUDA_ARCH_FEAT_SM120_ALL) || CUDA_ARCH_CONDITIONAL_OR_FAMILY(1200))
|
||||
# define ENABLE_SM90_KERNEL_LEVEL 1
|
||||
#endif
|
||||
# if (defined(__CUDA_ARCH_FEAT_SM90_ALL) || defined(__CUDA_ARCH_FEAT_SM120_ALL) || defined(__CUDA_ARCH_FEAT_SM121_ALL) ||\
|
||||
CUDA_ARCH_CONDITIONAL_OR_FAMILY(1200) || CUDA_ARCH_CONDITIONAL_OR_FAMILY(1210))
|
||||
# define ENABLE_SM90_KERNEL_LEVEL 1
|
||||
# endif
|
||||
|
||||
// Any Tensor Op MMA Atom in the ISA is arch conditional.
|
||||
#if ! defined(ENABLE_SM90_KERNEL_LEVEL)
|
||||
printf("ERROR : Arch conditional MMA instruction used without targeting appropriate compute capability. Aborting.\n");
|
||||
|
||||
@@ -354,9 +354,11 @@ public:
|
||||
using namespace cute;
|
||||
using X = Underscore;
|
||||
|
||||
#if (defined(__CUDA_ARCH_FEAT_SM90_ALL) || defined(__CUDA_ARCH_FEAT_SM120_ALL) || CUDA_ARCH_CONDITIONAL_OR_FAMILY(1200))
|
||||
# define ENABLE_SM90_KERNEL_LEVEL 1
|
||||
#endif
|
||||
# if (defined(__CUDA_ARCH_FEAT_SM90_ALL) || defined(__CUDA_ARCH_FEAT_SM120_ALL) || defined(__CUDA_ARCH_FEAT_SM121_ALL) ||\
|
||||
CUDA_ARCH_CONDITIONAL_OR_FAMILY(1200) || CUDA_ARCH_CONDITIONAL_OR_FAMILY(1210))
|
||||
# define ENABLE_SM90_KERNEL_LEVEL 1
|
||||
# endif
|
||||
|
||||
// Any Tensor Op MMA Atom in the ISA is arch conditional.
|
||||
#if ! defined(ENABLE_SM90_KERNEL_LEVEL)
|
||||
printf("ERROR : Arch conditional MMA instruction used without targeting appropriate compute capability. Aborting.\n");
|
||||
|
||||
@@ -298,6 +298,66 @@ struct TileSchedulerSelector<
|
||||
using Scheduler = StaticPersistentTileScheduler100;
|
||||
};
|
||||
|
||||
template <class TileShape, class ClusterShape, uint32_t SchedulerPipelineStageCount>
|
||||
struct TileSchedulerSelector<
|
||||
PersistentScheduler,
|
||||
arch::Sm103,
|
||||
TileShape,
|
||||
ClusterShape,
|
||||
SchedulerPipelineStageCount> {
|
||||
using Scheduler = PersistentTileSchedulerSm100<
|
||||
ClusterShape,
|
||||
SchedulerPipelineStageCount>;
|
||||
};
|
||||
|
||||
// Ptr-Array kernel may provide a specialized ArrayProblemShape type
|
||||
template <class TileShape,
|
||||
class ClusterShape,
|
||||
uint32_t SchedulerPipelineStageCount,
|
||||
class ProblemShape>
|
||||
struct TileSchedulerSelector<
|
||||
PersistentScheduler,
|
||||
arch::Sm103,
|
||||
TileShape,
|
||||
ClusterShape,
|
||||
SchedulerPipelineStageCount,
|
||||
ProblemShape> {
|
||||
using Scheduler = PersistentTileSchedulerSm100<
|
||||
ClusterShape,
|
||||
SchedulerPipelineStageCount>;
|
||||
};
|
||||
|
||||
// SM103 Group tile scheduler
|
||||
template <
|
||||
class TileShape,
|
||||
class ClusterShape,
|
||||
uint32_t SchedulerPipelineStageCount,
|
||||
class GroupProblemShape
|
||||
>
|
||||
struct TileSchedulerSelector<
|
||||
GroupScheduler,
|
||||
arch::Sm103,
|
||||
TileShape,
|
||||
ClusterShape,
|
||||
SchedulerPipelineStageCount,
|
||||
GroupProblemShape
|
||||
> {
|
||||
using Scheduler = PersistentTileSchedulerSm100Group<GroupProblemShape, SchedulerPipelineStageCount>;
|
||||
};
|
||||
|
||||
template <class TileShape, class ClusterShape, uint32_t SchedulerPipelineStageCount>
|
||||
struct TileSchedulerSelector<
|
||||
StreamKScheduler,
|
||||
arch::Sm103,
|
||||
TileShape,
|
||||
ClusterShape,
|
||||
SchedulerPipelineStageCount> {
|
||||
using Scheduler = PersistentTileSchedulerSm100StreamK<
|
||||
TileShape,
|
||||
ClusterShape,
|
||||
SchedulerPipelineStageCount>;
|
||||
};
|
||||
|
||||
// Default (void) for Sm120 maps to PersistentTileSchedulerSm100
|
||||
template <class TileShape, class ClusterShape, uint32_t SchedulerPipelineStageCount>
|
||||
struct TileSchedulerSelector<
|
||||
|
||||
Reference in New Issue
Block a user