leavelet/cutlass
1
0
Fork 0
Code Issues Pull Requests Actions 29 Packages Projects Releases Wiki Activity

v4.0 update. (#2371)

Browse Source
This commit is contained in:
Junkai-Wu
2025-06-06 02:39:20 -04:00
committed by GitHub
parent 2e2af190bd
commit 8bdbfca682
254 changed files with 29751 additions and 1980 deletions
Download Patch File Download Diff File Expand all files Collapse all files
include/cute/algorithm/axpby.hpp
+2 -3
View File
@@ -33,7 +33,6 @@
#include <cute/config.hpp>
#include <cute/tensor_impl.hpp>
#include <cute/tensor_predicate.hpp>
namespace cute
{
@@ -45,7 +44,7 @@ template <class Alpha,
class XEngine, class XLayout,
class Beta,
class YEngine, class YLayout,
class PrdTensor = TrivialPredTensor>
class PrdTensor = constant_fn<true_type>>
CUTE_HOST_DEVICE
void
axpby(Alpha const& alpha,
@@ -64,7 +63,7 @@ template <class Alpha,
class XEngine, class XLayout,
class Beta,
class YEngine, class YLayout,
class PrdTensor = TrivialPredTensor>
class PrdTensor = constant_fn<true_type>>
CUTE_HOST_DEVICE
void
axpby(Alpha const& alpha,
include/cute/algorithm/cooperative_copy.hpp
-1
View File
@@ -36,7 +36,6 @@
#include <cute/swizzle.hpp> // cute::Swizzle
#include <cute/swizzle_layout.hpp> // cute::get_nonswizzle_portion
#include <cute/tensor_impl.hpp> // cute::Tensor
#include <cute/tensor_predicate.hpp>
#include <cute/algorithm/copy.hpp>
#include <cute/atom/copy_atom.hpp>
include/cute/algorithm/copy.hpp
+77 -43
View File
@@ -32,7 +32,6 @@
#include <cute/config.hpp> // CUTE_HOST_DEVICE
#include <cute/tensor_impl.hpp> // cute::Tensor
#include <cute/tensor_predicate.hpp> // cute::TrivialPredTensor
#include <cute/atom/copy_atom.hpp> // cute::Copy_Atom
namespace cute
@@ -66,10 +65,45 @@ copy_if(PrdTensor const& pred,
// copy_if -- Predicated CopyAtom
//
// Predicate Tensor is an Actual Tensor
template <class... CopyArgs,
class PrdEngine, class PrdLayout,
class SrcEngine, class SrcLayout,
class DstEngine, class DstLayout>
CUTE_HOST_DEVICE
void
copy_if(Copy_Atom<CopyArgs...> const& copy_atom,
Tensor<PrdEngine, PrdLayout> const& prd, // ([V],Rest...)
Tensor<SrcEngine, SrcLayout> const& src, // ( V, Rest...)
Tensor<DstEngine, DstLayout> & dst) // ( V, Rest...)
{
if constexpr (PrdLayout::rank == SrcLayout::rank - 1) {
// Back-compat ONLY -- Delete?
copy_if(copy_atom, make_tensor(prd.data(), prepend(prd.layout(), Layout<_1,_0>{})), src, dst);
} else {
static_assert(SrcLayout::rank == DstLayout::rank, "CopyAtom rank-mismatch.");
static_assert(SrcLayout::rank == PrdLayout::rank, "CopyAtom rank-mismatch.");
if constexpr (SrcLayout::rank == 1) { // Dispatch the copy
copy_atom.call(prd, src, dst);
} else { // Loop over all but the first mode
constexpr int R = SrcLayout::rank;
Tensor prd_v = group_modes<1,R>(prd);
Tensor src_v = group_modes<1,R>(src);
Tensor dst_v = group_modes<1,R>(dst);
CUTE_UNROLL
for (int i = 0; i < size<1>(dst_v); ++i) {
copy_atom.call(prd_v(_,i), src_v(_,i), dst_v(_,i));
}
}
}
}
template <class... CopyArgs,
class PredTensor,
class SrcEngine, class SrcLayout,
class DstEngine, class DstLayout>
[[deprecated("Use a bool-tensor or transform-tensor as predication.")]]
CUTE_HOST_DEVICE
void
copy_if(Copy_Atom<CopyArgs...> const& copy_atom,
@@ -77,33 +111,14 @@ copy_if(Copy_Atom<CopyArgs...> const& copy_atom,
Tensor<SrcEngine, SrcLayout> const& src, // (V,Rest...)
Tensor<DstEngine, DstLayout> & dst) // (V,Rest...)
{
static_assert(SrcLayout::rank == DstLayout::rank, "CopyAtom rank-mismatch.");
auto has_with_bool = cute::is_valid([](auto t)->void_t<decltype(declval<typename decltype(t)::Traits>().with(true))>{}, copy_atom);
if constexpr (SrcLayout::rank == 1) { // Dispatch the copy
if constexpr (has_with_bool) {
copy_atom.with(pred()).call(src, dst);
} else {
if (pred()) { copy_atom.call(src, dst); }
}
} else { // Loop over all but the first mode
constexpr int R = SrcLayout::rank;
Tensor src_v = group_modes<1,R>(src);
Tensor dst_v = group_modes<1,R>(dst);
CUTE_UNROLL
for (int i = 0; i < size<1>(dst_v); ++i) {
if constexpr (has_with_bool) {
copy_atom.with(pred(i)).call(src_v(_,i), dst_v(_,i));
} else {
if (pred(i)) { copy_atom.call(src_v(_,i), dst_v(_,i)); }
}
}
}
Tensor tpred = cute::lazy::transform(make_tensor(counting_iterator<int>{}, replace<0>(shape(dst), _1{})), pred);
return copy_if(copy_atom, tpred, src, dst);
}
//
// copy_if -- AutoCopyAsync
//
template <class PrdTensor,
class SrcEngine, class SrcLayout,
class DstEngine, class DstLayout>
@@ -159,7 +174,7 @@ copy(AutoCopyAsync const& cpy,
Tensor<SrcEngine, SrcLayout> const& src, // (V,Rest...)
Tensor<DstEngine, DstLayout> & dst) // (V,Rest...)
{
copy_if(cpy, TrivialPredTensor{}, src, dst);
copy_if(cpy, constant_fn<true_type>{}, src, dst);
}
//
@@ -202,7 +217,7 @@ copy(Copy_Atom<CopyArgs...> const& copy_atom,
Tensor dst_c = dst_n(make_coord(_,Int<0>{}),make_coord(Int<0>{},_)); // (V, Rest)
Tensor src_c = src_n(make_coord(_,Int<0>{}),make_coord(Int<0>{},_)); // (V, Rest)
CUTE_STATIC_ASSERT_V(size<1>(src_c) == size<1>(dst_c));
CUTE_STATIC_ASSERT_V( size<1>(src_c) == size<1>(dst_c));
CUTE_STATIC_ASSERT_V(shape<0>(dst_c) == shape<0>(dst));
CUTE_STATIC_ASSERT_V(shape<0>(src_c) == shape<0>(src));
@@ -224,7 +239,7 @@ copy(Copy_Atom<CopyArgs...> const& copy_atom,
////////////////////////////////////////////////////////
// Specialization for AutoVectorizingCopyAssumedAlignment<MaxVecBits>
template <int MaxVecBits, class... Args,
template <int MaxVecBits,
class SrcEngine, class SrcLayout,
class DstEngine, class DstLayout>
CUTE_HOST_DEVICE
@@ -234,23 +249,30 @@ copy(AutoVectorizingCopyWithAssumedAlignment<MaxVecBits> const&,
Tensor<DstEngine, DstLayout> & dst)
{
constexpr int common_elem = CUTE_STATIC_V(max_common_vector(src, dst));
constexpr int align_bits = CUTE_STATIC_V(gcd(max_alignment(src), max_alignment(dst), Int<MaxVecBits>{}));
static_assert(is_integral<decltype(Int<common_elem>{} * sizeof_bits_v<typename SrcEngine::value_type>)>::value, "Error: Attempting a subbit copy!");
constexpr int vec_bits = gcd(common_elem * sizeof_bits_v<typename SrcEngine::value_type>, align_bits);
static_assert(is_integral<decltype(Int<common_elem>{} * sizeof_bits_v<typename DstEngine::value_type>)>::value, "Error: Attempting a subbit write!");
if constexpr (common_elem > 1 && ((vec_bits % 8) == 0)) {
// If more than one element vectorizes to 8bits or more, then recast and copy
using VecType = uint_bit_t<vec_bits>;
// Preserve volatility
using SrcVecType = conditional_t<is_volatile_v<typename SrcEngine::element_type>, VecType const volatile, VecType const>;
using DstVecType = conditional_t<is_volatile_v<typename DstEngine::element_type>, VecType volatile, VecType >;
if constexpr (common_elem > 1)
{
constexpr int align_bits = CUTE_STATIC_V(gcd(max_alignment(src), max_alignment(dst), Int<MaxVecBits>{}));
constexpr int vec_bits = gcd(common_elem * sizeof_bits_v<typename SrcEngine::value_type>, align_bits);
// Recast
Tensor src_v = recast<SrcVecType>(src);
Tensor dst_v = recast<DstVecType>(dst);
return copy_if(TrivialPredTensor{}, src_v, dst_v);
if constexpr ((vec_bits % 8) == 0)
{
// If more than one element vectorizes to 8bits or more, then recast and copy
using VecType = uint_bit_t<vec_bits>;
// Preserve volatility
using SrcVecType = conditional_t<is_volatile_v<typename SrcEngine::element_type>, VecType const volatile, VecType const>;
using DstVecType = conditional_t<is_volatile_v<typename DstEngine::element_type>, VecType volatile, VecType >;
// Recast
Tensor src_v = recast<SrcVecType>(src);
Tensor dst_v = recast<DstVecType>(dst);
return copy_if(constant_fn<true_type>{}, src_v, dst_v);
} else {
return copy_if(constant_fn<true_type>{}, src, dst);
}
} else {
return copy_if(TrivialPredTensor{}, src, dst);
return copy_if(constant_fn<true_type>{}, src, dst);
}
}
@@ -277,7 +299,7 @@ copy(AutoFilter<CopyOp> const& copy_op,
Tensor src_n = zipped_divide(src, dst_null);
CUTE_STATIC_ASSERT_V(cosize<0>(dst_n.layout()) == Int<1>{}, "Nullspace definition error");
CUTE_STATIC_ASSERT_V(cosize<0>(src_n.layout()) == Int<1>{}, "Error: Ambiguous scatter detected in copy");
CUTE_STATIC_ASSERT_V(cosize<0>(src_n.layout()) == Int<1>{}, "Error: Ambiguous race-condition detected.");
copy(copy_op.base, src_n(Int<0>{},_), dst_n(Int<0>{},_));
} else {
@@ -335,6 +357,18 @@ copy(Copy_Atom<AutoVectorizingCopyWithAssumedAlignment<MaxVecBits>, Args...> con
return copy(AutoVectorizingCopyWithAssumedAlignment<MaxVecBits>{}, src, dst);
}
template <int MaxVecBits, class... Args,
class SrcEngine, class SrcLayout,
class DstEngine, class DstLayout>
CUTE_HOST_DEVICE
void
copy(Copy_Atom<Copy_Traits<AutoVectorizingCopyWithAssumedAlignment<MaxVecBits>>, Args...> const&,
Tensor<SrcEngine, SrcLayout> const& src,
Tensor<DstEngine, DstLayout> & dst)
{
return copy(AutoVectorizingCopyWithAssumedAlignment<MaxVecBits>{}, src, dst);
}
#if defined(CUTE_COPY_ATOM_TMA_SM90_ENABLED)
template <class... CT_Args,
class SrcEngine, class SrcLayout,
@@ -375,8 +409,8 @@ template <class... CT_Args, class... CA_Args,
CUTE_HOST_DEVICE
void
copy(Copy_Atom<Copy_Traits<SM90_BULK_COPY_AUTO, CT_Args...>, CA_Args...> const& atom,
Tensor<SrcEngine, SrcLayout> const& src,
Tensor<DstEngine, DstLayout> & dst)
Tensor<SrcEngine, SrcLayout> const& src,
Tensor<DstEngine, DstLayout> & dst)
{
return copy(static_cast<Copy_Traits<SM90_BULK_COPY_AUTO, CT_Args...> const&>(atom), src, dst);
}
include/cute/algorithm/prefetch.hpp
+6 -5
View File
@@ -90,18 +90,19 @@ constexpr bool has_prefetch<CopyOp, void_t<typename CopyOp::PREFETCH>> = true;
} // end namespace detail
template <class CopyOp, class... CT_Args, class... CA_Args,
template <class CopyOp, class... CT_Args, class CopyType,
class GEngine, class GLayout>
CUTE_HOST_DEVICE
void
prefetch(Copy_Atom<Copy_Traits<CopyOp, CT_Args...>, CA_Args...> const& atom,
Tensor<GEngine, GLayout> const& src)
prefetch(Copy_Atom<Copy_Traits<CopyOp, CT_Args...>, CopyType> const& atom,
Tensor<GEngine, GLayout> const& src)
{
if constexpr (detail::has_prefetch<CopyOp>) {
using Prefetch_Traits = Copy_Traits<typename CopyOp::PREFETCH, CT_Args...>;
using Prefetch_Atom = Copy_Atom<Prefetch_Traits, CA_Args...>;
using Prefetch_Atom = Copy_Atom<Prefetch_Traits, CopyType>;
Prefetch_Atom prefetch_atom{atom};
auto& dst = const_cast<Tensor<GEngine, GLayout>&>(src); // dst is ignored for prefetch atoms
//auto& dst = const_cast<Tensor<GEngine, GLayout>&>(src); // dst is ignored for prefetch atoms
Tensor dst = make_tensor(make_smem_ptr<CopyType>(nullptr), shape(src));
return copy(prefetch_atom, src, dst);
} else {
return prefetch(src);
include/cute/algorithm/tensor_algorithms.hpp
+12
View File
@@ -163,4 +163,16 @@ transform(Tensor<EngineIn1,LayoutIn1> const& tensor_in1,
return transform(tensor_in1, tensor_in2, tensor_out, op);
}
namespace lazy {
template <class Engine, class Layout, class Fn>
CUTE_HOST_DEVICE constexpr
auto
transform(cute::Tensor<Engine,Layout> const& t, Fn const& fn)
{
return cute::make_tensor(cute::make_transform_iter(fn, t.data()), t.layout());
}
} // end namespace lazy
} // end namespace cute
include/cute/algorithm/tensor_reduce.hpp
+107
View File
@@ -0,0 +1,107 @@
/***************************************************************************************************
* Copyright (c) 2025 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-License-Identifier: BSD-3-Clause
*
* Redistribution and use in source and binary forms, with or without
* modification, are permitted provided that the following conditions are met:
*
* 1. Redistributions of source code must retain the above copyright notice, this
* list of conditions and the following disclaimer.
*
* 2. Redistributions in binary form must reproduce the above copyright notice,
* this list of conditions and the following disclaimer in the documentation
* and/or other materials provided with the distribution.
*
* 3. Neither the name of the copyright holder nor the names of its
* contributors may be used to endorse or promote products derived from
* this software without specific prior written permission.
*
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
* AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
* IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
* DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
* DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
* SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
* CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
* OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*
**************************************************************************************************/
#pragma once
#include <iostream>
#include <cute/config.hpp>
#include <cute/tensor_impl.hpp>
#include <cute/algorithm/functional.hpp>
#include <cute/algorithm/fill.hpp>
namespace cute
{
// Reduce @src tensor using binary reduction operator @op and initial value @init and return a scalar.
template <class SrcEngine, class SrcLayout, class T, class BinaryOp = cute::plus>
CUTE_HOST_DEVICE constexpr
T
reduce(Tensor<SrcEngine,SrcLayout> const& src, T init, BinaryOp op = {})
{
for (auto i = 0; i < size(src); ++i) {
init = op(init, src(i));
}
return init;
}
// Reduce @src tensor RedMode using binary reduction operator @op and store the result in @dst tensor
// for each index in @dst/BatchMode.
// @pre @src tensor has rank 2
// @pre size of @src batch mode is equal to size of @dst batch mode
template <class SrcEngine, class SrcLayout,
class DstEngine, class DstLayout,
class BinaryOp = cute::plus>
CUTE_HOST_DEVICE constexpr
void
batch_reduce(Tensor<SrcEngine, SrcLayout> const& src, // (RedMode, BatchMode)
Tensor<DstEngine, DstLayout> & dst, // (BatchMode)
BinaryOp op = {})
{
// Precondition
CUTE_STATIC_ASSERT_V(rank(src) == Int<2>{});
assert(size<1>(src) == size(dst));
for (int i = 0; i < size(dst); ++i) {
dst(i) = reduce(src(_,i), dst(i), op);
}
}
// Reduce @src tensor along selected modes specified in @target_profile using binary reduction operator @op
// and store the result in @dst tensor. @target_profile is a tuple where '_' indicates modes to keep and
// integers indicates modes to reduce.
// @pre @target_profile is compatible with @src layout
template <class SrcEngine, class SrcLayout,
class DstEngine, class DstLayout,
class TargetProfile,
class BinaryOp = cute::plus>
CUTE_HOST_DEVICE constexpr
void
logical_reduce(Tensor<SrcEngine, SrcLayout> const& src,
Tensor<DstEngine, DstLayout> & dst,
TargetProfile const& target_profile,
BinaryOp op = {})
{
// Precondition
assert(compatible(target_profile, shape(src)));
auto diced_layout = dice(target_profile, src.layout());
auto sliced_layout = slice(target_profile, src.layout());
auto red_mode = conditional_return<rank(diced_layout) == Int<0>{}>(Layout<_1,_0>{}, diced_layout);
auto batch_mode = conditional_return<rank(sliced_layout) == Int<0>{}>(Layout<_1,_0>{}, sliced_layout);
auto src_tensor = make_tensor(src.data(), make_layout(red_mode, batch_mode));
batch_reduce(src_tensor, dst, op);
}
} // end namespace cute
include/cute/atom/copy_atom.hpp
+52 -2
View File
@@ -123,6 +123,56 @@ struct Copy_Atom<Copy_Traits<Args...>, CopyInternalType>
{
return call(src, dst);
}
// Check and call instruction, or recurse
template <class PEngine, class PLayout,
class SEngine, class SLayout,
class DEngine, class DLayout>
CUTE_HOST_DEVICE
void
call(Tensor<PEngine,PLayout> const& prd,
Tensor<SEngine,SLayout> const& src,
Tensor<DEngine,DLayout> & dst) const
{
static_assert(PLayout::rank == 1, "Expected rank-1 prd tensor");
static_assert(SLayout::rank == 1, "Expected rank-1 src tensor");
static_assert(DLayout::rank == 1, "Expected rank-1 dst tensor");
if constexpr (is_constant<NumValSrc, decltype(size(src))>::value ||
is_constant<NumValDst, decltype(size(dst))>::value) {
// Dispatch to unpack to execute instruction
Traits const& traits = static_cast<Traits const&>(*this);
auto has_with_bool = cute::is_valid([](auto t)->void_t<decltype(t.with(true))>{}, traits);
if constexpr (has_with_bool) {
copy_unpack(traits.with(prd(Int<0>{})), src, dst);
} else {
if (prd(Int<0>{})) { copy_unpack(traits, src, dst); }
}
} else if constexpr (is_tuple<decltype(shape(prd))>::value &&
is_tuple<decltype(shape(src))>::value &&
is_tuple<decltype(shape(dst))>::value) {
// If the size of the src/dst doesn't match the instruction,
// recurse this rank-1 layout by peeling off the mode
// ((A,B,C,...)) -> (A,B,C,...)
return copy_if(*this, tensor<0>(prd), tensor<0>(src), tensor<0>(dst));
} else {
static_assert(dependent_false<SEngine>,
"CopyAtom: Src/Dst partitioning does not match the instruction requirement.");
}
}
// Accept mutable temporaries
template <class PEngine, class PLayout,
class SEngine, class SLayout,
class DEngine, class DLayout>
CUTE_HOST_DEVICE
void
call(Tensor<PEngine,PLayout> const& prd,
Tensor<SEngine,SLayout> const& src,
Tensor<DEngine,DLayout> && dst) const
{
return call(prd, src, dst);
}
};
//
@@ -733,13 +783,13 @@ print_latex_copy(LayoutS const& S, ThrIDS const& TS, // (m,n) -> (tid,vid) and
#include <cute/atom/copy_traits_sm75.hpp>
#include <cute/atom/copy_traits_sm80.hpp>
#include <cute/atom/copy_traits_sm90.hpp>
#include <cute/atom/copy_traits_sm100.hpp>
#include <cute/atom/copy_traits_sm100.hpp>
// Config
#if (__CUDACC_VER_MAJOR__ >= 12)
# define CUTE_COPY_ATOM_TMA_SM90_ENABLED
# define CUTE_COPY_ATOM_TMA_SM100_ENABLED
# define CUTE_COPY_ATOM_TMA_SM100_ENABLED
#endif
include/cute/pointer_base.hpp
+69
View File
@@ -235,6 +235,63 @@ raw_pointer_cast(counting_iterator<T> const& x) {
return x.n_;
}
//
// transform_iterator
//
template <class Fn, class Iter>
struct transform_iter
{
using iterator = Iter;
// using reference = typename iterator_traits<iterator>::reference;
// using element_type = typename iterator_traits<iterator>::element_type;
// using value_type = typename iterator_traits<iterator>::value_type;
Fn fn_;
iterator ptr_;
CUTE_HOST_DEVICE constexpr
transform_iter(Fn fn, iterator ptr = {}) : fn_(fn), ptr_(ptr) {}
CUTE_HOST_DEVICE constexpr
decltype(auto) operator*() const { return fn_(*ptr_); }
template <class Index>
CUTE_HOST_DEVICE constexpr
decltype(auto) operator[](Index const& i) const { return fn_(ptr_[i]); }
template <class Index>
CUTE_HOST_DEVICE constexpr
auto operator+(Index const& i) const { return transform_iter<Fn, decltype(ptr_+i)>{fn_, ptr_+i}; }
template <class IterY>
CUTE_HOST_DEVICE constexpr
friend bool operator==(transform_iter<Fn,Iter> const& x, transform_iter<Fn,IterY> const& y) { return x.ptr_ == y.ptr_; }
template <class IterY>
CUTE_HOST_DEVICE constexpr
friend bool operator!=(transform_iter<Fn,Iter> const& x, transform_iter<Fn,IterY> const& y) { return x.ptr_ != y.ptr_; }
template <class IterY>
CUTE_HOST_DEVICE constexpr
friend bool operator< (transform_iter<Fn,Iter> const& x, transform_iter<Fn,IterY> const& y) { return x.ptr_ < y.ptr_; }
template <class IterY>
CUTE_HOST_DEVICE constexpr
friend bool operator<=(transform_iter<Fn,Iter> const& x, transform_iter<Fn,IterY> const& y) { return x.ptr_ <= y.ptr_; }
template <class IterY>
CUTE_HOST_DEVICE constexpr
friend bool operator> (transform_iter<Fn,Iter> const& x, transform_iter<Fn,IterY> const& y) { return x.ptr_ > y.ptr_; }
template <class IterY>
CUTE_HOST_DEVICE constexpr
friend bool operator>=(transform_iter<Fn,Iter> const& x, transform_iter<Fn,IterY> const& y) { return x.ptr_ >= y.ptr_; }
};
template <class Fn, class Iterator>
CUTE_HOST_DEVICE constexpr
auto
make_transform_iter(Fn const& fn, Iterator const& ptr)
{
return transform_iter<Fn,Iterator>(fn,ptr);
}
//
// Display utilities
//
@@ -251,12 +308,24 @@ CUTE_HOST_DEVICE void print(counting_iterator<T> ptr)
printf("counting_iter("); print(ptr.n_); printf(")");
}
template <class Fn, class Iterator>
CUTE_HOST_DEVICE void print(transform_iter<Fn,Iterator> ptr)
{
printf("trans_"); print(ptr.ptr_);
}
#if !defined(__CUDACC_RTC__)
template <class T>
CUTE_HOST std::ostream& operator<<(std::ostream& os, counting_iterator<T> ptr)
{
return os << "counting_iter(" << ptr.n_ << ")";
}
template <class Fn, class Iterator>
CUTE_HOST std::ostream& operator<<(std::ostream& os, transform_iter<Fn,Iterator> ptr)
{
return os << "trans_" << ptr.ptr_;
}
#endif // !defined(__CUDACC_RTC__)
} // end namespace cute
include/cute/tensor_predicate.hpp
-78
View File
@@ -1,78 +0,0 @@
/***************************************************************************************************
* Copyright (c) 2023 - 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 <cute/config.hpp> // CUTE_HOST_DEVICE
#include <cute/numeric/integral_constant.hpp> // cute::true_type
namespace cute
{
template <class T>
struct ConstantTensor
{
template <class... Coords>
CUTE_HOST_DEVICE constexpr
T const&
operator()(Coords const&...) const {
return val_;
}
T val_;
};
struct TrivialPredTensor
{
template <class... Coords>
CUTE_HOST_DEVICE constexpr
true_type
operator()(Coords const&...) const {
return {};
}
};
template <class Fn>
struct FunctionPredTensor
{
CUTE_HOST_DEVICE constexpr
FunctionPredTensor(Fn const& fn) : fn_(fn) {}
template <class... Coords>
CUTE_HOST_DEVICE constexpr
auto
operator()(Coords const&... coords) const {
return fn_(coords...);
}
Fn const& fn_;
};
} // end namespace cute
include/cutlass/arch/grid_dependency_control.h
+2 -1
View File
@@ -41,7 +41,8 @@
#include "cutlass/gemm/dispatch_policy.hpp"
#ifndef CUTLASS_GDC_ENABLED
#if (defined(CUTLASS_ENABLE_GDC_FOR_SM90) && \
#if (CUDA_BARRIER_ENABLED && \
defined(CUTLASS_ENABLE_GDC_FOR_SM90) && \
__CUDACC_VER_MAJOR__ >= 12 && \
defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 && defined(__CUDA_ARCH_FEAT_SM90_ALL))
#define CUTLASS_GDC_ENABLED
include/cutlass/conv/collective/sm100_implicit_gemm_umma_warpspecialized.hpp
-1
View File
@@ -43,7 +43,6 @@
#include "cute/arch/cluster_sm90.hpp"
#include "cute/atom/mma_atom.hpp"
#include "cute/algorithm/gemm.hpp"
#include "cute/tensor_predicate.hpp"
#include "cute/numeric/arithmetic_tuple.hpp"
#include "cutlass/trace.h"
include/cutlass/conv/collective/sm90_implicit_gemm_gmma_ss_warpspecialized.hpp
+8 -9
View File
@@ -32,7 +32,6 @@
#include "cutlass/cutlass.h"
#include "cute/tensor_predicate.hpp"
#include "cute/arch/cluster_sm90.hpp"
#include "cute/arch/copy_sm90.hpp"
#include "cute/atom/mma_atom.hpp"
@@ -103,7 +102,7 @@ struct CollectiveConv<
using PipelineParams = typename MainloopPipeline::Params;
using PipelineState = typename cutlass::PipelineState<DispatchPolicy::Stages>;
using ProblemShape = ConvProblemShape<ConvOp, NumSpatialDimensions>;
static_assert(rank(SmemLayoutA{}) == 3, "SmemLayout must be rank 3 (M/N, K, PIPE)");
@@ -332,7 +331,7 @@ public:
TmaTransactionBytes
};
}
template <class ProblemShape>
static bool
can_implement(
@@ -409,7 +408,7 @@ public:
if constexpr (ConvOp == conv::Operator::kWgrad) {
#if defined(CUTLASS_DEBUG_TRACE_LEVEL) && (CUTLASS_DEBUG_TRACE_LEVEL > 1)
std::ostringstream os;
#endif
#endif
const auto & input_shape = problem_shape.shape_A;
const auto & input_stride = problem_shape.stride_A;
@@ -431,11 +430,11 @@ public:
<< "\n input_shape: " << input_shape
<< "\n input_stride: " << input_stride
<< "\n";
#endif
#endif
CUTLASS_TRACE_HOST(" CAN IMPLEMENT: Wgrad kernels don't support non-packed input strides.\n");
#if defined(CUTLASS_DEBUG_TRACE_LEVEL) && (CUTLASS_DEBUG_TRACE_LEVEL > 1)
CUTLASS_TRACE_HOST(os.str());
#endif
#endif
return false;
}
@@ -464,7 +463,7 @@ public:
CUTLASS_TRACE_HOST(" CAN IMPLEMENT: Wgrad kernels don't support non-packed output strides.\n");
#if defined(CUTLASS_DEBUG_TRACE_LEVEL) && (CUTLASS_DEBUG_TRACE_LEVEL > 1)
CUTLASS_TRACE_HOST(os.str());
#endif
#endif
return false;
}
}
@@ -516,8 +515,8 @@ public:
/// gA_mk - The tma tensor, A after a local tile so it has shape (BLK_M,BLK_K,m,k)
/// gB_nk - The tma tensor, B after a local tile so it has shape (BLK_N,BLK_K,n,k)
/// The rest of the tensors can be specified as needed by this collective.
/// The dimensions of gA_mk and gA_nk do not contain L to maintain consistency with
/// StrideA and StrideB set up for TMA
/// The dimensions of gA_mk and gA_nk do not contain L to maintain consistency with
/// StrideA and StrideB set up for TMA
template <class ProblemShapeMNKL>
CUTLASS_DEVICE auto
load_init(ProblemShapeMNKL const& problem_shape_MNKL, Params const& mainloop_params){
include/cutlass/conv/device/conv_universal_adapter.hpp
+25
View File
@@ -303,6 +303,16 @@ public:
dim3 cluster(cute::size<0>(typename ConvKernel::DispatchPolicy::ClusterShape{}),
cute::size<1>(typename ConvKernel::DispatchPolicy::ClusterShape{}),
cute::size<2>(typename ConvKernel::DispatchPolicy::ClusterShape{}));
// Dynamic cluster support
[[maybe_unused]] dim3 fallback_cluster = dim3{0,0,0};
if constexpr (ConvKernel::ArchTag::kMinComputeCapability == 100 ||
ConvKernel::ArchTag::kMinComputeCapability == 101) {
if constexpr (!cute::is_static_v<typename ConvKernel::DispatchPolicy::ClusterShape>) {
fallback_cluster = params.hw_info.cluster_shape_fallback;
cluster = params.hw_info.cluster_shape;
}
}
void* kernel_params[] = {&params};
if constexpr (kEnableCudaHostAdapter) {
//
@@ -313,6 +323,7 @@ public:
launch_result = cuda_adapter->launch(grid,
cluster,
fallback_cluster,
block,
smem_size,
stream,
@@ -338,6 +349,20 @@ public:
grid, cluster, block, smem_size, stream, kernel, kernel_params);
}
}
else {
if constexpr (ConvKernel::ArchTag::kMinComputeCapability == 100 ||
ConvKernel::ArchTag::kMinComputeCapability == 101) {
launch_result = ClusterLauncher::launch_with_fallback_cluster(
grid,
cluster,
fallback_cluster,
block,
smem_size,
stream,
kernel,
kernel_params);
}
}
}
}
else {
include/cutlass/detail/blockwise_scale_layout.hpp
+15 -4
View File
@@ -48,7 +48,7 @@ namespace cutlass::detail{
using namespace cute;
template<int SFVecSizeM, int SFVecSizeN, int SFVecSizeK, UMMA::Major majorSFA = UMMA::Major::MN, UMMA::Major majorSFB = UMMA::Major::MN>
struct Sm100BlockwiseScaleConfig {
struct Sm1xxBlockwiseScaleConfig {
using ShapeSFA = Shape<Shape<Int<SFVecSizeM>, int32_t>, Shape<Int<SFVecSizeK>, int32_t>, int32_t>;
using ShapeSFB = Shape<Shape<Int<SFVecSizeN>, int32_t>, Shape<Int<SFVecSizeK>, int32_t>, int32_t>;
@@ -271,7 +271,18 @@ struct RuntimeBlockwiseScaleConfig {
// Sm90 only supports MN major for SFA and SFB for now
template<int SFVecSizeM, int SFVecSizeN, int SFVecSizeK>
using Sm90BlockwiseScaleConfig = Sm100BlockwiseScaleConfig<SFVecSizeM, SFVecSizeN, SFVecSizeK>;
using Sm90BlockwiseScaleConfig = Sm1xxBlockwiseScaleConfig<SFVecSizeM, SFVecSizeN, SFVecSizeK>;
template<int SFVecSizeM, int SFVecSizeN, int SFVecSizeK, UMMA::Major majorSFA = UMMA::Major::MN, UMMA::Major majorSFB = UMMA::Major::MN>
using Sm100BlockwiseScaleConfig = Sm1xxBlockwiseScaleConfig<SFVecSizeM, SFVecSizeN, SFVecSizeK, majorSFA, majorSFB>;
template<int SFVecSizeM, int SFVecSizeN, int SFVecSizeK, UMMA::Major majorSFA = UMMA::Major::MN, UMMA::Major majorSFB = UMMA::Major::MN>
using Sm120BlockwiseScaleConfig = Sm1xxBlockwiseScaleConfig<SFVecSizeM, SFVecSizeN, SFVecSizeK, majorSFA, majorSFB>;
template<class MmaTileShape_MNK>
constexpr auto sm90_trivial_blockwise_scale_config(MmaTileShape_MNK) {
return Sm90BlockwiseScaleConfig<size<0>(MmaTileShape_MNK{}), size<1>(MmaTileShape_MNK{}), size<2>(MmaTileShape_MNK{})>{};
}
template<class MmaTileShape_MNK>
constexpr auto sm100_trivial_blockwise_scale_config(MmaTileShape_MNK) {
@@ -279,8 +290,8 @@ constexpr auto sm100_trivial_blockwise_scale_config(MmaTileShape_MNK) {
}
template<class MmaTileShape_MNK>
constexpr auto sm90_trivial_blockwise_scale_config(MmaTileShape_MNK) {
return Sm90BlockwiseScaleConfig<size<0>(MmaTileShape_MNK{}), size<1>(MmaTileShape_MNK{}), size<2>(MmaTileShape_MNK{})>{};
constexpr auto sm120_trivial_blockwise_scale_config(MmaTileShape_MNK) {
return Sm120BlockwiseScaleConfig<size<0>(MmaTileShape_MNK{}), size<1>(MmaTileShape_MNK{}), size<2>(MmaTileShape_MNK{})>{};
}
/////////////////////////////////////////////////////////////////////////////////////////////////
include/cutlass/detail/layout.hpp
+5 -3
View File
@@ -371,11 +371,14 @@ template <
constexpr int
get_input_alignment_bits() {
if constexpr (IsF8F6F4SubBytes && sizeof_bits<ElementType>::value == 4) {
// 16U4 format: The inner tensor size dimension should be multiple of 64B.
return 64 * 8;
}
else if constexpr (IsF8F6F4SubBytes && sizeof_bits<ElementType>::value == 6) {
// 16U6 format : The inner tensor size dimension must be a multiple of 96B.
return 96 * 8;
}
// TMA 16B alignment requirement
return 128;
}
@@ -383,12 +386,11 @@ get_input_alignment_bits() {
template <class ElementType>
constexpr int
get_output_alignment_bits() {
if constexpr (sizeof_bits<ElementType>::value == 6) {
// U6 format : The inner tensor size dimension must be a multiple of 96B.
// 16U6 format : The inner tensor size dimension must be a multiple of 96B.
return 96 * 8;
}
// TMA 16B alignment requirement
return 128;
}
include/cutlass/epilogue/collective/builders/sm100_builder.inl
+3
View File
@@ -981,6 +981,9 @@ private:
static constexpr bool Is2SmMma = is_base_of_v<TmaWarpSpecialized2Sm, Schedule>;
static_assert(Is1SmMma ^ Is2SmMma, "unsupported schedule");
static_assert(not (Is2SmMma && size<0>(ClusterShape_MNK{}) % 2 == 1), "schedule + cluster mismatch");
// C/D should meet TMA alignment requirement if not void
static_assert(detail::is_aligned<ElementC_, AlignmentC, ElementD_, AlignmentD>(),
"C/D Should meet TMA alignment requirement\n");
static constexpr bool DisableDestination = cute::is_void_v<ElementD_>;
using ElementD = cute::conditional_t<DisableDestination,fusion::get_element_aux_t<FusionOpOrCallbacks>,ElementD_>; // prevents void ref breakages
include/cutlass/epilogue/collective/builders/sm90_builder.inl
+3
View File
@@ -293,6 +293,9 @@ template <
class DispatchPolicy
>
struct Sm90TmaBuilderImpl {
// C/D should meet TMA alignment requirement if not void
static_assert(detail::is_aligned<ElementC_, AlignmentC, ElementD_, AlignmentD>(),
"C/D Should meet TMA alignment requirement\n");
// Passing void D disables destination store + smem allocation
using ElementD = cute::conditional_t<cute::is_void_v<ElementD_>,
fusion::get_element_aux_t<FusionOpOrCallbacks>, ElementD_>;
include/cutlass/epilogue/collective/builders/sm90_common.inl
+8
View File
@@ -91,6 +91,14 @@ sm90_get_smem_load_op_for_source() {
}
}
// C/D should meet TMA alignment requirement if not void
template <class ElementC, int AlignmentC, class ElementD, int AlignmentD>
constexpr bool
is_aligned() {
return (cute::is_void_v<ElementC> || (cute::sizeof_bits_v<ElementC> * AlignmentC) % cutlass::detail::get_output_alignment_bits<ElementC>() == 0) &&
(cute::is_void_v<ElementD> || (cute::sizeof_bits_v<ElementD> * AlignmentD) % cutlass::detail::get_output_alignment_bits<ElementD>() == 0);
}
///////////////////////////////////////////////////////////////////////////////
} // namespace cutlass::epilogue::collective::detail
include/cutlass/epilogue/collective/detail.hpp
+10
View File
@@ -217,6 +217,16 @@ struct IsThreadEpilogueOpWithActivation <ThreadEpilogueOp, cute::enable_if_t<Thr
using type = typename ThreadEpilogueOp::ActivationFn;
};
template <typename ThreadEpilogueOp, typename = void>
struct IsThreadEpilogueOpWithPerChannelScaled {
static constexpr bool value = false;
};
template <typename ThreadEpilogueOp>
struct IsThreadEpilogueOpWithPerChannelScaled <ThreadEpilogueOp, cute::void_t<decltype(ThreadEpilogueOp::IsPerRowScaleSupported)>> {
static constexpr bool value = ThreadEpilogueOp::IsPerRowScaleSupported || ThreadEpilogueOp::IsPerColScaleSupported;
};
template <typename ThreadEpilogueOp, typename = void>
struct IsThreadEpilogueOpWithElementwiseArguments : cute::false_type {};
include/cutlass/epilogue/collective/sm100_epilogue_nosmem.hpp
+17 -25
View File
@@ -57,7 +57,7 @@ struct IsDefaultFusionOp {
};
template<
class ElementD, class ElementCompute,
class ElementD, class ElementCompute,
class ElementC, FloatRoundStyle RoundStyle
>
struct IsDefaultFusionOp<
@@ -69,7 +69,7 @@ struct IsDefaultFusionOp<
template<
class ElementOutput, int Count, class ElementAccumulator,
class ElementCompute, epilogue::thread::ScaleType::Kind Scale,
class ElementCompute, epilogue::thread::ScaleType::Kind Scale,
FloatRoundStyle Round, class ElementSource
>
struct IsDefaultFusionOp<
@@ -133,7 +133,7 @@ public:
constexpr static int ThreadCount = 128;
constexpr static int kOutputAlignment = ThreadEpilogueOp::kCount;
constexpr static bool isEpilogueBiasSupported = detail::IsThreadEpilogueOpWithBias<ThreadEpilogueOp>::value;
using AlignmentType = typename cute::uint_bit<sizeof_bits<ElementOutput>::value * kOutputAlignment>::type;
constexpr static uint32_t TmaTransactionBytes = 0;
@@ -240,7 +240,7 @@ public:
Tensor tTR_rAcc = make_tensor<ElementAccumulator>(shape(tTR_gD)); // (T2R,T2R_M,T2R_N)
Tensor tTR_rC = make_tensor<GmemElementC>(shape(tTR_gC)); // (T2R,T2R_M,T2R_N)
Tensor coordCD = make_identity_tensor(problem_shape_mnl); // (M,N,L) -> (m,n,l)
Tensor cCD = local_tile(coordCD, cta_tiler, cta_coord_mnl); // (CTA_M,CTA_N) -> (m,n,l)
Tensor tTR_cCD = thread_t2r.partition_D(cCD); // (T2R,T2R_M,T2R_N) -> (m,n,l)
@@ -250,7 +250,7 @@ public:
Tensor tTR_rD_frag = make_tensor<ElementD>(shape(tTR_rAcc));
Tensor tTR_rD_src = recast<Array<ElementD, VD>>(coalesce(tTR_rD_frag));
Tensor tR2G_rD_dst = recast<Array<ElementD, VD>>(coalesce(tTR_gD));
Tensor tTR_cD_mn_frg = tensor<1>(zipped_divide(coalesce(tTR_cCD), mclD.compose(Int<VD>{})));
Tensor tDpD = make_tensor<bool>(shape(tR2G_rD_dst));
@@ -325,7 +325,7 @@ public:
copy_if(tDpD, tTR_rD_src, tR2G_rD_dst);
}
// source is not needed, avoid load
else
else
{
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < size(tTR_rAcc); i++) {
@@ -382,7 +382,7 @@ public:
auto thread_t2r = tiled_t2r.get_slice(threadIdx.x % size(tiled_t2r));
Tensor tTR_gC = thread_t2r.partition_D(gC); // (T2R,T2R_M,T2R_N)
Tensor tTR_gD = thread_t2r.partition_D(gD); // (T2R,T2R_M,T2R_N)
Tensor coordCD = make_identity_tensor(problem_shape_mnl); // (M,N,L) -> (m,n,l)
Tensor cCD = local_tile(coordCD, cta_tiler, cta_coord_mnl); // (CTA_M,CTA_N) -> (m,n,l)
@@ -498,7 +498,7 @@ public:
// Constructor and Data Members
//
CUTLASS_DEVICE
CollectiveEpilogue(Params const& params_, SharedStorage& shared_tensors)
CollectiveEpilogue(Params const& params_, SharedStorage& shared_tensors)
: fusion_callbacks(params_.thread, shared_tensors.thread)
, smem_buffer_ptr(shared_tensors.buffer.data())
, params(params_) {};
@@ -506,7 +506,7 @@ public:
protected:
FusionCallbacks fusion_callbacks;
uint8_t* smem_buffer_ptr;
Params const& params;
Params const& params;
public:
@@ -543,7 +543,7 @@ public:
can_implement(
[[maybe_unused]] ProblemShape const& problem_shape,
[[maybe_unused]] Arguments const& args) {
bool fusion_implementable = FusionCallbacks::can_implement(problem_shape, args.thread);
if (!fusion_implementable) {
CUTLASS_TRACE_HOST(" CAN IMPLEMENT: Problem Size doesn't meet the minimum requirements for FusionCallbacks.\n");
@@ -636,7 +636,7 @@ public:
Tensor tTR_rD_frg = recast<Array<ElementD, FragmentSize>>(coalesce(tTR_rD));
auto cst_args = cutlass::epilogue::fusion::detail::ConsumerStoreArgs{
problem_shape_mnkl,
problem_shape_mnkl,
cta_tile_mnk,
cta_coord_mnkl,
int(0),
@@ -693,20 +693,17 @@ public:
}
Tensor tTR_cCD_mn = tTR_cCD(_,_,_,epi_m,epi_n);
Tensor tTR_pCD_mn = cute::lazy::transform(tTR_cCD_mn, [&] (auto const& c) CUTLASS_LAMBDA_FUNC_INLINE { return elem_less(c, problem_shape_mnl); });
cst_callbacks.begin_loop(epi_m, epi_n);
if constexpr (not cute::is_void_v<ElementC>) {
if (is_C_load_needed) {
using CVecType = uint_bit_t<VC * sizeof_bits_v<ElementC>>;
Tensor tTR_cC_frag = tensor<1>(zipped_divide(coalesce(tTR_cCD_mn), mclC.compose(Int<VC>{})));
auto pred_fn_C = [&] (auto const&... coords) CUTLASS_LAMBDA_FUNC_INLINE {
return elem_less(tTR_cC_frag(coords...), problem_shape_mnl);
};
Tensor tTR_gC_frg = recast<CVecType>(coalesce(tTR_gC(_,_,_,epi_m,epi_n)));
Tensor tTR_rC_frg = recast<CVecType>(coalesce(tCrC));
copy_if(pred_fn_C, tTR_gC_frg, tTR_rC_frg);
Tensor tTR_pC_frg = tensor<1>(zipped_divide(coalesce(tTR_pCD_mn), mclC.compose(Int<VC>{})));
copy_if(tTR_pC_frg, tTR_gC_frg, tTR_rC_frg);
}
}
@@ -717,7 +714,7 @@ public:
Tensor tTR_rAcc_frg = recast<Array<ElementAccumulator, FragmentSize>>(coalesce(tTR_rAcc));
copy(tiled_t2r, tTR_tAcc_mn, tTR_rAcc);
// After the last tmem load, signal that tmem buffer is consumed and empty
if (do_acc_release) {
cutlass::arch::fence_view_async_tmem_load();
@@ -737,16 +734,11 @@ public:
cst_callbacks.end_loop(epi_m, epi_n);
Tensor tTR_cD_frag = tensor<1>(zipped_divide(coalesce(tTR_cCD_mn), mclD.compose(Int<VD>{})));
auto pred_fn_D = [&] (auto const&... coords) CUTLASS_LAMBDA_FUNC_INLINE {
return elem_less(tTR_cD_frag(coords...), problem_shape_mnl);
};
using VecType = uint_bit_t<VD * sizeof_bits_v<ElementD>>;
Tensor tTR_gD_frg = recast<VecType>(coalesce(tTR_gD(_,_,_,epi_m,epi_n)));
Tensor tTR_rD_frg = recast<VecType>(coalesce(tTR_rD));
copy_if(pred_fn_D, tTR_rD_frg, tTR_gD_frg);
Tensor tTR_pD_frg = tensor<1>(zipped_divide(coalesce(tTR_pCD_mn), mclD.compose(Int<VD>{})));
copy_if(tTR_pD_frg, tTR_rD_frg, tTR_gD_frg);
} // for epi_m
} // for epi_n
include/cutlass/epilogue/collective/sm70_epilogue_vectorized.hpp
+6 -6
View File
@@ -59,7 +59,7 @@ template <
>
class Epilogue {
static_assert(cute::is_same_v<EpilogueScheduleType, EpilogueSimtVectorized> ||
cute::is_same_v<EpilogueScheduleType, EpiloguePtrArraySimtVectorized>,
cute::is_same_v<EpilogueScheduleType, EpiloguePtrArraySimtVectorized>,
"Could not find an epilogue specialization.");
};
@@ -141,7 +141,7 @@ public:
ElementScalar const* beta_ptr = nullptr;
ElementBias const* bias_ptr = nullptr;
StrideBias dBias{};
};
};
template<class ThreadEpiOp>
struct ThreadEpilogueOpArguments<
@@ -202,7 +202,7 @@ public:
to_underlying_arguments(
[[maybe_unused]] ProblemShape const& _,
Arguments const& args,
[[maybe_unused]] void* workspace) {
[[maybe_unused]] void* workspace) {
typename ThreadEpilogueOp::Params thread_op_args;
thread_op_args.alpha = args.thread.alpha;
thread_op_args.beta = args.thread.beta;
@@ -317,7 +317,7 @@ public:
Tensor gC = gC_mnl(_,_,m_coord,n_coord,l_coord); // (BLK_M,BLK_N)
Tensor gD = gD_mnl(_,_,m_coord,n_coord,l_coord); // (BLK_M,BLK_N)
Tensor gBias = gBias_mnl(_,_,m_coord,n_coord,l_coord); // (BLK_M,BLK_N)
// Construct a tensor in SMEM that we can partition for rearranging data
SharedStorage& storage = *reinterpret_cast<SharedStorage*>(smem_buf);
Tensor sAcc = make_tensor(make_smem_ptr(storage.smem_epilogue.data()), SmemLayout{}); // (SMEM_M,SMEM_N)
@@ -389,10 +389,10 @@ public:
Tensor tSR_gBias_flt = filter_zeros(tSR_gBias);
Tensor tSR_rBias_flt = filter_zeros(tSR_rBias);
Tensor tSR_cD_flt = filter_zeros(tSR_cD, tSR_gBias.stride());
Tensor tSR_pD_flt = cute::lazy::transform(tSR_cD_flt, [&](auto const& c){ return elem_less(c, take<0,2>(residue_mnk)); });
// Step 0. Copy Bias from GMEM to fragment
auto pred_fn = [&] (auto const&... coords) { return elem_less(tSR_cD_flt(coords...), take<0, 2>(residue_mnk)); };
copy_if(pred_fn, tSR_gBias_flt, tSR_rBias_flt);
copy_if(tSR_pD_flt, tSR_gBias_flt, tSR_rBias_flt);
}
}
include/cutlass/epilogue/fusion/sm90_visitor_compute_tma_warpspecialized.hpp
+10 -10
View File
@@ -560,18 +560,18 @@ struct Sm90TreeVisitor<
Tensor tC_rAux_vec = recast<VecType>(tC_rAux);
Tensor tC_gAux_vec = recast<VecType>(tC_gAux);
Tensor tC_cAux_vec = tensor<1>(zipped_divide(tC_cAux, MCL.compose(Int<V>{})));
auto predicate_fn = [&] (auto&&... coords) CUTLASS_LAMBDA_FUNC_INLINE { return elem_less(tC_cAux_vec(coords...), residue_tC_cAux); };
copy_if(predicate_fn, tC_rAux_vec, tC_gAux_vec);
Tensor tC_pAux_vec = cute::lazy::transform(tC_cAux_vec, [&](auto const& c){ return elem_less(c, residue_tC_cAux); });
copy_if(tC_pAux_vec, tC_rAux_vec, tC_gAux_vec);
}
// sub-byte vectorization, must serialize threads
else {
// Assumes no inter-warp sharing of bytes (most copy layouts should satisfy this)
int lane_idx = canonical_lane_idx();
auto predicate_fn = [&] (auto&&... coords) CUTLASS_LAMBDA_FUNC_INLINE { return elem_less(tC_cAux(coords...), residue_tC_cAux); };
Tensor tC_pAux = cute::lazy::transform(tC_cAux, [&](auto const& c){ return elem_less(c, residue_tC_cAux); });
CUTLASS_PRAGMA_NO_UNROLL
for (int i = 0; i < NumThreadsPerWarp; ++i) {
if (lane_idx == i) {
copy_if(predicate_fn, tC_rAux, tC_gAux);
copy_if(tC_pAux, tC_rAux, tC_gAux);
}
__syncwarp();
}
@@ -719,12 +719,12 @@ struct Sm90AuxLoad<
Tensor tC_gAux_vec = recast<VecType>(tC_gAux);
Tensor tC_rAux_vec = recast<VecType>(tC_rAux);
Tensor tC_cAux_vec = tensor<1>(zipped_divide(tC_cAux, MCL.compose(Int<V>{})));
auto predicate_fn = [&] (auto&&... coords) CUTLASS_LAMBDA_FUNC_INLINE { return elem_less(tC_cAux_vec(coords...), residue_tC_cAux); };
copy_if(predicate_fn, tC_gAux_vec, tC_rAux_vec);
Tensor tC_pAux_vec = cute::lazy::transform(tC_cAux_vec, [&](auto const& c){ return elem_less(c, residue_tC_cAux); });
copy_if(tC_pAux_vec, tC_gAux_vec, tC_rAux_vec);
}
else {
auto predicate_fn = [&] (auto&&... coords) CUTLASS_LAMBDA_FUNC_INLINE { return elem_less(tC_cAux(coords...), residue_tC_cAux); };
copy_if(predicate_fn, tC_gAux, tC_rAux);
Tensor tC_pAux = cute::lazy::transform(tC_cAux, [&](auto const& c){ return elem_less(c, residue_tC_cAux); });
copy_if(tC_pAux, tC_gAux, tC_rAux);
}
}
}
@@ -738,8 +738,8 @@ struct Sm90AuxLoad<
}
}
auto predicate_fn = [&] (auto&&... coords) CUTLASS_LAMBDA_FUNC_INLINE { return elem_less(tC_cAux(_,_,_,epi_m,epi_n)(coords...), residue_tC_cAux); };
copy_if(predicate_fn, tC_gAux(_,_,_,epi_m,epi_n), tC_rAux);
Tensor tC_pAux = cute::lazy::transform(tC_cAux(_,_,_,epi_m,epi_n), [&](auto const& c){ return elem_less(c, residue_tC_cAux); });
copy_if(tC_pAux, tC_gAux(_,_,_,epi_m,epi_n), tC_rAux);
}
}
include/cutlass/epilogue/fusion/sm90_visitor_load_tma_warpspecialized.hpp
+33 -37
View File
@@ -449,7 +449,7 @@ template <
bool EnableNullptr
>
struct Sm90AuxLoad<
0, EpilogueTile, Element, LayoutOrStrideMNL,
0, EpilogueTile, Element, LayoutOrStrideMNL,
SmemLayoutAtom, CopyOpS2R, Alignment, EnableNullptr
> {
using ElementAux = Element;
@@ -496,7 +496,7 @@ struct Sm90AuxLoad<
CUTLASS_HOST_DEVICE
Sm90AuxLoad(Params const& params, SharedStorage const& shared_storage)
: params_ptr(&params) { }
Params const* params_ptr;
CUTLASS_DEVICE bool
@@ -533,7 +533,7 @@ struct Sm90AuxLoad<
tC_cAux(cute::forward<CTensorG2R>(tC_cAux)),
problem_shape_mnl(problem_shape_mnl),
params_ptr(params_ptr) {}
GTensorG2R tC_gAux;
RTensor tC_rAux;
CTensorG2R tC_cAux;
@@ -551,17 +551,13 @@ struct Sm90AuxLoad<
constexpr auto MCL = decltype(max_common_layout(tC_gAux(_,_,_,_0{},_0{}), tC_rAux)){};
constexpr int V = cute::min(Alignment, size(MCL));
Tensor tC_cAux_mn = tC_cAux(_,_,_,epi_m,epi_n);
Tensor tC_cAux_vec = tensor<1>(zipped_divide(coalesce(tC_cAux_mn), MCL.compose(Int<V>{})));
Tensor tC_gAux_vec = recast<Array<Element, V>>(coalesce(tC_gAux(_,_,_,epi_m,epi_n)));
Tensor tC_rAux_vec = recast<Array<Element, V>>(coalesce(tC_rAux));
auto pred_fn = [&] (auto const&... coords) CUTLASS_LAMBDA_FUNC_INLINE {
return elem_less(tC_cAux_vec(coords...), problem_shape_mnl);
};
Tensor tC_cAux_vec = tensor<1>(zipped_divide(coalesce(tC_cAux(_,_,_,epi_m,epi_n)), MCL.compose(Int<V>{})));
Tensor tC_pAux_vec = cute::lazy::transform(tC_cAux_vec, [&](auto const& c){ return elem_less(c, problem_shape_mnl); });
copy_if(pred_fn, tC_gAux_vec, tC_rAux_vec);
copy_if(tC_pAux_vec, tC_gAux_vec, tC_rAux_vec);
}
template <typename ElementAccumulator, int FragmentSize>
@@ -647,7 +643,7 @@ struct Sm90ScalarBroadcast {
can_implement(ProblemShape const& problem_shape, Arguments const& args) {
return true;
}
template <class ProblemShape>
static size_t
get_workspace_size(ProblemShape const& problem_shape, Arguments const& args) {
@@ -674,11 +670,11 @@ struct Sm90ScalarBroadcast {
// This must be called after update_scalar is called
CUTLASS_DEVICE bool
is_zero() const {
if (get<2>(params_ptr->dScalar[0]) == 0) {
if (get<2>(params_ptr->dScalar[0]) == 0) {
// Only 1 batch
return scalar == Element(0);
}
else {
else {
// multiple batch
if (valid_scalar == false) {
// for stridedBatch kernel, if ptr has a valid address, we need to enable the epi_load warps.
@@ -761,7 +757,7 @@ private:
if (params_ptr->scalar_ptrs[0] != nullptr) {
scalar = params_ptr->scalar_ptrs[0][l_offset];
}
}
else {
// batch stride is ignored for nullptr fallback
scalar = params_ptr->scalars[0];
@@ -774,7 +770,7 @@ private:
if (params_ptr->scalar_ptrs[i] != nullptr) {
int rest_l_offset = l_coord * size<2>(params_ptr->dScalar[i]);
scalar = reduction_fn(scalar, params_ptr->scalar_ptrs[i][rest_l_offset]);
}
}
else {
// batch stride is ignored for nullptr fallback
scalar = reduction_fn(scalar, params_ptr->scalars[i]);
@@ -826,7 +822,7 @@ struct Sm90ScalarBroadcastPtrArray {
can_implement(ProblemShape const& problem_shape, Arguments const& args) {
return true;
}
template <class ProblemShape>
static size_t
get_workspace_size(ProblemShape const& problem_shape, Arguments const& args) {
@@ -946,7 +942,7 @@ private:
if (params_ptr->scalar_ptrs[i] != nullptr) {
int rest_l_offset = l_coord * size<2>(params_ptr->dScalar[i]);
scalar = reduction_fn(scalar, params_ptr->scalar_ptrs[i][rest_l_offset]);
}
}
else {
// batch stride is ignored for nullptr fallback
scalar = reduction_fn(scalar, params_ptr->scalars[i]);
@@ -992,7 +988,7 @@ struct Sm90RowBroadcast {
static_assert(is_static_v<decltype(take<0,2>(StrideMNL{}))> || IsDynamicBroadcast); // batch stride can be dynamic or static
static_assert(take<0,2>(StrideMNL{}) == Stride<_0,_1>{} || IsDynamicBroadcast);
struct SharedStorage {
struct SharedStorage {
array_aligned<ElementInput, size<1>(CtaTileShapeMNK{})> smem;
};
@@ -1078,8 +1074,8 @@ struct Sm90RowBroadcast {
struct ConsumerStoreCallbacks : EmptyConsumerStoreCallbacks {
CUTLASS_DEVICE
ConsumerStoreCallbacks(
GS_GTensor tGS_gRow_, GS_STensor tGS_sRow_,
GS_CTensor tGS_cRow_, Tiled_G2S tiled_g2s_,
GS_GTensor tGS_gRow_, GS_STensor tGS_sRow_,
GS_CTensor tGS_cRow_, Tiled_G2S tiled_g2s_,
SR_STensor tSR_sRow_, SR_RTensor tSR_rRow_,
Residue residue_cRow_, Params const& params_)
: tGS_gRow(tGS_gRow_)
@@ -1098,8 +1094,8 @@ struct Sm90RowBroadcast {
Tiled_G2S tiled_G2S;
SR_STensor tSR_sRow; // (CPY,CPY_M,CPY_N,EPI_M,EPI_N)
SR_RTensor tSR_rRow; // (CPY,CPY_M,CPY_N,EPI_M,EPI_N)
SR_RTensor tSR_rRow; // (CPY,CPY_M,CPY_N,EPI_M,EPI_N)
Residue residue_cRow; // (m, n)
Params const& params;
@@ -1113,7 +1109,7 @@ struct Sm90RowBroadcast {
for (int i = 0; i < size(tGS_gRow_flt); ++i) {
if (get<1>(tGS_cRow_flt(i)) >= size<1>(CtaTileShapeMNK{})) {
continue; // OOB of SMEM,
continue; // OOB of SMEM,
}
if (not is_nullptr && elem_less(tGS_cRow_flt(i), residue_cRow)) {
tGS_sRow_flt(i) = tGS_gRow_flt(i); // issue async gmem to smem load
@@ -1201,18 +1197,18 @@ struct Sm90RowBroadcast {
}
Tensor mRow = make_tensor(make_gmem_ptr(ptr_row), make_layout(layout_M,layout_N,layout_L));
Tensor gRow = local_tile(mRow(_,_,l), take<0,2>(args.tile_shape_mnk), make_coord(m, n)); // (CTA_M, CTA_N)
Tensor sRow = make_tensor(make_smem_ptr(smem),
Tensor sRow = make_tensor(make_smem_ptr(smem),
make_shape(size<0>(CtaTileShapeMNK{}), size<1>(CtaTileShapeMNK{})), make_shape(_0{}, _1{})); // (CTA_M, CTA_N)
//// G2S: Gmem to Smem
auto tiled_g2s = make_tiled_copy(Copy_Atom<DefaultCopy, ElementInput>{},
Layout< Shape<_1, ThreadCount>,
Stride<_0, _1>>{},
Layout<_1>{});
Layout< Shape<_1, ThreadCount>,
Stride<_0, _1>>{},
Layout<_1>{});
auto thr_g2s = tiled_g2s.get_slice(args.thread_idx);
Tensor tGS_gRow = thr_g2s.partition_S(gRow);
Tensor tGS_sRow = thr_g2s.partition_D(sRow);
//// G2S: Coord
//// G2S: Coord
Tensor tGS_cRow = thr_g2s.partition_S(args.cD);
//// S2R: Smem to Reg
@@ -1220,11 +1216,11 @@ struct Sm90RowBroadcast {
Tensor tSR_rRow = make_tensor_like<ElementCompute>(take<0,3>(tSR_sRow)); // (CPY,CPY_M,CPY_N)
return ConsumerStoreCallbacks(
tGS_gRow,
tGS_sRow,
tGS_cRow, tiled_g2s,
tSR_sRow,
tSR_rRow,
tGS_gRow,
tGS_sRow,
tGS_cRow, tiled_g2s,
tSR_sRow,
tSR_rRow,
args.residue_cD,
params);
}
@@ -1378,12 +1374,12 @@ struct Sm90ColBroadcast {
Tensor tCgCol_vec = recast<VecType>(coalesce(tCgCol_flt));
Tensor tCrCol_vec = recast<VecType>(coalesce(tCrCol_flt));
Tensor tCcCol_vec = tensor<1>(zipped_divide(tCcCol_flt, MCL.compose(Int<V>{})));
auto pred_fn = [&] (auto const&... coords) CUTLASS_LAMBDA_FUNC_INLINE { return elem_less(tCcCol_vec(coords...), residue_tCcCol); };
copy_if(pred_fn, tCgCol_vec, tCrCol_vec);
Tensor tCpCol_vec = cute::lazy::transform(tCcCol_vec, [&](auto const& c){ return elem_less(c, residue_tCcCol); });
copy_if(tCpCol_vec, tCgCol_vec, tCrCol_vec);
}
else {
auto pred_fn = [&] (auto const&... coords) CUTLASS_LAMBDA_FUNC_INLINE { return elem_less(tCcCol_flt(coords...), residue_tCcCol); };
copy_if(pred_fn, tCgCol_flt, tCrCol_flt);
Tensor tCpCol_flt = cute::lazy::transform(tCcCol_flt, [&](auto const& c){ return elem_less(c, residue_tCcCol); });
copy_if(tCpCol_flt, tCgCol_flt, tCrCol_flt);
}
constexpr int FrgSize = size(tCrCol_flt);
include/cutlass/epilogue/fusion/sm90_visitor_store_tma_warpspecialized.hpp
+3 -7
View File
@@ -412,17 +412,13 @@ struct Sm90AuxStore<
constexpr auto MCL = decltype(max_common_layout(tC_gAux(_,_,_,_0{},_0{}), tC_rAux)){};
constexpr int V = cute::min(Alignment, size(MCL));
Tensor tC_cAux_mn = tC_cAux(_,_,_,epi_m,epi_n);
Tensor tC_cAux_vec = tensor<1>(zipped_divide(coalesce(tC_cAux_mn), MCL.compose(Int<V>{})));
Tensor tC_gAux_vec = recast<Array<Element, V>>(coalesce(tC_gAux(_,_,_,epi_m,epi_n)));
Tensor tC_rAux_vec = recast<Array<Element, V>>(coalesce(tC_rAux));
auto pred_fn = [&] (auto const&... coords) {
return elem_less(tC_cAux_vec(coords...), problem_shape_mnl);
};
Tensor tC_cAux_vec = tensor<1>(zipped_divide(coalesce(tC_cAux(_,_,_,epi_m,epi_n)), MCL.compose(Int<V>{})));
Tensor tC_pAux_vec = cute::lazy::transform(tC_cAux_vec, [&](auto const& c){ return elem_less(c, problem_shape_mnl); });
copy_if(pred_fn, tC_rAux_vec, tC_gAux_vec);
copy_if(tC_pAux_vec, tC_rAux_vec, tC_gAux_vec);
}
};
include/cutlass/epilogue/thread/activation.h
+17
View File
@@ -540,6 +540,23 @@ struct HardSwish<Array<T, N> > {
}
};
template <int N>
struct HardSwish<Array<half_t, N> > {
using T = half_t;
static const bool kIsHeavy = false;
static constexpr float kOneSixth = 0.16666667f;
CUTLASS_HOST_DEVICE
Array<T, N> operator()(Array<T, N> const &value) const {
minimum<Array<T, N> > mn;
maximum<Array<T, N> > mx;
multiplies<Array<T, N> > mul;
plus<Array<T, N> > add;
return mul(mul(mn(mx(add(value, T(3)), T(0)), T(6)), value), T(kOneSixth));
}
};
template <typename T>
using ScaledHardSwish = Scale<HardSwish<T>>;
include/cutlass/epilogue/threadblock/fusion/visitor_load.hpp
+2 -6
View File
@@ -542,12 +542,8 @@ struct VisitorColBroadcast {
}
}
clear(tC_rCol);
Tensor pred = make_tensor<bool>(shape(tC_gCol));
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < size(pred); ++i) {
pred(i) = get<0>(tC_cCol(i)) < m;
}
copy_if(pred, tC_gCol, tC_rCol);
Tensor tC_pCol = cute::lazy::transform(tC_cCol, [&] (auto const& c) { return get<0>(c) < m; });
copy_if(tC_pCol, tC_gCol, tC_rCol);
}
template <class ElementAccumulator, int FragmentSize>
include/cutlass/experimental/distributed/device/dist_gemm_universal_wrapper.hpp
+1 -1
View File
@@ -446,7 +446,7 @@ public:
Status
construct_graph(bool launch_with_pdl) {
#if ((__CUDACC_VER_MAJOR__ >= 12) && (__CUDACC_VER_MINOR__ >= 4))
#if (__CUDACC_VER_MAJOR__ > 12 || (__CUDACC_VER_MAJOR__ == 12 && __CUDACC_VER_MINOR__ >= 6))
Status status = Status::kSuccess;
// Destroy existing graph, if created
include/cutlass/experimental/distributed/device/full_barrier.hpp
+1 -1
View File
@@ -47,7 +47,7 @@ void launch_full_barrier(
cudaStream_t stream,
bool launch_with_pdl) {
#if (__CUDACC_VER_MAJOR__ > 12 || (__CUDACC_VER_MAJOR__ == 12 && __CUDACC_VER_MINOR__ >= 4))
#if (__CUDACC_VER_MAJOR__ > 12 || (__CUDACC_VER_MAJOR__ == 12 && __CUDACC_VER_MINOR__ >= 6))
// Legacy (kernel) launch with PDL
cudaLaunchAttribute attributes[1];
attributes[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
include/cutlass/gemm/collective/builders/sm100_blockscaled_sparse_umma_builder.inl
+1 -1
View File
@@ -268,7 +268,7 @@ struct CollectiveBuilder<
// Calculate SMEM matrix A and B buffers' pipeline stages and the accumulator stages.
static constexpr uint32_t AccumulatorNPerCta = cute::size<1>(TileShape_MNK{});
static constexpr uint32_t AccumulatorPipelineStageCount = (AccumulatorNPerCta == 256) ? 1 : 2;
static constexpr uint32_t SchedulerPipelineStageCount = 1;
static constexpr uint32_t SchedulerPipelineStageCount = 2;
using SmemTileShape = cute::Shape<BlockTileA_M, BlockTileB_N, BlockTileA_K>;
include/cutlass/gemm/collective/builders/sm100_blockscaled_umma_builder.inl
+1 -1
View File
@@ -238,7 +238,7 @@ struct CollectiveBuilder<
static constexpr bool IsArrayOfPointersGemm = cute::is_base_of_v<KernelSchedulePtrArrayBlockScaledGemmSm100, BuilderScheduleTag>;
// Grouped GEMM(where Stride type is Stride*) uses specific static tile scheduler.
static constexpr bool IsGroupGemm = !cute::is_same_v<StrideA, InternalStrideA>;
static constexpr uint32_t SchedulerPipelineStageCount = cute::conditional_return<IsGroupGemm>(8, 1);
static constexpr uint32_t SchedulerPipelineStageCount = cute::conditional_return<IsGroupGemm>(8, 2);
static constexpr uint32_t KernelSmemCarveout = detail::Sm100DenseGemmTmaUmmaCarveout<
ClusterShape_MNK,
include/cutlass/gemm/collective/builders/sm100_pipeline_carveout.inl
+6
View File
@@ -51,6 +51,8 @@ struct Sm100DenseGemmTmaUmmaCarveout {
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
@@ -64,6 +66,7 @@ struct Sm100DenseGemmTmaUmmaCarveout {
CLCPipelineStorage +
LoadOrderBarrierStorage +
TmemDeallocStorage +
CLCThrottlePipelineStorage +
CLCResponseStorage +
TmemBasePtrsStorage +
TensorMapStorage
@@ -80,6 +83,8 @@ struct Sm100SparseGemmTmaUmmaCarveout {
static constexpr auto CLCPipelineStorage = sizeof(typename cutlass::PipelineCLCFetchAsync<SchedulerPipelineStageCount, ClusterShape_MNK>::SharedStorage);
// AccumulatorPipeline = PipelineUmmaAsync
static constexpr auto AccumulatorPipelineStorage = sizeof(typename cutlass::PipelineUmmaAsync<AccumulatorPipelineStageCount>::SharedStorage);
// CLC Throttle pipeline storage
static constexpr auto CLCThrottlePipelineStorage = sizeof(typename cutlass::PipelineAsync<SchedulerPipelineStageCount>::SharedStorage);
// Tmem dealloc
static constexpr auto TmemDeallocStorage = sizeof(cutlass::arch::ClusterBarrier);
@@ -87,6 +92,7 @@ struct Sm100SparseGemmTmaUmmaCarveout {
cutlass::round_up(LoadOrderBarrierStorage, 16) +
cutlass::round_up(CLCPipelineStorage, 16) +
cutlass::round_up(AccumulatorPipelineStorage, 16) +
cutlass::round_up(CLCThrottlePipelineStorage, 16) +
cutlass::round_up(TmemDeallocStorage, 16),
16));
include/cutlass/gemm/collective/builders/sm100_sparse_umma_builder.inl
+1 -1
View File
@@ -371,7 +371,7 @@ struct CollectiveBuilder<
// Calculate SMEM matrix A and B buffers' pipeline stages and the accumulator stages.
static constexpr uint32_t AccumulatorNPerCta = cute::size<1>(TileShape_MNK{});
static constexpr uint32_t AccumulatorPipelineStageCount = AccumulatorNPerCta > 224 ? 1 : 2;
static constexpr uint32_t SchedulerPipelineStageCount = 1;
static constexpr uint32_t SchedulerPipelineStageCount = 2;
using SmemTileShape = cute::Shape<BlockTileA_M, BlockTileB_N, BlockTileA_K>;
include/cutlass/gemm/collective/builders/sm100_umma_builder.inl
+1 -1
View File
@@ -267,7 +267,7 @@ struct CollectiveBuilder<
static constexpr bool IsArrayOfPointersGemm = (cute::is_base_of_v<KernelScheduleSm100PtrArrayDenseGemm, BuilderScheduleTag>);
// Grouped GEMM(where Stride type is Stride*) uses specific static tile scheduler.
static constexpr bool IsGroupGemm = !cute::is_same_v<StrideA, InternalStrideA>;
static constexpr uint32_t SchedulerPipelineStageCount = cute::conditional_return<IsGroupGemm>(8, 1);
static constexpr uint32_t SchedulerPipelineStageCount = cute::conditional_return<IsGroupGemm>(8, 2);
static constexpr uint32_t KernelSmemCarveout = detail::Sm100DenseGemmTmaUmmaCarveout<
ClusterShape_MNK,
include/cutlass/gemm/collective/builders/sm120_blockwise_mma_builder.inl
+264
View File
@@ -0,0 +1,264 @@
/***************************************************************************************************
* Copyright (c) 2025 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-License-Identifier: BSD-3-Clause
*
* Redistribution and use in source and binary forms, with or without
* modification, are permitted provided that the following conditions are met:
*
* 1. Redistributions of source code must retain the above copyright notice, this
* list of conditions and the following disclaimer.
*
* 2. Redistributions in binary form must reproduce the above copyright notice,
* this list of conditions and the following disclaimer in the documentation
* and/or other materials provided with the distribution.
*
* 3. Neither the name of the copyright holder nor the names of its
* contributors may be used to endorse or promote products derived from
* this software without specific prior written permission.
*
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
* AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
* IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
* DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
* DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
* SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
* CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
* OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*
**************************************************************************************************/
#pragma once
#include "cutlass/gemm/collective/builders/sm120_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 ElementB,
class ElementScalar,
class TileShapeMNK,
class ScaleShapeMNK,
class MainloopPipelineStorage,
int stages
>
constexpr int
sm120_compute_stage_count_or_override_blockwise(StageCount<stages> stage_count) {
return stages;
}
// Returns the maximum number of smem tiles that can be used with a given smem capacity.
template <
int CapacityBytes,
class ElementA,
class ElementB,
class ElementScalar,
class TileShapeMNK,
class ScaleShapeMNK,
class MainloopPipelineStorage,
int carveout_bytes
>
constexpr auto
sm120_compute_stage_count_or_override_blockwise(StageCountAutoCarveout<carveout_bytes> stage_count) {
// For F6/F4 sub-bytes, ElementA/B will be passed in as uint8_t
constexpr auto a_bits = cute::sizeof_bits_v<ElementA>;
constexpr auto b_bits = cute::sizeof_bits_v<ElementB>;
constexpr auto scale_bits = cute::sizeof_bits_v<ElementScalar>;
constexpr auto mainloop_pipeline_bytes = sizeof(MainloopPipelineStorage);
constexpr int stage_bytes =
cutlass::bits_to_bytes(a_bits * size<0>(TileShapeMNK{}) * size<2>(TileShapeMNK{})) +
cutlass::bits_to_bytes(b_bits * size<1>(TileShapeMNK{}) * size<2>(TileShapeMNK{})) +
cutlass::bits_to_bytes(scale_bits * size<0>(ScaleShapeMNK{}) * size<2>(ScaleShapeMNK{})) +
cutlass::bits_to_bytes(scale_bits * size<1>(ScaleShapeMNK{}) * size<2>(ScaleShapeMNK{})) +
static_cast<int>(mainloop_pipeline_bytes);
return (CapacityBytes - carveout_bytes) / stage_bytes;
}
} // namespace detail
/////////////////////////////////////////////////////////////////////////////////////////////////
template <
class ElementA,
class GmemLayoutATagPair,
int AlignmentA,
class ElementB,
class GmemLayoutBTagPair,
int AlignmentB,
class ElementAccumulator,
class TileShape_MNK,
class ClusterShape_MNK,
class StageCountType,
class BuilderScheduleTag
>
struct CollectiveBuilder<
arch::Sm120,
arch::OpClassTensorOp,
ElementA,
GmemLayoutATagPair,
AlignmentA,
ElementB,
GmemLayoutBTagPair,
AlignmentB,
ElementAccumulator,
TileShape_MNK,
ClusterShape_MNK,
StageCountType,
BuilderScheduleTag,
cute::enable_if_t<
not cute::is_tuple_v<ElementA> && not cute::is_tuple_v<ElementB> &&
not cute::is_complex_v<ElementA> && not cute::is_complex_v<ElementB> &&
cute::is_tuple_v<GmemLayoutATagPair> && cute::is_tuple_v<GmemLayoutBTagPair> &&
(cute::is_base_of_v<KernelScheduleSm120Blockwise, BuilderScheduleTag> ||
cute::is_same_v<KernelScheduleAuto, BuilderScheduleTag>) &&
detail::sm1xx_gemm_is_aligned<ElementA, AlignmentA, ElementB, AlignmentB, BuilderScheduleTag>()>>
{
static_assert(detail::is_sm10x_f8f6f4_element<ElementA>() && detail::is_sm10x_f8f6f4_element<ElementB>(),
"SM120 TmaWarpSpecialized blockwise scaling builder currently only supports F8F6F4 MMA.");
static_assert(cute::is_static_v<TileShape_MNK>, "TileShape has to be static");
static_assert(cute::is_static_v<ClusterShape_MNK>, "Cluster has to be static");
using GmemLayoutATag = cute::remove_cvref_t<decltype(get<0>(GmemLayoutATagPair{}))>;
using GmemLayoutSFATag = cute::remove_cvref_t<decltype(get<1>(GmemLayoutATagPair{}))>;
using GmemLayoutBTag = cute::remove_cvref_t<decltype(get<0>(GmemLayoutBTagPair{}))>;
using GmemLayoutSFBTag = cute::remove_cvref_t<decltype(get<1>(GmemLayoutBTagPair{}))>;
static_assert(cute::depth(cute::remove_pointer_t<GmemLayoutSFATag>{}) == 2 and
cute::depth(cute::remove_pointer_t<GmemLayoutSFBTag>{}) == 2,
"Expect SFA and SFB layout to be depth of two with shape ((SFVecMN, restMN),(SFVecK, restK), L)");
static_assert(size<1, 0>(cute::remove_pointer_t<GmemLayoutSFATag>{}) ==
size<1, 0>(cute::remove_pointer_t<GmemLayoutSFBTag>{}),
"SFA and SFB must have equivalent SF vector sizes along K");
static constexpr cute::UMMA::Major UmmaMajorA = detail::tag_to_umma_major_A<GmemLayoutATag>();
static constexpr cute::UMMA::Major UmmaMajorB = detail::tag_to_umma_major_B<GmemLayoutBTag>();
static_assert((UmmaMajorA == UMMA::Major::K && UmmaMajorB == UMMA::Major::K), "Only TN layout is supported.");
using PermTileM = decltype(cute::min(size<0>(TileShape_MNK{}), _128{}));
using PermTileN = decltype(cute::min(size<1>(TileShape_MNK{}), _32{}));
static constexpr bool IsCooperative = !cute::is_base_of_v<KernelTmaWarpSpecializedPingpong, BuilderScheduleTag>;
using AtomLayoutMNK = cute::conditional_t<IsCooperative,
Layout<Shape<_4,_2,_1>>, Layout<Shape<_2,_2,_1>>>;
// 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>());
static_assert(detail::sm1xx_gemm_check_for_f8f6f4_mix8bit_requirement<ElementAMma, ElementBMma,
TileShape_MNK, ClusterShape_MNK,
GmemLayoutATag, GmemLayoutBTag, false /*IsSparse*/>(),
"TileSize and MNK Major does not met with MMA Mix 8-bit TMA load requirement" );
// Setup TiledMma
using TiledMma = decltype(cute::make_tiled_mma(
cute::rr_op_selector_sm120<ElementA, ElementB, ElementAccumulator>(),
AtomLayoutMNK{},
Tile<PermTileM, PermTileN, _32>{}
));
// DType check
static constexpr bool UseF8f6f4 = detail::is_sm120_f8f6f4<TiledMma, ElementA, ElementB>();
static_assert(UseF8f6f4, "Non-blockscaled collective builder only supports F8F6F4 MMA.\n");
// Element type
using SmemAllocTypeA = cute::conditional_t<UseF8f6f4, uint8_t, typename TiledMma::ValTypeA>;
using SmemAllocTypeB = cute::conditional_t<UseF8f6f4, uint8_t, typename TiledMma::ValTypeB>;
using GmemTiledCopyA = decltype(detail::sm90_cluster_shape_to_tma_atom(shape<1>(ClusterShape_MNK{})));
using GmemTiledCopyB = decltype(detail::sm90_cluster_shape_to_tma_atom(shape<0>(ClusterShape_MNK{})));
using SmemLayoutAtomA = decltype(detail::sm120_rr_smem_selector<SmemAllocTypeA, decltype(size<2>(TileShape_MNK{}))>());
using SmemLayoutAtomB = decltype(detail::sm120_rr_smem_selector<SmemAllocTypeB, decltype(size<2>(TileShape_MNK{}))>());
using StrideA = cutlass::gemm::TagToStrideA_t<GmemLayoutATag>;
using StrideB = cutlass::gemm::TagToStrideB_t<GmemLayoutBTag>;
using StrideSFA = cutlass::gemm::TagToStrideA_t<GmemLayoutSFATag>;
using StrideSFB = cutlass::gemm::TagToStrideB_t<GmemLayoutSFBTag>;
static constexpr int ScaleGranularityM = size<0,0>(cute::remove_pointer_t<GmemLayoutSFATag>{});
static constexpr int ScaleGranularityN = size<0,0>(cute::remove_pointer_t<GmemLayoutSFBTag>{});
static constexpr int ScaleGranularityK = size<1,0>(cute::remove_pointer_t<GmemLayoutSFBTag>{});
static_assert(size<0>(TileShape_MNK{}) % ScaleGranularityM == 0, "Scale Granularity M must evenly divide the tile shape M.");
static_assert(size<1>(TileShape_MNK{}) % ScaleGranularityN == 0, "Scale Granularity N must evenly divide the tile shape N.");
static_assert(size<2>(TileShape_MNK{}) == ScaleGranularityK , "Scale Granularity K must be equal to the tile shape K.");
using BlockTileScale_M = Int<size<0>(TileShape_MNK{}) / ScaleGranularityM>;
using BlockTileScale_N = Int<size<1>(TileShape_MNK{}) / ScaleGranularityN>;
using BlockTileScale_K = Int<size<2>(TileShape_MNK{}) / ScaleGranularityK>;
using ScaleTileShape = cute::Shape<BlockTileScale_M, BlockTileScale_N, BlockTileScale_K>;
// Setup Stages and DispatchPolicy
using MainloopPipelineStorage = typename cutlass::PipelineTmaUmmaAsync<1>::SharedStorage;
static constexpr int PipelineStages = detail::sm120_compute_stage_count_or_override_blockwise<
detail::sm120_smem_capacity_bytes, SmemAllocTypeA,
SmemAllocTypeB, ElementAccumulator,
TileShape_MNK, ScaleTileShape, MainloopPipelineStorage>(StageCountType{});
static constexpr uint32_t SchedulerPipelineStageCount = 2;
static constexpr bool IsGroupedGemmKernel = !cute::is_same_v<cute::remove_pointer_t<StrideA>, StrideA>;
using KernelSchedule = cute::conditional_t<IsGroupedGemmKernel,
// PtrArray
cute::conditional_t<IsCooperative,
KernelPtrArrayTmaWarpSpecializedCooperativeBlockwiseScalingSm120<SchedulerPipelineStageCount>,
KernelPtrArrayTmaWarpSpecializedPingpongBlockwiseScalingSm120<SchedulerPipelineStageCount>>,
// Non-PtrArray
cute::conditional_t<IsCooperative,
KernelTmaWarpSpecializedCooperativeBlockwiseScalingSm120<SchedulerPipelineStageCount>,
KernelTmaWarpSpecializedPingpongBlockwiseScalingSm120<SchedulerPipelineStageCount>>>;
using DispatchPolicy = cute::conditional_t<IsGroupedGemmKernel,
MainloopSm120ArrayTmaWarpSpecializedBlockwiseScaling<PipelineStages,
SchedulerPipelineStageCount,
ClusterShape_MNK,
KernelSchedule>,
MainloopSm120TmaWarpSpecializedBlockwiseScaling<PipelineStages,
SchedulerPipelineStageCount,
ClusterShape_MNK,
KernelSchedule>>;
using SmemCopyAtomA = Copy_Atom<decltype(detail::sm120_rr_smem_copy_selector_A<ElementA, ElementB, UseF8f6f4>()), SmemAllocTypeA>;
using SmemCopyAtomB = Copy_Atom<decltype(detail::sm120_rr_smem_copy_selector_B<ElementA, ElementB, UseF8f6f4>()), SmemAllocTypeB>;
using CollectiveOp = CollectiveMma<
DispatchPolicy,
TileShape_MNK,
ElementA,
cute::tuple<StrideA, StrideSFA>,
ElementB,
cute::tuple<StrideB, StrideSFB>,
TiledMma,
GmemTiledCopyA,
SmemLayoutAtomA,
SmemCopyAtomA,
cute::identity,
GmemTiledCopyB,
SmemLayoutAtomB,
SmemCopyAtomB,
cute::identity
>;
};
} // namespace cutlass::gemm::collective
/////////////////////////////////////////////////////////////////////////////////////////////////
include/cutlass/gemm/collective/builders/sm120_mma_builder.inl
+2
View File
@@ -66,6 +66,8 @@ struct CollectiveBuilder<
StageCountType,
BuilderScheduleTag,
cute::enable_if_t<
not cute::is_tuple_v<ElementA> && not cute::is_tuple_v<ElementB> &&
not cute::is_tuple_v<GmemLayoutATag> && not cute::is_tuple_v<GmemLayoutBTag> &&
// Dense Gemm
(cute::is_base_of_v<KernelScheduleSm120DenseGemm, BuilderScheduleTag> ||
cute::is_base_of_v<KernelTmaWarpSpecializedPingpong, BuilderScheduleTag> ||
include/cutlass/gemm/collective/collective_builder.hpp
+1
View File
@@ -50,6 +50,7 @@
#include "cutlass/gemm/collective/builders/sm120_blockscaled_mma_builder.inl"
#include "cutlass/gemm/collective/builders/sm120_sparse_mma_builder.inl"
#include "cutlass/gemm/collective/builders/sm120_blockscaled_sparse_mma_builder.inl"
#include "cutlass/gemm/collective/builders/sm120_blockwise_mma_builder.inl"
#endif
/////////////////////////////////////////////////////////////////////////////////////////////////
include/cutlass/gemm/collective/collective_mma.hpp
+2
View File
@@ -67,6 +67,8 @@
#include "cutlass/gemm/collective/sm120_blockscaled_mma_array_tma.hpp"
#include "cutlass/gemm/collective/sm120_sparse_mma_tma.hpp"
#include "cutlass/gemm/collective/sm120_blockscaled_sparse_mma_tma.hpp"
#include "cutlass/gemm/collective/sm120_mma_tma_blockwise_scaling.hpp"
#include "cutlass/gemm/collective/sm120_mma_array_tma_blockwise_scaling.hpp"
#endif // !defined(__CUDACC_RTC__)
include/cutlass/gemm/collective/sm100_blockscaled_mma_array_warpspecialized.hpp
+114 -76
View File
@@ -28,10 +28,6 @@
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*
**************************************************************************************************/
#pragma once
#include "cutlass/cutlass.h"
@@ -51,7 +47,6 @@
#include "cute/arch/cluster_sm90.hpp"
#include "cute/atom/mma_atom.hpp"
#include "cute/algorithm/gemm.hpp"
#include "cute/tensor_predicate.hpp"
#include "cute/numeric/arithmetic_tuple.hpp"
/////////////////////////////////////////////////////////////////////////////////////////////////
@@ -169,7 +164,6 @@ struct CollectiveMma<
using InternalStrideB = cute::remove_pointer_t<StrideB>;
static constexpr bool IsRuntimeDataTypeA = cutlass::gemm::collective::detail::is_sm10x_runtime_f8f6f4<ElementA>();
static constexpr bool IsRuntimeDataTypeB = cutlass::gemm::collective::detail::is_sm10x_runtime_f8f6f4<ElementB>();
static_assert((IsRuntimeDataTypeA && IsRuntimeDataTypeB) ||
@@ -210,19 +204,15 @@ struct CollectiveMma<
AtomThrShapeMNK>;
using MainloopPipelineState = typename MainloopPipeline::PipelineState;
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(rank(SmemLayoutAtomA{}) == 2, "SmemLayoutAtom must be rank 2 (M/N, K)");
static_assert((size<0>(TileShape{}) % size<0>(SmemLayoutAtomA{})) == 0, "SmemLayoutAtomA must evenly divide the tile shape.");
static_assert((size<2>(TileShape{}) % size<1>(SmemLayoutAtomA{})) == 0, "SmemLayoutAtomA must evenly divide the 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(rank(SmemLayoutAtomB{}) == 2, "SmemLayoutAtom must be rank 2 (M/N, K)");
static_assert((size<1>(TileShape{}) % size<0>(SmemLayoutAtomB{})) == 0, "SmemLayoutAtomB must evenly divide the tile shape.");
static_assert((size<2>(TileShape{}) % size<1>(SmemLayoutAtomB{})) == 0, "SmemLayoutAtomB must evenly divide the tile shape.");
static_assert(cute::is_void_v<SmemCopyAtomB>,
"SM100 UMMA cannot have a non-void copy atom for smem sourced instructions.");
@@ -275,8 +265,8 @@ struct CollectiveMma<
using SmemAllocTypeA = cute::conditional_t<IsF8F6F4 && cute::sizeof_bits_v<ElementAMma> < 8, uint8_t, ElementAMma>;
using SmemAllocTypeB = cute::conditional_t<IsF8F6F4 && cute::sizeof_bits_v<ElementBMma> < 8, uint8_t, ElementBMma>;
using BitTypeElementA = uint_bit_t<cute::sizeof_bits_v<ElementA>>;
using BitTypeElementB = uint_bit_t<cute::sizeof_bits_v<ElementB>>;
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>;
@@ -308,15 +298,22 @@ struct CollectiveMma<
using TensorMapStorage = typename SharedStorage::TensorMapStorage;
using PipelineStorage = typename SharedStorage::PipelineStorage;
// Only one thread issues the TMA and updates the barriers in a 2SM MMA, adjust bytes accordingly
static constexpr uint32_t SFTransactionBytes =
cutlass::bits_to_bytes(size(AtomThrShapeMNK{}) * cosize(take<0,3>(SmemLayoutSFA{})) * cute::sizeof_bits_v<ElementSF>) +
cutlass::bits_to_bytes(size(AtomThrShapeMNK{}) * cosize(take<0,3>(SmemLayoutSFB{})) * cute::sizeof_bits_v<ElementSF>);
// Only one thread issues the TMA and updates the barriers in a 2SM MMA, adjust bytes accordingly
static constexpr uint32_t ABTmaTransactionBytes =
cutlass::bits_to_bytes(size(AtomThrShapeMNK{}) * cosize(take<0,3>(SmemLayoutA{})) * cute::sizeof_bits_v<ElementA>) +
cutlass::bits_to_bytes(size(AtomThrShapeMNK{}) * cosize(take<0,3>(SmemLayoutB{})) * cute::sizeof_bits_v<ElementB>);
static constexpr uint32_t TmaTransactionBytes = ABTmaTransactionBytes + SFTransactionBytes;
template <class AccTensor, class SfaTensor, class SfbTensor>
struct TmemStorage {
AccTensor accumulators;
SfaTensor tCtSFA;
SfbTensor tCtSFB;
};
// Host side kernel arguments
struct Arguments {
ArrayElementA const** ptr_A{nullptr};
@@ -401,7 +398,11 @@ struct CollectiveMma<
CUTLASS_DEVICE
CollectiveMma(Params const& params, ClusterShape cluster_shape, uint32_t block_rank_in_cluster)
: cluster_shape_(cluster_shape)
, block_rank_in_cluster_(block_rank_in_cluster) {
, block_rank_in_cluster_(block_rank_in_cluster)
, layout_SFA_(params.layout_SFA)
, layout_SFB_(params.layout_SFB)
, runtime_data_type_a_(params.runtime_data_type_a)
, runtime_data_type_b_(params.runtime_data_type_b) {
if constexpr (IsDynamicCluster) {
const bool is_fallback_cluster = (cute::size<0>(cluster_shape_) == params.cluster_shape_fallback.x &&
cute::size<1>(cluster_shape_) == params.cluster_shape_fallback.y);
@@ -613,18 +614,48 @@ struct CollectiveMma<
}
/// Construct A Single Stage's Accumulator Shape
CUTLASS_DEVICE auto
CUTLASS_DEVICE static
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;
}
template <class TmemStorage>
CUTLASS_DEVICE static
auto
slice_accumulator(TmemStorage tmem_storage, int stage) {
return tmem_storage.accumulators(_,_,_,stage);
}
template <class FrgEngine, class FrgLayout>
CUTLASS_DEVICE auto
slice_accumulator(cute::Tensor<FrgEngine, FrgLayout> const& accumulators, int stage) {
return accumulators(_,_,_,stage);
template <class EpilogueTile, bool IsOverlappingAccum = false>
CUTLASS_DEVICE static
auto
init_tmem_tensors(EpilogueTile epi_tile) {
TiledMma tiled_mma;
auto acc_shape = partition_accumulator_shape();
// ((MMA_TILE_M,MMA_TILE_N),MMA_M,MMA_N,ACC_PIPE) where ACC_PIPE=2 so we can double buffer our accumulators for mainloop and epilogue.
Tensor accumulators = cutlass::detail::make_sm100_accumulator<AccumulatorPipelineStageCount, IsOverlappingAccum>(
tiled_mma, acc_shape, EpilogueTile{});
Tensor tCtSFA = make_tensor<typename TiledMma::FrgTypeSFA>(shape(SmemLayoutAtomSFA{}));
Tensor tCtSFB = make_tensor<typename TiledMma::FrgTypeSFB>(shape(SmemLayoutAtomSFB{}));
TmemStorage<decltype(accumulators), decltype(tCtSFA), decltype(tCtSFB)> tmem_storage;
tmem_storage.accumulators = accumulators;
tmem_storage.tCtSFA = tCtSFA;
tmem_storage.tCtSFB = tCtSFB;
return tmem_storage;
}
template <class TmemStorage>
CUTLASS_DEVICE static
void
set_tmem_offsets(TmemStorage& tmem_storage, uint32_t tmem_base_addr) {
tmem_storage.accumulators.data() = tmem_base_addr;
tmem_storage.tCtSFA.data() = tmem_storage.accumulators.data().get() + cutlass::detail::find_tmem_tensor_col_offset(tmem_storage.accumulators);
tmem_storage.tCtSFB.data() = tmem_storage.tCtSFA.data().get() + cutlass::detail::find_tmem_tensor_col_offset(tmem_storage.tCtSFA);
}
/// Set up the data needed by this collective for load.
@@ -693,9 +724,9 @@ struct CollectiveMma<
}
else if constexpr (IsCtaN64) {
Tensor mSFB_tmp = observed_tma_load_sfb_->get_tma_tensor(shape(layout_SFB));
auto new_shape = make_shape(make_shape(shape<0,0>(mSFB_tmp),
auto new_shape = make_shape(make_shape(shape<0,0>(mSFB_tmp),
make_shape(_2{} , shape<0,1>(mSFB_tmp))), shape<1>(mSFB_tmp), shape<2>(mSFB_tmp));
auto new_stride = make_stride(make_stride(stride<0,0>(mSFB_tmp),
auto new_stride = make_stride(make_stride(stride<0,0>(mSFB_tmp),
make_stride(_0{}, stride<0,1>(mSFB_tmp))), stride<1>(mSFB_tmp), stride<2>(mSFB_tmp));
return make_tensor(mSFB_tmp.data(), make_layout(new_shape, new_stride));
}
@@ -707,7 +738,6 @@ struct CollectiveMma<
Tensor gSFA_mkl = local_tile(mSFA_mkl, TileShape{}, make_coord(_,_,_), Step<_1, X,_1>{}); // (TILE_M,TILE_K,m,k,l)
Tensor gSFB_nkl = local_tile(mSFB_nkl, TileShape_SF{}, make_coord(_,_,_), Step< X,_1,_1>{}); // (TILE_N,TILE_K,n,k,l)
// Partition for this CTA
ThrMMA cta_mma = TiledMma{}.get_slice(blockIdx.x % size(typename TiledMma::AtomThrID{}));
@@ -770,17 +800,15 @@ struct CollectiveMma<
}
/// Set up the data needed by this collective for mma compute.
template <class FrgEngine, class FrgLayout>
template <class TmemStorage>
CUTLASS_DEVICE auto
mma_init(
Params const& params,
[[maybe_unused]] cute::Tensor<FrgEngine, FrgLayout> const& accumulators,
TensorStorage& shared_tensors,
uint32_t const tmem_offset) const {
TmemStorage tmem_storage,
TensorStorage& shared_tensors) const {
// Allocate "fragments/descriptors" for A and B matrices
Tensor sA = make_tensor(make_smem_ptr(shared_tensors.smem_A.begin()), SmemLayoutA{}); // (BLK_M,BLK_K,PIPE)
Tensor sB = make_tensor(make_smem_ptr(shared_tensors.smem_B.begin()), SmemLayoutB{}); // (BLK_N,BLK_K,PIPE)
Tensor sA = make_tensor(make_smem_ptr(shared_tensors.smem_A.begin()), SmemLayoutA{}); // (BLK_M,BLK_K,PIPE)
Tensor sB = make_tensor(make_smem_ptr(shared_tensors.smem_B.begin()), SmemLayoutB{}); // (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)
@@ -792,13 +820,8 @@ struct CollectiveMma<
//
// Scale Factor
//
Tensor tCtSFA = make_tensor<typename TiledMma::FrgTypeSFA>(shape(SmemLayoutAtomSFA{}));
// Set tCtSFA and tCtSFB start addresses. Only update the TMEM column address by masking the address with 0x000001FF.
// TMEM allocations for SFA and SFB will always start at DP 0.
tCtSFA.data() = tmem_offset;
Tensor tCtSFB = make_tensor<typename TiledMma::FrgTypeSFB>(shape(SmemLayoutAtomSFB{}));
tCtSFB.data() = tCtSFA.data().get() + cutlass::detail::find_tmem_tensor_col_offset(tCtSFA);
Tensor tCtSFA = tmem_storage.tCtSFA;
Tensor tCtSFB = tmem_storage.tCtSFB;
// Setup smem descriptors for UTCCP
Tensor tCsSFA = make_tensor(make_smem_ptr(shared_tensors.smem_SFA.begin()), SmemLayoutSFA{});
Tensor tCsSFB = make_tensor(make_smem_ptr(shared_tensors.smem_SFB.begin()), SmemLayoutSFB{});
@@ -831,8 +854,10 @@ struct CollectiveMma<
TiledMma tiled_mma;
if constexpr (IsRuntimeDataType) {
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;
// 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(runtime_data_type_a_) & 0b111;
tiled_mma.idesc_.b_format_ = uint8_t(runtime_data_type_b_) & 0b111;
}
return cute::make_tuple(
@@ -997,45 +1022,52 @@ struct CollectiveMma<
//
tiled_mma.accumulate_ = UMMA::ScaleOut::Zero;
if (k_tile_count > 0) { // first iteraion
// WAIT on mainloop_pipe_consumer_state until its data are available
// (phase bit flips from mainloop_pipe_consumer_state.phase() value)
mainloop_pipeline.consumer_wait(mainloop_pipe_consumer_state, barrier_token);
if constexpr (IsOverlappingAccum) {
// first iteration manual unroll for tmem overlap kernel
if (k_tile_count > 0) {
// WAIT on mainloop_pipe_consumer_state until its data are available
// (phase bit flips from mainloop_pipe_consumer_state.phase() value)
mainloop_pipeline.consumer_wait(mainloop_pipe_consumer_state, barrier_token);
// Compute on k_tile
int read_stage = mainloop_pipe_consumer_state.index();
// Save current mainlop pipeline read state
auto curr_mainloop_pipe_consumer_state = mainloop_pipe_consumer_state;
// Compute on k_tile
int read_stage = mainloop_pipe_consumer_state.index();
// Save current mainlop pipeline read state
auto curr_mainloop_pipe_consumer_state = mainloop_pipe_consumer_state;
// Advance mainloop_pipe
++mainloop_pipe_consumer_state;
--k_tile_count;
skip_wait = k_tile_count <= 0;
// Peek at next iteration
barrier_token = mainloop_pipeline.consumer_try_wait(mainloop_pipe_consumer_state, skip_wait);
// Advance mainloop_pipe
++mainloop_pipe_consumer_state;
--k_tile_count;
skip_wait = k_tile_count <= 0;
// Peek at next iteration
barrier_token = mainloop_pipeline.consumer_try_wait(mainloop_pipe_consumer_state, skip_wait);
if (cute::elect_one_sync()) {
copy(tiled_copy_s2t_SFA, thr_tCsSFA_s2t(_,_,_,_,read_stage), thr_tCtSFA_s2t);
copy(tiled_copy_s2t_SFB, thr_tCsSFB_s2t(_,_,_,_,read_stage), thr_tCtSFB_s2t);
}
if (cute::elect_one_sync()) {
copy(tiled_copy_s2t_SFA, thr_tCsSFA_s2t(_,_,_,_,read_stage), thr_tCtSFA_s2t);
copy(tiled_copy_s2t_SFB, thr_tCsSFB_s2t(_,_,_,_,read_stage), thr_tCtSFB_s2t);
}
if constexpr (IsOverlappingAccum) {
// Wait for tmem accumulator buffer to become empty with a flipped phase
accumulator_pipeline.producer_acquire(accumulator_pipe_producer_state);
}
// Unroll the K mode manually so we can set scale C to 1
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.with(tiled_mma.accumulate_,
tCtSFA(_,_,k_block),
tCtSFB_mma(_,_,k_block)),
tCrA(_,_,k_block,read_stage),
tCrB(_,_,k_block,read_stage),
accumulators);
tiled_mma.accumulate_ = UMMA::ScaleOut::One;
// Unroll the K mode manually so we can set scale C to 1
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.with(tiled_mma.accumulate_,
tCtSFA(_,_,k_block),
tCtSFB_mma(_,_,k_block)),
tCrA(_,_,k_block,read_stage),
tCrB(_,_,k_block,read_stage),
accumulators);
tiled_mma.accumulate_ = UMMA::ScaleOut::One;
}
mainloop_pipeline.consumer_release(curr_mainloop_pipe_consumer_state);
}
mainloop_pipeline.consumer_release(curr_mainloop_pipe_consumer_state);
}
else {
// Wait for tmem accumulator buffer to become empty with a flipped phase
accumulator_pipeline.producer_acquire(accumulator_pipe_producer_state);
}
CUTLASS_PRAGMA_NO_UNROLL
@@ -1073,6 +1105,7 @@ struct CollectiveMma<
accumulators);
tiled_mma.accumulate_ = UMMA::ScaleOut::One;
}
mainloop_pipeline.consumer_release(curr_mainloop_pipe_consumer_state);
}
@@ -1273,6 +1306,11 @@ protected:
typename Params::TMA_SFA const* observed_tma_load_sfa_{nullptr};
typename Params::TMA_SFB const* observed_tma_load_sfb_{nullptr};
LayoutSFA layout_SFA_;
LayoutSFB layout_SFB_;
RuntimeDataTypeA runtime_data_type_a_{};
RuntimeDataTypeB runtime_data_type_b_{};
ClusterShape cluster_shape_;
uint32_t block_rank_in_cluster_;
};
include/cutlass/gemm/collective/sm100_blockscaled_mma_warpspecialized.hpp
+3 -4
View File
@@ -47,7 +47,6 @@
#include "cute/arch/cluster_sm90.hpp"
#include "cute/atom/mma_atom.hpp"
#include "cute/algorithm/gemm.hpp"
#include "cute/tensor_predicate.hpp"
#include "cute/numeric/arithmetic_tuple.hpp"
/////////////////////////////////////////////////////////////////////////////////////////////////
@@ -123,7 +122,7 @@ struct CollectiveMma<
"Static cluster shape used: TileShape should be evenly divided by TiledMma");
using CtaShape_MNK = decltype(shape_div(TileShape{}, AtomThrShapeMNK{}));
static_assert(shape<1>(CtaShape_MNK{}) == 192 or shape<1>(CtaShape_MNK{}) == 64 or
static_assert(shape<1>(CtaShape_MNK{}) == 192 or shape<1>(CtaShape_MNK{}) == 64 or
shape<1>(CtaShape_MNK{}) == 128 or shape<1>(CtaShape_MNK{}) == 256,
"Cta N should be one of 64/128/192/256");
@@ -726,9 +725,9 @@ struct CollectiveMma<
}
else if constexpr (IsCtaN64) {
Tensor mSFB_tmp = observed_tma_load_sfb_->get_tma_tensor(shape(layout_SFB_));
auto new_shape = make_shape(make_shape(shape<0,0>(mSFB_tmp),
auto new_shape = make_shape(make_shape(shape<0,0>(mSFB_tmp),
make_shape(_2{} , shape<0,1>(mSFB_tmp))), shape<1>(mSFB_tmp), shape<2>(mSFB_tmp));
auto new_stride = make_stride(make_stride(stride<0,0>(mSFB_tmp),
auto new_stride = make_stride(make_stride(stride<0,0>(mSFB_tmp),
make_stride(_0{}, stride<0,1>(mSFB_tmp))), stride<1>(mSFB_tmp), stride<2>(mSFB_tmp));
return make_tensor(mSFB_tmp.data(), make_layout(new_shape, new_stride));
}
include/cutlass/gemm/collective/sm100_blockscaled_sparse_mma_warpspecialized.hpp
+2 -3
View File
@@ -48,7 +48,6 @@
#include "cute/arch/cluster_sm90.hpp"
#include "cute/atom/mma_atom.hpp"
#include "cute/algorithm/gemm.hpp"
#include "cute/tensor_predicate.hpp"
#include "cute/numeric/arithmetic_tuple.hpp"
/////////////////////////////////////////////////////////////////////////////////////////////////
@@ -911,9 +910,9 @@ struct CollectiveMma<
}
else if constexpr (IsCtaN64) {
Tensor mSFB_tmp = observed_tma_load_sfb_->get_tma_tensor(shape(layout_SFB_));
auto new_shape = make_shape(make_shape(shape<0,0>(mSFB_tmp),
auto new_shape = make_shape(make_shape(shape<0,0>(mSFB_tmp),
make_shape(_2{} , shape<0,1>(mSFB_tmp))), shape<1>(mSFB_tmp), shape<2>(mSFB_tmp));
auto new_stride = make_stride(make_stride(stride<0,0>(mSFB_tmp),
auto new_stride = make_stride(make_stride(stride<0,0>(mSFB_tmp),
make_stride(_0{}, stride<0,1>(mSFB_tmp))), stride<1>(mSFB_tmp), stride<2>(mSFB_tmp));
return make_tensor(mSFB_tmp.data(), make_layout(new_shape, new_stride));
}
include/cutlass/gemm/collective/sm100_mma_array_warpspecialized.hpp
+58 -28
View File
@@ -30,7 +30,6 @@
**************************************************************************************************/
#pragma once
#include "cutlass/cutlass.h"
@@ -43,12 +42,12 @@
#include "cutlass/trace.h"
#include "cutlass/kernel_hardware_info.hpp"
#include "cutlass/cuda_host_adapter.hpp"
#include "cutlass/detail/sm100_tmem_helper.hpp"
#include "cute/algorithm/functional.hpp"
#include "cute/arch/cluster_sm90.hpp"
#include "cute/atom/mma_atom.hpp"
#include "cute/algorithm/gemm.hpp"
#include "cute/tensor_predicate.hpp"
#include "cute/numeric/arithmetic_tuple.hpp"
/////////////////////////////////////////////////////////////////////////////////////////////////
@@ -131,7 +130,7 @@ struct CollectiveMma<
using ElementB = ElementB_;
using ElementBMma = typename TiledMma::ValTypeB;
using StrideB = StrideB_;
using InternalStrideB = cute::remove_pointer_t<StrideB>;
using InternalStrideB = cute::remove_pointer_t<StrideB>;
static constexpr bool IsRuntimeDataTypeA = cutlass::gemm::collective::detail::is_sm10x_runtime_f8f6f4<ElementA>();
@@ -212,8 +211,8 @@ struct CollectiveMma<
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 = uint_bit_t<cute::sizeof_bits_v<ElementA>>;
using BitTypeElementB = uint_bit_t<cute::sizeof_bits_v<ElementB>>;
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>;
@@ -221,6 +220,8 @@ struct CollectiveMma<
using RuntimeDataTypeA = cute::conditional_t<IsRuntimeDataTypeA, cute::UMMA::MXF8F6F4Format, void*>;
using RuntimeDataTypeB = cute::conditional_t<IsRuntimeDataTypeB, cute::UMMA::MXF8F6F4Format, void*>;
static constexpr bool IsGroupedGemmKernel = !cute::is_same_v<InternalStrideA, StrideA>;
struct SharedStorage {
struct TensorStorage : cute::aligned_struct<128, _0> {
cute::ArrayEngine<SmemAllocTypeA, cute::cosize_v<SmemLayoutA>> smem_A;
@@ -246,7 +247,10 @@ struct CollectiveMma<
cutlass::bits_to_bytes(size(AtomThrShapeMNK{}) * cosize(take<0,3>(SmemLayoutA{})) * cute::sizeof_bits_v<ElementA>) +
cutlass::bits_to_bytes(size(AtomThrShapeMNK{}) * cosize(take<0,3>(SmemLayoutB{})) * cute::sizeof_bits_v<ElementB>);
static constexpr bool IsGroupedGemmKernel = !cute::is_same_v<InternalStrideA, StrideA>;
template <class AccTensor>
struct TmemStorage {
AccTensor accumulators;
};
// Host side kernel arguments
struct Arguments {
@@ -298,9 +302,11 @@ struct CollectiveMma<
CUTLASS_DEVICE
CollectiveMma(Params const& params, ClusterShape cluster_shape, uint32_t block_rank_in_cluster)
: cluster_shape_(cluster_shape)
, block_rank_in_cluster_(block_rank_in_cluster) {
, block_rank_in_cluster_(block_rank_in_cluster)
, runtime_data_type_a_(params.runtime_data_type_a)
, runtime_data_type_b_(params.runtime_data_type_b) {
if constexpr (IsDynamicCluster) {
const bool is_fallback_cluster = (cute::size<0>(cluster_shape_) == params.cluster_shape_fallback.x &&
const bool is_fallback_cluster = (cute::size<0>(cluster_shape_) == params.cluster_shape_fallback.x &&
cute::size<1>(cluster_shape_) == params.cluster_shape_fallback.y);
observed_tma_load_a_ = is_fallback_cluster ? &params.tma_load_a_fallback : &params.tma_load_a;
observed_tma_load_b_ = is_fallback_cluster ? &params.tma_load_b_fallback : &params.tma_load_b;
@@ -357,7 +363,6 @@ struct CollectiveMma<
auto cluster_layout_vmnk = tiled_divide(make_layout(cluster_shape), make_tile(typename TiledMma::AtomThrID{}));
auto cluster_shape_fallback = cutlass::detail::select_cluster_shape(ClusterShape{}, hw_info.cluster_shape_fallback);
auto cluster_layout_vmnk_fallback = tiled_divide(make_layout(cluster_shape_fallback), make_tile(typename TiledMma::AtomThrID{}));
typename Params::TMA_A tma_load_a = make_tma_atom_A_sm100<TmaInternalElementA>(
GmemTiledCopyA{},
tensor_a,
@@ -421,7 +426,7 @@ struct CollectiveMma<
return cutlass::Status::kSuccess;
}
template<class ProblemShape>
template <class ProblemShape>
static bool
can_implement(
ProblemShape problem_shapes,
@@ -450,17 +455,38 @@ struct CollectiveMma<
}
/// Construct A Single Stage's Accumulator Shape
CUTLASS_DEVICE auto
CUTLASS_DEVICE static
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;
return partition_shape_C(TiledMma{}, take<0,2>(TileShape{})); // ((MMA_TILE_M,MMA_TILE_N),MMA_M,MMA_N)
}
template <class FrgEngine, class FrgLayout>
CUTLASS_DEVICE auto
slice_accumulator(cute::Tensor<FrgEngine, FrgLayout> const& accumulators, int stage) {
return accumulators(_,_,_,stage);
template <class TmemStorage>
CUTLASS_DEVICE static
auto
slice_accumulator(TmemStorage tmem_storage, int stage) {
return tmem_storage.accumulators(_,_,_,stage);
}
template <class EpilogueTile, bool IsOverlappingAccum = false>
CUTLASS_DEVICE static
auto
init_tmem_tensors(EpilogueTile epi_tile) {
TiledMma tiled_mma;
auto acc_shape = partition_accumulator_shape();
// ((MMA_TILE_M,MMA_TILE_N),MMA_M,MMA_N,ACC_PIPE) where ACC_PIPE=2 so we can double buffer our accumulators for mainloop and epilogue.
Tensor accumulators = cutlass::detail::make_sm100_accumulator<AccumulatorPipelineStageCount, IsOverlappingAccum>(
tiled_mma, acc_shape, EpilogueTile{});
TmemStorage<decltype(accumulators)> tmem_storage;
tmem_storage.accumulators = accumulators;
return tmem_storage;
}
template <class TmemStorage>
CUTLASS_DEVICE static
void
set_tmem_offsets(TmemStorage& tmem_storage, uint32_t tmem_base_addr) {
tmem_storage.accumulators.data() = tmem_base_addr;
}
/// Set up the data needed by this collective for load.
@@ -535,13 +561,13 @@ struct CollectiveMma<
}
/// Set up the data needed by this collective for mma compute.
template <class FrgEngine, class FrgLayout>
template <class TmemStorage>
CUTLASS_DEVICE auto
mma_init(
Params const& params,
[[maybe_unused]] cute::Tensor<FrgEngine, FrgLayout> const& accumulators,
TensorStorage& shared_tensors,
[[maybe_unused]] uint32_t const tmem_nonaccum_offset) const {
[[maybe_unused]] TmemStorage tmem_storage,
TensorStorage& shared_tensors) const {
// Allocate "fragments/descriptors" for A and B matrices
Tensor sA = make_tensor(make_smem_ptr(shared_tensors.smem_A.begin()), SmemLayoutA{}); // (BLK_M,BLK_K,PIPE)
Tensor sB = make_tensor(make_smem_ptr(shared_tensors.smem_B.begin()), SmemLayoutB{}); // (BLK_N,BLK_K,PIPE)
@@ -550,15 +576,15 @@ struct CollectiveMma<
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));
CUTE_STATIC_ASSERT_V(Int<DispatchPolicy::Stages>{} == size<3>(sB)); // PIPE
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;
tiled_mma.idesc_.a_format_ = uint8_t(runtime_data_type_a_) & 0b111;
tiled_mma.idesc_.b_format_ = uint8_t(runtime_data_type_b_) & 0b111;
}
return cute::make_tuple(tiled_mma, tCrA, tCrB);
@@ -672,6 +698,8 @@ struct CollectiveMma<
// PIPELINED MAIN LOOP
//
tiled_mma.accumulate_ = UMMA::ScaleOut::Zero;
// Wait for tmem accumulator buffer to become empty with a flipped phase
accumulator_pipeline.producer_acquire(accumulator_pipe_producer_state);
CUTLASS_PRAGMA_NO_UNROLL
while (k_tile_count > 0) {
@@ -776,9 +804,9 @@ struct CollectiveMma<
TmaInternalElementB const* ptr_B = nullptr;
Tensor tensor_b = make_tensor(ptr_B, make_shape(N,K,Int<1>{}), mainloop_params.dB[next_group]);
cute::detail::fill_tma_gmem_shape_stride(*observed_tma_load_a_, tensor_a,
cute::detail::fill_tma_gmem_shape_stride(*observed_tma_load_a_, tensor_a,
prob_shape_A, prob_stride_A);
cute::detail::fill_tma_gmem_shape_stride(*observed_tma_load_b_, tensor_b,
cute::detail::fill_tma_gmem_shape_stride(*observed_tma_load_b_, tensor_b,
prob_shape_B, prob_stride_B);
// Convert strides to byte strides
@@ -852,6 +880,8 @@ protected:
typename Params::TMA_A const* observed_tma_load_a_{nullptr};
typename Params::TMA_B const* observed_tma_load_b_{nullptr};
RuntimeDataTypeA runtime_data_type_a_{};
RuntimeDataTypeB runtime_data_type_b_{};
ClusterShape cluster_shape_;
uint32_t block_rank_in_cluster_;
include/cutlass/gemm/collective/sm100_mma_array_warpspecialized_blockwise_scaling.hpp
+49 -50
View File
@@ -45,7 +45,6 @@
#include "cute/arch/cluster_sm90.hpp"
#include "cute/atom/mma_atom.hpp"
#include "cute/algorithm/gemm.hpp"
#include "cute/tensor_predicate.hpp"
#include "cute/numeric/arithmetic_tuple.hpp"
/////////////////////////////////////////////////////////////////////////////////////////////////
@@ -131,7 +130,7 @@ struct CollectiveMma<
using ElementBMma = typename TiledMma::ValTypeB;
using StrideB = cute::remove_cvref_t<decltype(get<0>(StridePairB_{}))>;
using LayoutSFB = cute::remove_cvref_t<decltype(get<1>(StridePairB_{}))>;
using InternalStrideB = cute::remove_pointer_t<StrideB>;
using InternalStrideB = cute::remove_pointer_t<StrideB>;
using InternalLayoutSFB = cute::remove_pointer_t<LayoutSFB>;
static constexpr bool IsRuntimeDataTypeA = cutlass::gemm::collective::detail::is_sm10x_runtime_f8f6f4<ElementA>();
@@ -143,9 +142,9 @@ struct CollectiveMma<
"ElementA and ElementB should be both runtime or both static.");
static constexpr bool IsRuntimeDataType = IsRuntimeDataTypeA && IsRuntimeDataTypeB;
static constexpr int ScaleGranularityM = size<0,0>(InternalLayoutSFA{});
static constexpr int ScaleMsPerTile = size<0>(TileShape{}) / ScaleGranularityM;
static_assert(size<0>(TileShape{}) % ScaleGranularityM == 0 and ScaleGranularityM <= size<0>(TileShape{}), "Scale Granularity M must divide Tile Shape");
@@ -166,12 +165,12 @@ struct CollectiveMma<
static_assert(size<1>(CtaShape_MNK{}) >= ScaleGranularityN, "Scale Granularity must be smaller than or equal to the tile shape");
static_assert(size<2>(CtaShape_MNK{}) >= ScaleGranularityK, "Scale Granularity must be smaller than or equal to the tile shape");
using ScaleConfig = cutlass::detail::Sm100BlockwiseScaleConfig<ScaleGranularityM,
ScaleGranularityN,
ScaleGranularityK,
using ScaleConfig = cutlass::detail::Sm100BlockwiseScaleConfig<ScaleGranularityM,
ScaleGranularityN,
ScaleGranularityK,
size<0,1>(InternalLayoutSFA{}.stride()) == 1 ? UMMA::Major::MN : UMMA::Major::K,
size<0,1>(InternalLayoutSFB{}.stride()) == 1 ? UMMA::Major::MN : UMMA::Major::K>;
using SmemLayoutAtomSFA = decltype(ScaleConfig::smem_atom_layoutSFA(CtaShape_MNK{}));
using SmemLayoutAtomSFB = decltype(ScaleConfig::smem_atom_layoutSFB(CtaShape_MNK{}));
@@ -193,9 +192,9 @@ struct CollectiveMma<
static constexpr int CopyAlignmentSFA = GmemTiledCopySFA::AtomNumVal::value * sizeof(typename GmemTiledCopySFA::ValType) / sizeof(ElementAccumulator);
static constexpr int CopyAlignmentSFB = GmemTiledCopySFB::AtomNumVal::value * sizeof(typename GmemTiledCopySFB::ValType) / sizeof(ElementAccumulator);
static constexpr int AlignmentSFA = CopyAlignmentSFA * (GmemTiledCopySFA::AtomNumVal::value > 1 ?
static constexpr int AlignmentSFA = CopyAlignmentSFA * (GmemTiledCopySFA::AtomNumVal::value > 1 ?
(size<0,1>(InternalLayoutSFA{}.stride()) == 1 ? ScaleGranularityM : ScaleGranularityK) : 1);
static constexpr int AlignmentSFB = CopyAlignmentSFB * (GmemTiledCopySFB::AtomNumVal::value > 1 ?
static constexpr int AlignmentSFB = CopyAlignmentSFB * (GmemTiledCopySFB::AtomNumVal::value > 1 ?
(size<0,1>(InternalLayoutSFB{}.stride()) == 1 ? ScaleGranularityN : ScaleGranularityK) : 1);
@@ -383,7 +382,7 @@ struct CollectiveMma<
: cluster_shape_(cluster_shape)
, block_rank_in_cluster_(block_rank_in_cluster) {
if constexpr (IsDynamicCluster) {
const bool is_fallback_cluster = (cute::size<0>(cluster_shape_) == params.cluster_shape_fallback.x &&
const bool is_fallback_cluster = (cute::size<0>(cluster_shape_) == params.cluster_shape_fallback.x &&
cute::size<1>(cluster_shape_) == params.cluster_shape_fallback.y);
observed_tma_load_a_ = is_fallback_cluster ? &params.tma_load_a_fallback : &params.tma_load_a;
observed_tma_load_b_ = is_fallback_cluster ? &params.tma_load_b_fallback : &params.tma_load_b;
@@ -666,7 +665,7 @@ struct CollectiveMma<
auto layout_SFA = [&]() CUTLASS_LAMBDA_FUNC_INLINE {
if constexpr (IsGroupedGemmKernel) {
return params.layout_SFA[current_group];
}
}
else {
return params.layout_SFA;
}
@@ -675,7 +674,7 @@ struct CollectiveMma<
auto layout_SFB = [&]() CUTLASS_LAMBDA_FUNC_INLINE {
if constexpr (IsGroupedGemmKernel) {
return params.layout_SFB[current_group];
}
}
else {
return params.layout_SFB;
}
@@ -690,14 +689,14 @@ struct CollectiveMma<
Tensor SFB_nkl_ident = make_identity_tensor(shape(layout_SFB));
// Tile the tensors and defer the slice
Tensor gSFA_mkl = local_tile(mSFA_mkl, CtaShape_MNK{},
Tensor gSFA_mkl = local_tile(mSFA_mkl, CtaShape_MNK{},
make_coord(_,_,_), Step<_1, X,_1>{}); // (BLK_M, BLK_K, m, k, l)
Tensor gSFB_nkl = local_tile(mSFB_nkl, CtaShape_MNK{},
Tensor gSFB_nkl = local_tile(mSFB_nkl, CtaShape_MNK{},
make_coord(_,_,_), Step< X,_1,_1>{}); // (BLK_N, BLK_K, n, k, l)
Tensor identSFA_mkl = local_tile(SFA_mkl_ident, CtaShape_MNK{},
Tensor identSFA_mkl = local_tile(SFA_mkl_ident, CtaShape_MNK{},
make_coord(_,_,_), Step<_1, X,_1>{}); // (BLK_M, BLK_K, m, k, l)
Tensor identSFB_nkl = local_tile(SFB_nkl_ident, CtaShape_MNK{},
Tensor identSFB_nkl = local_tile(SFB_nkl_ident, CtaShape_MNK{},
make_coord(_,_,_), Step< X,_1,_1>{}); // (BLK_N, BLK_K, n, k, l)
static_assert(rank(decltype(gSFA_mkl){}) == 5);
@@ -710,16 +709,16 @@ struct CollectiveMma<
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()),
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()),
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);
@@ -731,16 +730,16 @@ struct CollectiveMma<
tSFAgSFA_mkl, tSFBgSFB_nkl,
tSFAsSFA, tSFBsSFB,
tSFAIdentSFA_mkl, tSFBIdentSFB_nkl,
layout_SFA, layout_SFB);
layout_SFA, layout_SFB);
}
/// Setup data needed for transform
CUTLASS_DEVICE auto
accum_init(
TensorStorage& shared_tensors) const {
Tensor sSFA = make_tensor(make_smem_ptr(shared_tensors.smem_SFA.begin()),
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()),
Tensor sSFB = make_tensor(make_smem_ptr(shared_tensors.smem_SFB.begin()),
SmemLayoutScaleB{}); // (CTA_M,CTA_K,P)
return cute::make_tuple(sSFA, sSFB);
@@ -763,20 +762,20 @@ struct CollectiveMma<
CUTE_STATIC_ASSERT_V(rank(tCrA_) == _4{});
auto mma_tile_shape_A = make_shape(get<0>(shape(tCrA_.layout())),
get<1>(shape(tCrA_.layout())),
Int<K_BLOCK_MMAS_PER_SCALE_K>{},
auto mma_tile_shape_A = make_shape(get<0>(shape(tCrA_.layout())),
get<1>(shape(tCrA_.layout())),
Int<K_BLOCK_MMAS_PER_SCALE_K>{},
_1{});
auto mma_tile_shape_B = make_shape(get<0>(shape(tCrB_.layout())),
get<1>(shape(tCrB_.layout())),
Int<K_BLOCK_MMAS_PER_SCALE_K>{},
auto mma_tile_shape_B = make_shape(get<0>(shape(tCrB_.layout())),
get<1>(shape(tCrB_.layout())),
Int<K_BLOCK_MMAS_PER_SCALE_K>{},
_1{});
Tensor tCrA = flat_divide(tCrA_,
Tensor tCrA = flat_divide(tCrA_,
mma_tile_shape_A)(_,_,_,_0{},_0{},_0{},_,_); // (MMA,MMA_M,MMA_K_PER_SCALE,MMA_K_REST,PIPE)
Tensor tCrB = flat_divide(tCrB_,
Tensor tCrB = flat_divide(tCrB_,
mma_tile_shape_B)(_,_,_,_0{},_0{},_0{},_,_); // (MMA,MMA_N,MMA_K_PER_SCALE,MMA_K_REST,PIPE)
CUTE_STATIC_ASSERT_V(Int<DispatchPolicy::Stages>{} == size<3>(sA)); // PIPE
@@ -884,10 +883,10 @@ struct CollectiveMma<
load_sf(
MainloopSFPipeline mainloop_sf_pipeline,
MainloopSFPipelineState mainloop_sf_pipe_producer_state,
cute::tuple<UnusedGTensorA,
cute::tuple<UnusedGTensorA,
GTensorPartitionedSFA, GTensorPartitionedSFB,
STensorSFA, STensorSFB,
IdentPartitionedSFA,
IdentPartitionedSFA,
IdentPartitionedSFB,
InternalLayoutSFA,
InternalLayoutSFB> const& mainloop_sf_inputs,
@@ -921,19 +920,19 @@ 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());
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);
}
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());
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);
}
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())));
mainloop_sf_pipeline.producer_commit(mainloop_sf_pipe_producer_state, cutlass::arch::cpasync_barrier_arrive_noinc);
mainloop_sf_pipeline.producer_commit(mainloop_sf_pipe_producer_state, cutlass::arch::cpasync_barrier_arrive_noinc);
__syncwarp();
@@ -949,7 +948,7 @@ struct CollectiveMma<
/// Perform a Producer Epilogue to prevent early exit of ctas in a Cluster
CUTLASS_DEVICE void
load_sf_tail(
MainloopSFPipeline mainloop_sf_pipeline,
MainloopSFPipeline mainloop_sf_pipeline,
MainloopSFPipelineState mainloop_sf_pipe_producer_state) {
// Issue the epilogue waits
// This helps avoid early exit of ctas in Cluster
@@ -1050,7 +1049,7 @@ struct CollectiveMma<
class CopyOpT2R,
class EpilogueTile
>
CUTLASS_DEVICE auto
CUTLASS_DEVICE auto
accum(
cute::tuple<AccumulatorPipeline, MainloopSFPipeline> pipelines,
cute::tuple<AccumulatorPipelineState, MainloopSFPipelineState> consumer_states,
@@ -1068,7 +1067,7 @@ struct CollectiveMma<
//
// PIPELINED Transform
//
Tensor acc = slice_accumulator(accumulators, _0{});
Tensor tAcc = acc(make_coord(_,_),_0{},_0{});
Tensor tAcc_epi = flat_divide(tAcc, EpilogueTile{}); // (EPI_TILE_M,EPI_TILE_N,EPI_M,EPI_N)
@@ -1100,15 +1099,15 @@ struct CollectiveMma<
int thread_idx = threadIdx.x % size(tiled_t2r_epi);
ThrCopy thread_t2r_epi = tiled_t2r_epi.get_slice(thread_idx);
ThrCopy thread_t2r_epi = tiled_t2r_epi.get_slice(thread_idx);
Tensor acc_ident_epi = make_identity_tensor(shape(tAcc_epi));
Tensor tTR_rAcc_epi = thread_t2r_epi.partition_D(acc_ident_epi); // (T2R, T2R_M, T2R_N, EPI_M, EPI_N)
Tensor tTR_sSFA_epi = thread_t2r_epi.partition_D(sSFA_epi); // (T2R, T2R_M, T2R_N, EPI_M, EPI_N)
Tensor tTR_sSFB_epi = thread_t2r_epi.partition_D(sSFB_epi); // (T2R, T2R_M, T2R_N, EPI_M, EPI_N)
static_assert(rank(decltype(tTR_sSFA_epi){}) == 7);
Tensor tTR_FullAcc = make_tensor<ElementAccumulator>(shape(tTR_rAcc_epi));
@@ -1137,10 +1136,10 @@ struct CollectiveMma<
CUTE_STATIC_ASSERT_V(cosize(tTR_rSFA_layout) == size(tTR_rSFA_compact));
CUTE_STATIC_ASSERT_V(cosize(tTR_rSFB_layout) == size(tTR_rSFB_compact));
Tensor tTR_rSFA = make_tensor(tTR_rSFA_compact.data(), tTR_rSFA_layout);
Tensor tTR_rSFB = make_tensor(tTR_rSFB_compact.data(), tTR_rSFB_layout);
mainloop_sf_pipeline.consumer_release(mainloop_sf_pipe_state);
++mainloop_sf_pipe_state;
@@ -1166,19 +1165,19 @@ struct CollectiveMma<
// Compute tmem load predication if necessary
copy(tiled_t2r_epi, tTR_tAcc(_,_,_,epi_m,epi_n), tTR_PartAcc);
cutlass::arch::fence_view_async_tmem_load();
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < size(full_acc); ++i) {
ElementAccumulator scale = scale_a(i) * scale_b(i);
full_acc(i) += scale * tTR_PartAcc(i);
}
}
}
}
cutlass::arch::fence_view_async_tmem_load();
accumulator_pipeline.consumer_release(accumulator_pipe_state);
// release acc
++accumulator_pipe_state;
}
}
--k_tile_count;
}
@@ -1255,9 +1254,9 @@ struct CollectiveMma<
TmaInternalElementB const* ptr_B = nullptr;
Tensor tensor_b = make_tensor(ptr_B, make_shape(N,K,Int<1>{}), mainloop_params.dB[next_group]);
cute::detail::fill_tma_gmem_shape_stride(*observed_tma_load_a_, tensor_a,
cute::detail::fill_tma_gmem_shape_stride(*observed_tma_load_a_, tensor_a,
prob_shape_A, prob_stride_A);
cute::detail::fill_tma_gmem_shape_stride(*observed_tma_load_b_, tensor_b,
cute::detail::fill_tma_gmem_shape_stride(*observed_tma_load_b_, tensor_b,
prob_shape_B, prob_stride_B);
// Convert strides to byte strides
include/cutlass/gemm/collective/sm100_mma_array_warpspecialized_emulated.hpp
+10 -11
View File
@@ -48,7 +48,6 @@
#include "cute/atom/mma_atom.hpp"
#include "cute/atom/copy_atom.hpp"
#include "cute/algorithm/gemm.hpp"
#include "cute/tensor_predicate.hpp"
#include "cute/arch/mma_sm100.hpp"
#include "cutlass/trace.h"
#include "cutlass/kernel_hardware_info.hpp"
@@ -138,13 +137,13 @@ struct CollectiveMma<
using ElementA = float;
using PackedElementA = float2;
using StrideA = StrideA_;
using InternalStrideA = cute::remove_pointer_t<StrideA>;
using InternalStrideA = cute::remove_pointer_t<StrideA>;
using ElementAMma = typename TiledMma::ValTypeA;
using PackedElementAMma = uint32_t;
using ElementB = float;
using PackedElementB = float2;
using StrideB = StrideB_;
using InternalStrideB = cute::remove_pointer_t<StrideB>;
using InternalStrideB = cute::remove_pointer_t<StrideB>;
using ElementBMma = typename TiledMma::ValTypeB;
using PackedElementBMma = uint32_t;
using ElementAccumulator = typename TiledMma::ValTypeC;
@@ -308,7 +307,7 @@ struct CollectiveMma<
// Device side kernel params
struct Params {
using ClusterLayout_VMNK = decltype(tiled_divide(make_layout(conditional_return<IsDynamicCluster>(make_shape(uint32_t(0), uint32_t(0), Int<1>{}), ClusterShape{})),
using ClusterLayout_VMNK = decltype(tiled_divide(make_layout(conditional_return<IsDynamicCluster>(make_shape(uint32_t(0), uint32_t(0), Int<1>{}), ClusterShape{})),
make_tile(typename TiledMma::AtomThrID{})));
using TMA_A = decltype(make_tma_atom_A_sm100<ElementA>(
@@ -342,11 +341,11 @@ struct CollectiveMma<
: cluster_shape_(cluster_shape)
, block_rank_in_cluster_(block_rank_in_cluster) {
if constexpr (IsDynamicCluster) {
const bool is_fallback_cluster = (cute::size<0>(cluster_shape_) == params.cluster_shape_fallback.x &&
const bool is_fallback_cluster = (cute::size<0>(cluster_shape_) == params.cluster_shape_fallback.x &&
cute::size<1>(cluster_shape_) == params.cluster_shape_fallback.y);
observed_tma_load_a_ = is_fallback_cluster ? &params.tma_load_a_fallback : &params.tma_load_a;
observed_tma_load_b_ = is_fallback_cluster ? &params.tma_load_b_fallback : &params.tma_load_b;
}
}
else {
observed_tma_load_a_ = &params.tma_load_a;
observed_tma_load_b_ = &params.tma_load_b;
@@ -369,7 +368,7 @@ struct CollectiveMma<
Tensor tensor_a = make_tensor(ptr_A_first_batch, make_layout(make_shape(M,K,mock_L), args.dA));
Tensor tensor_b = make_tensor(ptr_B_first_batch, make_layout(make_shape(N,K,mock_L), args.dB));
auto cluster_shape = cutlass::detail::select_cluster_shape(ClusterShape{}, hw_info.cluster_shape);
// Cluster layout for TMA construction
auto cluster_layout_vmnk = tiled_divide(make_layout(cluster_shape), make_tile(typename TiledMma::AtomThrID{}));
@@ -458,7 +457,7 @@ struct CollectiveMma<
}
/// Construct A Single Stage's Accumulator Shape
CUTLASS_DEVICE auto
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)
@@ -925,7 +924,7 @@ struct CollectiveMma<
CUTLASS_DEVICE auto
accum_init(cute::Tensor<FrgEngine, FrgLayout> const& accumulators, TmemCopyAtom tmem_cp_atom, EpilogueTile epilogue_tile) {
// Obtain a single accumulator
Tensor tAcc = tensor<0>(accumulators(_,_,_,_0{}));
Tensor tAcc = tensor<0>(accumulators(_,_,_,_0{}));
// Apply epilogue subtiling
Tensor tAcc_epi = flat_divide(tAcc, EpilogueTile{}); // (EPI_TILE_M,EPI_TILE_N,EPI_M,EPI_N)
// Create the TMEM copy for single EpilogueTile.
@@ -937,7 +936,7 @@ struct CollectiveMma<
Tensor tTR_rGlobAcc = make_tensor<ElementAccumulator>(shape(tTR_gC)); // (T2R,T2R_M,T2R_N)
Tensor tTR_rAcc_float2 = recast<Array<ElementAccumulator,2>>(tTR_rAcc); // (T2R/2,T2R_M,T2R_N)
Tensor tTR_rGlobAcc_float2 = recast<Array<ElementAccumulator,2>>(tTR_rGlobAcc); // (T2R/2,T2R_M,T2R_N)
// Apply epilogue subtiling to bulk accumulator
// We need to tile the whole bulk_tmem allocation with EpilogueTile.
// The accumulation should be aware of the AccumulatorPipelineStages
@@ -967,7 +966,7 @@ struct CollectiveMma<
uint32_t skip_wait = 0;
auto mma2accum_flag = mma2accum_pipeline.consumer_try_wait(mma2accum_pipeline_consumer_state, skip_wait);
// 1. Global periodic accumulation in registers
CUTLASS_PRAGMA_NO_UNROLL
for (; k_tile_count > 0; --k_tile_count) {
include/cutlass/gemm/collective/sm100_mma_warpspecialized.hpp
-1
View File
@@ -47,7 +47,6 @@
#include "cute/arch/cluster_sm90.hpp"
#include "cute/atom/mma_atom.hpp"
#include "cute/algorithm/gemm.hpp"
#include "cute/tensor_predicate.hpp"
#include "cute/numeric/arithmetic_tuple.hpp"
/////////////////////////////////////////////////////////////////////////////////////////////////
include/cutlass/gemm/collective/sm100_mma_warpspecialized_blockwise_scaling.hpp
+48 -49
View File
@@ -47,7 +47,6 @@
#include "cute/arch/cluster_sm90.hpp"
#include "cute/atom/mma_atom.hpp"
#include "cute/algorithm/gemm.hpp"
#include "cute/tensor_predicate.hpp"
#include "cute/numeric/arithmetic_tuple.hpp"
/////////////////////////////////////////////////////////////////////////////////////////////////
@@ -142,9 +141,9 @@ struct CollectiveMma<
static constexpr int K_BLOCK_MMAS_PER_SCALE_K = ScaleGranularityK / size<2>(typename TiledMma::AtomShape_MNK{});
using ScaleConfig = cutlass::detail::Sm100BlockwiseScaleConfig<ScaleGranularityM,
ScaleGranularityN,
ScaleGranularityK,
using ScaleConfig = cutlass::detail::Sm100BlockwiseScaleConfig<ScaleGranularityM,
ScaleGranularityN,
ScaleGranularityK,
size<0,1>(LayoutSFA{}.stride()) == 1 ? UMMA::Major::MN : UMMA::Major::K,
size<0,1>(LayoutSFB{}.stride()) == 1 ? UMMA::Major::MN : UMMA::Major::K>;
@@ -204,9 +203,9 @@ struct CollectiveMma<
static constexpr int CopyAlignmentSFA = GmemTiledCopySFA::AtomNumVal::value * sizeof(typename GmemTiledCopySFA::ValType) / sizeof(ElementAccumulator);
static constexpr int CopyAlignmentSFB = GmemTiledCopySFB::AtomNumVal::value * sizeof(typename GmemTiledCopySFB::ValType) / sizeof(ElementAccumulator);
static constexpr int AlignmentSFA = CopyAlignmentSFA * (GmemTiledCopySFA::AtomNumVal::value > 1 ?
static constexpr int AlignmentSFA = CopyAlignmentSFA * (GmemTiledCopySFA::AtomNumVal::value > 1 ?
(size<0,1>(LayoutSFA{}.stride()) == 1 ? ScaleGranularityM : ScaleGranularityK) : 1);
static constexpr int AlignmentSFB = CopyAlignmentSFB * (GmemTiledCopySFB::AtomNumVal::value > 1 ?
static constexpr int AlignmentSFB = CopyAlignmentSFB * (GmemTiledCopySFB::AtomNumVal::value > 1 ?
(size<0,1>(LayoutSFB{}.stride()) == 1 ? ScaleGranularityN : ScaleGranularityK) : 1);
@@ -399,7 +398,7 @@ struct CollectiveMma<
>
struct AccumTransformParams {
// for scheduler
STensorScaleA sSFA;
STensorScaleB sSFB;
@@ -468,7 +467,7 @@ struct CollectiveMma<
, runtime_data_type_a_(params.runtime_data_type_a)
, runtime_data_type_b_(params.runtime_data_type_b) {
if constexpr (IsDynamicCluster) {
const bool is_fallback_cluster = (cute::size<0>(cluster_shape_) == params.cluster_shape_fallback.x &&
const bool is_fallback_cluster = (cute::size<0>(cluster_shape_) == params.cluster_shape_fallback.x &&
cute::size<1>(cluster_shape_) == params.cluster_shape_fallback.y);
observed_tma_load_a_ = is_fallback_cluster ? &params.tma_load_a_fallback : &params.tma_load_a;
observed_tma_load_b_ = is_fallback_cluster ? &params.tma_load_b_fallback : &params.tma_load_b;
@@ -567,7 +566,7 @@ struct CollectiveMma<
implementable = implementable && cutlass::detail::check_alignment<min_tma_aligned_elements_A>(cute::make_shape(M,K,L), StrideA{});
constexpr int min_tma_aligned_elements_B = tma_alignment_bits_B / cute::sizeof_bits<ElementB>::value;
implementable = implementable && cutlass::detail::check_alignment<min_tma_aligned_elements_B>(cute::make_shape(N,K,L), StrideB{});
if (!implementable) {
CUTLASS_TRACE_HOST(" CAN IMPLEMENT: Problem Size doesn't meet the minimum alignment requirements for TMA.\n");
}
@@ -627,7 +626,7 @@ struct CollectiveMma<
}
/// Set up the data needed by this collective for load.
/// Return load params containing
/// Return load params containing
/// gA_mkl - The tiled tma tensor for input A
/// gB_nkl - The tiled tma tensor for input B
/// tAsA - partitioned smem tensor for A
@@ -691,12 +690,12 @@ struct CollectiveMma<
}
/// Set up the data needed by this collective for load.
/// Return load params containing
/// Return load params containing
/// tSFAgSFA_mkl - partitioned gmem tensor for SFA
/// tSFBgSFB_nkl - partitioned gmem tensor for SFB
/// tSFAIdentSFA_mkl - partitioned identity tensor for SFA in gmem
/// tSFBIdentSFB_nkl - partitioned identity tensor for SFB in gmem
/// tSFAsSFA - partitioned smem tensor for SFA
/// tSFAsSFA - partitioned smem tensor for SFA
/// tSFBsSFB - partitioned smem tensor for SFB
/// layout_SFA - layout of SFA in gmem
/// layout_SFB - layout of SFB in gmem
@@ -720,14 +719,14 @@ struct CollectiveMma<
Tensor SFB_nkl_ident = make_identity_tensor(shape(mainloop_params.layout_SFB));
// Tile the tensors and defer the slice
Tensor gSFA_mkl = local_tile(mSFA_mkl, CtaShape_MNK{},
Tensor gSFA_mkl = local_tile(mSFA_mkl, CtaShape_MNK{},
make_coord(_,_,_), Step<_1, X,_1>{}); // (BLK_M, BLK_K, m, k, l)
Tensor gSFB_nkl = local_tile(mSFB_nkl, CtaShape_MNK{},
Tensor gSFB_nkl = local_tile(mSFB_nkl, CtaShape_MNK{},
make_coord(_,_,_), Step< X,_1,_1>{}); // (BLK_N, BLK_K, n, k, l)
Tensor identSFA_mkl = local_tile(SFA_mkl_ident, CtaShape_MNK{},
Tensor identSFA_mkl = local_tile(SFA_mkl_ident, CtaShape_MNK{},
make_coord(_,_,_), Step<_1, X,_1>{}); // (BLK_M, BLK_K, m, k, l)
Tensor identSFB_nkl = local_tile(SFB_nkl_ident, CtaShape_MNK{},
Tensor identSFB_nkl = local_tile(SFB_nkl_ident, CtaShape_MNK{},
make_coord(_,_,_), Step< X,_1,_1>{}); // (BLK_N, BLK_K, n, k, l)
static_assert(rank(decltype(gSFA_mkl){}) == 5);
@@ -740,16 +739,16 @@ struct CollectiveMma<
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()),
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()),
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);
@@ -784,20 +783,20 @@ struct CollectiveMma<
CUTE_STATIC_ASSERT_V(rank(tCrA_) == _4{});
auto mma_tile_shape_A = make_shape(get<0>(shape(tCrA_.layout())),
get<1>(shape(tCrA_.layout())),
Int<K_BLOCK_MMAS_PER_SCALE_K>{},
auto mma_tile_shape_A = make_shape(get<0>(shape(tCrA_.layout())),
get<1>(shape(tCrA_.layout())),
Int<K_BLOCK_MMAS_PER_SCALE_K>{},
_1{});
auto mma_tile_shape_B = make_shape(get<0>(shape(tCrB_.layout())),
get<1>(shape(tCrB_.layout())),
Int<K_BLOCK_MMAS_PER_SCALE_K>{},
auto mma_tile_shape_B = make_shape(get<0>(shape(tCrB_.layout())),
get<1>(shape(tCrB_.layout())),
Int<K_BLOCK_MMAS_PER_SCALE_K>{},
_1{});
Tensor tCrA = flat_divide(tCrA_,
Tensor tCrA = flat_divide(tCrA_,
mma_tile_shape_A)(_,_,_,_0{},_0{},_0{},_,_); // (MMA,MMA_M,MMA_K_PER_SCALE,MMA_K_REST,PIPE)
Tensor tCrB = flat_divide(tCrB_,
Tensor tCrB = flat_divide(tCrB_,
mma_tile_shape_B)(_,_,_,_0{},_0{},_0{},_,_); // (MMA,MMA_N,MMA_K_PER_SCALE,MMA_K_REST,PIPE)
@@ -830,9 +829,9 @@ struct CollectiveMma<
// Separate out problem shape for convenience
auto [M,N,K,L] = problem_shape_MNKL;
Tensor sSFA = make_tensor(cute::make_smem_ptr(shared_tensors.smem_SFA.begin()),
Tensor sSFA = make_tensor(cute::make_smem_ptr(shared_tensors.smem_SFA.begin()),
SmemLayoutScaleA{}); // (ScaleMsPerTile,ScakeKsPerTile,P)
Tensor sSFB = make_tensor(cute::make_smem_ptr(shared_tensors.smem_SFB.begin()),
Tensor sSFB = make_tensor(cute::make_smem_ptr(shared_tensors.smem_SFB.begin()),
SmemLayoutScaleB{}); // (ScaleNsPerTile,ScaleKsPerTile,P)
@@ -852,7 +851,7 @@ struct CollectiveMma<
CUTLASS_DEVICE auto
load_ab(
MainloopABPipeline mainloop_pipeline,
MainloopABPipelineState mainloop_pipe_producer_state,
MainloopABPipelineState mainloop_pipe_producer_state,
LoadABParams const& load_inputs,
TileCoordMNKL const& cta_coord_mnkl,
KTileIterator k_tile_iter, int k_tile_count) {
@@ -896,7 +895,7 @@ struct CollectiveMma<
/// Perform a Producer Epilogue to prevent early exit of ctas in a Cluster
CUTLASS_DEVICE void
load_ab_tail(
MainloopABPipeline mainloop_pipeline,
MainloopABPipeline mainloop_pipeline,
MainloopABPipelineState mainloop_pipe_producer_state) {
// Issue the epilogue waits
// This helps avoid early exit of ctas in Cluster
@@ -923,7 +922,7 @@ struct CollectiveMma<
KTileIterator k_tile_iter, int k_tile_count) {
auto [unused_k_tiles,
tSFAgSFA_mkl, tSFBgSFB_nkl,
tSFAgSFA_mkl, tSFBgSFB_nkl,
tSFAIdentSFA_mkl, tSFBIdentSFB_nkl,
tSFAsSFA, tSFBsSFB,
layout_SFA, layout_SFB] = load_inputs;
@@ -950,19 +949,19 @@ 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());
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);
}
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());
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);
}
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())));
mainloop_sf_pipeline.producer_commit(mainloop_sf_pipe_producer_state, cutlass::arch::cpasync_barrier_arrive_noinc);
mainloop_sf_pipeline.producer_commit(mainloop_sf_pipe_producer_state, cutlass::arch::cpasync_barrier_arrive_noinc);
__syncwarp();
@@ -977,7 +976,7 @@ struct CollectiveMma<
/// Perform a Producer Epilogue to prevent early exit of ctas in a Cluster
CUTLASS_DEVICE void
load_sf_tail(
MainloopSFPipeline mainloop_sf_pipeline,
MainloopSFPipeline mainloop_sf_pipeline,
MainloopSFPipelineState mainloop_sf_pipe_producer_state) {
// Issue the epilogue waits
// This helps avoid early exit of ctas in Cluster
@@ -1007,10 +1006,10 @@ struct CollectiveMma<
int k_tile_count) {
auto [tiled_mma, tCrA, tCrB] = mma_inputs;
auto [mainloop_pipeline,
auto [mainloop_pipeline,
accumulator_pipeline] = pipelines;
auto [mainloop_pipe_consumer_state,
auto [mainloop_pipe_consumer_state,
accumulator_pipe_producer_state] = pipeline_states;
uint32_t skip_wait = k_tile_count <= 0;
@@ -1077,7 +1076,7 @@ struct CollectiveMma<
class CopyOpT2R,
class EpilogueTile
>
CUTLASS_DEVICE auto
CUTLASS_DEVICE auto
accum(
cute::tuple<AccumulatorPipeline, MainloopSFPipeline> pipelines,
cute::tuple<AccumulatorPipelineState, MainloopSFPipelineState> consumer_states,
@@ -1095,7 +1094,7 @@ struct CollectiveMma<
//
// PIPELINED Transform
//
Tensor acc = get<0>(slice_accumulator(tmem_storage, _0{}));
Tensor tAcc = acc(make_coord(_,_),_0{},_0{});
@@ -1130,15 +1129,15 @@ struct CollectiveMma<
int thread_idx = threadIdx.x % size(tiled_t2r_epi);
ThrCopy thread_t2r_epi = tiled_t2r_epi.get_slice(thread_idx);
ThrCopy thread_t2r_epi = tiled_t2r_epi.get_slice(thread_idx);
Tensor acc_ident_epi = make_identity_tensor(shape(tAcc_epi));
Tensor tTR_rAcc_epi = thread_t2r_epi.partition_D(acc_ident_epi); // (T2R, T2R_M, T2R_N, EPI_M, EPI_N)
Tensor tTR_sSFA_epi = thread_t2r_epi.partition_D(sSFA_epi); // (T2R, T2R_M, T2R_N, EPI_M, EPI_N)
Tensor tTR_sSFB_epi = thread_t2r_epi.partition_D(sSFB_epi); // (T2R, T2R_M, T2R_N, EPI_M, EPI_N)
static_assert(rank(decltype(tTR_sSFA_epi){}) == 7);
Tensor tTR_FullAcc = make_tensor<ElementAccumulator>(shape(tTR_rAcc_epi));
@@ -1167,10 +1166,10 @@ struct CollectiveMma<
CUTE_STATIC_ASSERT_V(cosize(tTR_rSFA_layout) == size(tTR_rSFA_compact));
CUTE_STATIC_ASSERT_V(cosize(tTR_rSFB_layout) == size(tTR_rSFB_compact));
Tensor tTR_rSFA = make_tensor(tTR_rSFA_compact.data(), tTR_rSFA_layout);
Tensor tTR_rSFB = make_tensor(tTR_rSFB_compact.data(), tTR_rSFB_layout);
mainloop_sf_pipeline.consumer_release(mainloop_sf_pipe_state);
++mainloop_sf_pipe_state;
@@ -1196,19 +1195,19 @@ struct CollectiveMma<
// Compute tmem load predication if necessary
copy(tiled_t2r_epi, tTR_tAcc(_,_,_,epi_m,epi_n), tTR_PartAcc);
cutlass::arch::fence_view_async_tmem_load();
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < size(full_acc); ++i) {
ElementAccumulator scale = scale_a(i) * scale_b(i);
full_acc(i) += scale * tTR_PartAcc(i);
}
}
}
}
cutlass::arch::fence_view_async_tmem_load();
accumulator_pipeline.consumer_release(accumulator_pipe_state);
// release acc
++accumulator_pipe_state;
}
}
--k_tile_count;
}
include/cutlass/gemm/collective/sm100_mma_warpspecialized_emulated.hpp
+8 -9
View File
@@ -48,7 +48,6 @@
#include "cute/atom/mma_atom.hpp"
#include "cute/atom/copy_atom.hpp"
#include "cute/algorithm/gemm.hpp"
#include "cute/tensor_predicate.hpp"
#include "cute/arch/mma_sm100.hpp"
#include "cutlass/trace.h"
#include "cutlass/kernel_hardware_info.hpp"
@@ -313,7 +312,7 @@ struct CollectiveMma<
// Device side kernel params
struct Params {
using ClusterLayout_VMNK = decltype(tiled_divide(make_layout(conditional_return<IsDynamicCluster>(make_shape(uint32_t(0), uint32_t(0), Int<1>{}), ClusterShape{})),
using ClusterLayout_VMNK = decltype(tiled_divide(make_layout(conditional_return<IsDynamicCluster>(make_shape(uint32_t(0), uint32_t(0), Int<1>{}), ClusterShape{})),
make_tile(typename TiledMma::AtomThrID{})));
using TMA_A = decltype(make_tma_atom_A_sm100<ElementA>(
@@ -344,11 +343,11 @@ struct CollectiveMma<
: cluster_shape_(cluster_shape)
, block_rank_in_cluster_(block_rank_in_cluster) {
if constexpr (IsDynamicCluster) {
const bool is_fallback_cluster = (cute::size<0>(cluster_shape_) == params.cluster_shape_fallback.x &&
const bool is_fallback_cluster = (cute::size<0>(cluster_shape_) == params.cluster_shape_fallback.x &&
cute::size<1>(cluster_shape_) == params.cluster_shape_fallback.y);
observed_tma_load_a_ = is_fallback_cluster ? &params.tma_load_a_fallback : &params.tma_load_a;
observed_tma_load_b_ = is_fallback_cluster ? &params.tma_load_b_fallback : &params.tma_load_b;
}
}
else {
observed_tma_load_a_ = &params.tma_load_a;
observed_tma_load_b_ = &params.tma_load_b;
@@ -366,7 +365,7 @@ struct CollectiveMma<
Tensor tensor_a = make_tensor(args.ptr_A, make_layout(make_shape(M,K,L), args.dA));
Tensor tensor_b = make_tensor(args.ptr_B, make_layout(make_shape(N,K,L), args.dB));
auto cluster_shape = cutlass::detail::select_cluster_shape(ClusterShape{}, hw_info.cluster_shape);
// Cluster layout for TMA construction
auto cluster_layout_vmnk = tiled_divide(make_layout(cluster_shape), make_tile(typename TiledMma::AtomThrID{}));
@@ -459,7 +458,7 @@ struct CollectiveMma<
}
/// Construct A Single Stage's Accumulator Shape
CUTLASS_DEVICE auto
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)
@@ -917,7 +916,7 @@ struct CollectiveMma<
CUTLASS_DEVICE auto
accum_init(cute::Tensor<FrgEngine, FrgLayout> const& accumulators, TmemCopyAtom tmem_cp_atom, EpilogueTile epilogue_tile) {
// Obtain a single accumulator
Tensor tAcc = tensor<0>(accumulators(_,_,_,_0{}));
Tensor tAcc = tensor<0>(accumulators(_,_,_,_0{}));
// Apply epilogue subtiling
Tensor tAcc_epi = flat_divide(tAcc, EpilogueTile{}); // (EPI_TILE_M,EPI_TILE_N,EPI_M,EPI_N)
// Create the TMEM copy for single EpilogueTile.
@@ -929,7 +928,7 @@ struct CollectiveMma<
Tensor tTR_rGlobAcc = make_tensor<ElementAccumulator>(shape(tTR_gC)); // (T2R,T2R_M,T2R_N)
Tensor tTR_rAcc_float2 = recast<Array<ElementAccumulator,2>>(tTR_rAcc); // (T2R/2,T2R_M,T2R_N)
Tensor tTR_rGlobAcc_float2 = recast<Array<ElementAccumulator,2>>(tTR_rGlobAcc); // (T2R/2,T2R_M,T2R_N)
// Apply epilogue subtiling to bulk accumulator
// We need to tile the whole bulk_tmem allocation with EpilogueTile.
// The accumulation should be aware of the AccumulatorPipelineStages
@@ -959,7 +958,7 @@ struct CollectiveMma<
uint32_t skip_wait = 0;
auto mma2accum_flag = mma2accum_pipeline.consumer_try_wait(mma2accum_pipeline_consumer_state, skip_wait);
// 1. Global periodic accumulation in registers
CUTLASS_PRAGMA_NO_UNROLL
for (; k_tile_count > 0; --k_tile_count) {
include/cutlass/gemm/collective/sm100_sparse_mma_warpspecialized.hpp
-1
View File
@@ -46,7 +46,6 @@
#include "cute/arch/cluster_sm90.hpp"
#include "cute/atom/mma_atom.hpp"
#include "cute/algorithm/gemm.hpp"
#include "cute/tensor_predicate.hpp"
#include "cute/numeric/arithmetic_tuple.hpp"
/////////////////////////////////////////////////////////////////////////////////////////////////
include/cutlass/gemm/collective/sm120_blockscaled_mma_array_tma.hpp
+8 -9
View File
@@ -45,7 +45,6 @@
#include "cute/atom/mma_atom.hpp"
#include "cute/algorithm/functional.hpp"
#include "cute/algorithm/gemm.hpp"
#include "cute/tensor_predicate.hpp"
#include "cute/numeric/arithmetic_tuple.hpp"
/////////////////////////////////////////////////////////////////////////////////////////////////
@@ -114,7 +113,7 @@ struct CollectiveMma<
using ElementB = remove_cvref_t<decltype(get<0>(ElementPairB{}))>;
using StrideB = remove_cvref_t<decltype(get<0>(StridePairB{}))>;
using InternalStrideB = cute::remove_pointer_t<StrideB>;
// SFA and SFB
using ElementSF = remove_cvref_t<decltype(get<1>(ElementPairA{}))>;
using LayoutSFA = remove_cvref_t<decltype(get<1>(StridePairA{}))>;
@@ -466,7 +465,7 @@ struct CollectiveMma<
constexpr int tma_alignment_bits_B = cutlass::detail::get_input_alignment_bits<ElementB, IsF8F6F4>();
constexpr int min_tma_aligned_elements_A = tma_alignment_bits_A / cutlass::sizeof_bits<ElementA>::value;
constexpr int min_tma_aligned_elements_B = tma_alignment_bits_B / cutlass::sizeof_bits<ElementB>::value;
bool implementable = true;
if (problem_shapes.is_host_problem_shape_available()) {
// Check alignment for all problem sizes
@@ -642,7 +641,7 @@ struct CollectiveMma<
// Represent the full tensors -- get these from TMA
Tensor mA_mkl = params.tma_load_a.get_tma_tensor(make_shape(M,K,init_L)); // (m,k,l)
Tensor mB_nkl = params.tma_load_b.get_tma_tensor(make_shape(N,K,init_L)); // (n,k,l)
// Represent the full tensor of Scale factors
InternalLayoutSFA layout_SFA{};
InternalLayoutSFB layout_SFB{};
@@ -883,7 +882,7 @@ struct CollectiveMma<
auto tCsB_stage = tCsB(_,_,_,read_stage);
auto tCsSFA_stage = tCsSFA(_,_,_,read_stage);
auto tCsSFB_stage = tCsSFB(_,_,_,read_stage);
auto copy_kblock = [&](auto k_block) {
// copy smem->rmem for A/B operand
copy(smem_tiled_copy_A, tCsA_stage(_,_,k_block), tCrA_copy_view(_,_,k_block));
@@ -894,7 +893,7 @@ struct CollectiveMma<
fp4_shift_A(MMAOp{}, tCrA_copy_view(_,_,k_block));
fp4_shift_B(MMAOp{}, tCrB_copy_view(_,_,k_block));
// Copy smem->rmem for SFA/SFB operand
copy(tCsSFA_stage(_,_,k_block), tCrSFA_copy_view(_,_,k_block));
copy(tCsSFB_stage(_,_,k_block), tCrSFB_copy_view(_,_,k_block));
@@ -916,7 +915,7 @@ struct CollectiveMma<
for_each(make_int_sequence<K_BLOCK_MAX>{}, [&] (auto k_block) {
auto k_block_next = ((k_block + 1) == K_BLOCK_MAX) ? 0 : (k_block + 1);
if (k_block == K_BLOCK_MAX - 1) {
cutlass::arch::NamedBarrier::sync(
thr_size(tiled_mma), cutlass::arch::ReservedNamedBarriers::Sm120MainloopBarrier);
@@ -943,7 +942,7 @@ struct CollectiveMma<
for_each(make_int_sequence<K_BLOCK_MAX>{}, [&] (auto k_block) {
auto k_block_next = ((k_block + 1) == K_BLOCK_MAX) ? 0 : (k_block + 1);
if (k_block == K_BLOCK_MAX - 1) {
cutlass::arch::NamedBarrier::sync(
thr_size(tiled_mma), cutlass::arch::ReservedNamedBarriers::Sm120MainloopBarrier);
@@ -1154,7 +1153,7 @@ struct CollectiveMma<
[[maybe_unused]] int32_t next_batch) {
return input_tensors;
}
};
/////////////////////////////////////////////////////////////////////////////////////////////////
include/cutlass/gemm/collective/sm120_blockscaled_mma_tma.hpp
+4 -5
View File
@@ -45,7 +45,6 @@
#include "cute/atom/mma_atom.hpp"
#include "cute/algorithm/functional.hpp"
#include "cute/algorithm/gemm.hpp"
#include "cute/tensor_predicate.hpp"
#include "cute/numeric/arithmetic_tuple.hpp"
/////////////////////////////////////////////////////////////////////////////////////////////////
@@ -798,7 +797,7 @@ struct CollectiveMma<
auto tCsB_stage = tCsB(_,_,_,read_stage);
auto tCsSFA_stage = tCsSFA(_,_,_,read_stage);
auto tCsSFB_stage = tCsSFB(_,_,_,read_stage);
auto copy_kblock = [&](auto k_block) {
// copy smem->rmem for A/B operand
copy(smem_tiled_copy_A, tCsA_stage(_,_,k_block), tCrA_copy_view(_,_,k_block));
@@ -809,7 +808,7 @@ struct CollectiveMma<
fp4_shift_A(MMAOp{}, tCrA_copy_view(_,_,k_block));
fp4_shift_B(MMAOp{}, tCrB_copy_view(_,_,k_block));
// Copy smem->rmem for SFA/SFB operand
copy(tCsSFA_stage(_,_,k_block), tCrSFA_copy_view(_,_,k_block));
copy(tCsSFB_stage(_,_,k_block), tCrSFB_copy_view(_,_,k_block));
@@ -831,7 +830,7 @@ struct CollectiveMma<
for_each(make_int_sequence<K_BLOCK_MAX>{}, [&] (auto k_block) {
auto k_block_next = ((k_block + 1) == K_BLOCK_MAX) ? 0 : (k_block + 1);
if (k_block == K_BLOCK_MAX - 1) {
cutlass::arch::NamedBarrier::sync(
thr_size(tiled_mma), cutlass::arch::ReservedNamedBarriers::Sm120MainloopBarrier);
@@ -858,7 +857,7 @@ struct CollectiveMma<
for_each(make_int_sequence<K_BLOCK_MAX>{}, [&] (auto k_block) {
auto k_block_next = ((k_block + 1) == K_BLOCK_MAX) ? 0 : (k_block + 1);
if (k_block == K_BLOCK_MAX - 1) {
cutlass::arch::NamedBarrier::sync(
thr_size(tiled_mma), cutlass::arch::ReservedNamedBarriers::Sm120MainloopBarrier);
include/cutlass/gemm/collective/sm120_blockscaled_sparse_mma_tma.hpp
+3 -5
View File
@@ -46,7 +46,6 @@
#include "cute/atom/mma_atom.hpp"
#include "cute/algorithm/functional.hpp"
#include "cute/algorithm/gemm.hpp"
#include "cute/tensor_predicate.hpp"
#include "cute/numeric/arithmetic_tuple.hpp"
/////////////////////////////////////////////////////////////////////////////////////////////////
@@ -707,7 +706,7 @@ struct CollectiveMma<
Tensor mSFB_tmp = mainloop_params.tma_load_sfb.get_tma_tensor(shape(mainloop_params.layout_SFB));
auto x = stride<0,1>(mSFB_tmp);
auto y = ceil_div(shape<0,1>(mSFB_tmp), _2{});
auto new_shape = make_shape (make_shape( shape<0,0>(mSFB_tmp),
auto new_shape = make_shape (make_shape( shape<0,0>(mSFB_tmp),
make_shape( make_shape(_2{}), y)), shape<1>(mSFB_tmp), shape<2>(mSFB_tmp));
auto new_stride = make_stride(make_stride(stride<0,0>(mSFB_tmp),
make_stride(make_stride(_0{}), x)), stride<1>(mSFB_tmp), stride<2>(mSFB_tmp));
@@ -964,10 +963,9 @@ struct CollectiveMma<
Tensor tCcE = gmem_thr_copy_E.partition_S(cE_mk); // (CPY,CPY_M,CPY_K)
auto [atom, vec] = get_copy_atom_and_common_vec();
// Coordinate comparison for out of bound (OOB) predication
Tensor tZcE = zipped_divide(tCcE, vec);
auto pred_fn = [&](auto coord){ return cute::elem_less(tZcE(Int<0>{}, coord), Shape_MK); };
Tensor tZpE = cute::lazy::transform(zipped_divide(tCcE, vec), [&](auto const& c){ return cute::elem_less(c, Shape_MK); });
// Copy
cute::copy_if(atom, pred_fn, zipped_divide(tCgE, vec), zipped_divide(tCrE_copy_view, vec));
cute::copy_if(atom, tZpE, zipped_divide(tCgE, vec), zipped_divide(tCrE_copy_view, vec));
}
else {
// Copy
include/cutlass/gemm/collective/sm120_mma_array_tma_blockwise_scaling.hpp
+1001
View File
File diff suppressed because it is too large Load Diff
include/cutlass/gemm/collective/sm120_mma_tma.hpp
+4 -5
View File
@@ -44,7 +44,6 @@
#include "cute/atom/mma_atom.hpp"
#include "cute/algorithm/functional.hpp"
#include "cute/algorithm/gemm.hpp"
#include "cute/tensor_predicate.hpp"
#include "cute/numeric/arithmetic_tuple.hpp"
/////////////////////////////////////////////////////////////////////////////////////////////////
@@ -113,7 +112,7 @@ struct CollectiveMma<
using RuntimeDataTypeA = void*;
using RuntimeDataTypeB = void*;
static constexpr int ThreadCount = size(TiledMma{});
using MainloopPipeline = cutlass::PipelineTmaAsync<DispatchPolicy::Stages>;
@@ -505,7 +504,7 @@ struct CollectiveMma<
int read_stage = smem_pipe_read.index();
auto tCsA_stage = tCsA(_,_,_,read_stage);
auto tCsB_stage = tCsB(_,_,_,read_stage);
auto copy_kblock = [&](auto k_block) {
// copy smem->rmem for A/B operand
copy(smem_tiled_copy_A, tCsA_stage(_,_,k_block), tCrA_copy_view(_,_,k_block));
@@ -533,7 +532,7 @@ struct CollectiveMma<
for_each(make_int_sequence<K_BLOCK_MAX>{}, [&] (auto k_block) {
auto k_block_next = ((k_block + 1) == K_BLOCK_MAX) ? 0 : (k_block + 1);
if (k_block == K_BLOCK_MAX - 1) {
cutlass::arch::NamedBarrier::sync(
thr_size(tiled_mma), cutlass::arch::ReservedNamedBarriers::Sm120MainloopBarrier);
@@ -558,7 +557,7 @@ struct CollectiveMma<
for_each(make_int_sequence<K_BLOCK_MAX>{}, [&] (auto k_block) {
auto k_block_next = ((k_block + 1) == K_BLOCK_MAX) ? 0 : (k_block + 1);
if (k_block == K_BLOCK_MAX - 1) {
cutlass::arch::NamedBarrier::sync(
thr_size(tiled_mma), cutlass::arch::ReservedNamedBarriers::Sm120MainloopBarrier);
include/cutlass/gemm/collective/sm120_mma_tma_blockwise_scaling.hpp
+779
View File
@@ -0,0 +1,779 @@
/***************************************************************************************************
* Copyright (c) 2025 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-License-Identifier: BSD-3-Clause
*
* Redistribution and use in source and binary forms, with or without
* modification, are permitted provided that the following conditions are met:
*
* 1. Redistributions of source code must retain the above copyright notice, this
* list of conditions and the following disclaimer.
*
* 2. Redistributions in binary form must reproduce the above copyright notice,
* this list of conditions and the following disclaimer in the documentation
* and/or other materials provided with the distribution.
*
* 3. Neither the name of the copyright holder nor the names of its
* contributors may be used to endorse or promote products derived from
* this software without specific prior written permission.
*
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
* AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
* IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
* DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
* DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
* SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
* CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
* OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*
**************************************************************************************************/
#pragma once
#include "cutlass/cutlass.h"
#include "cutlass/gemm/gemm.h"
#include "cutlass/pipeline/pipeline.hpp"
#include "cutlass/gemm/dispatch_policy.hpp"
#include "cutlass/detail/dependent_false.hpp"
#include "cutlass/trace.h"
#include "cutlass/numeric_types.h"
#include "cute/arch/cluster_sm90.hpp"
#include "cute/arch/copy_sm90.hpp"
#include "cute/atom/mma_atom.hpp"
#include "cute/algorithm/functional.hpp"
#include "cute/algorithm/gemm.hpp"
#include "cute/numeric/arithmetic_tuple.hpp"
/////////////////////////////////////////////////////////////////////////////////////////////////
namespace cutlass::gemm::collective {
using namespace cute;
/////////////////////////////////////////////////////////////////////////////////////////////////
template <
int Stages,
int SchedulerPipelineStageCount,
class ClusterShape,
class KernelScheduleType,
class TileShape_,
class ElementA_,
class StridePairA_,
class ElementB_,
class StridePairB_,
class TiledMma_,
class GmemTiledCopyA_,
class SmemLayoutAtomA_,
class SmemCopyAtomA_,
class TransformA_,
class GmemTiledCopyB_,
class SmemLayoutAtomB_,
class SmemCopyAtomB_,
class TransformB_>
struct CollectiveMma<
MainloopSm120TmaWarpSpecializedBlockwiseScaling<Stages, SchedulerPipelineStageCount, ClusterShape, KernelScheduleType>,
TileShape_,
ElementA_,
StridePairA_,
ElementB_,
StridePairB_,
TiledMma_,
GmemTiledCopyA_,
SmemLayoutAtomA_,
SmemCopyAtomA_,
TransformA_,
GmemTiledCopyB_,
SmemLayoutAtomB_,
SmemCopyAtomB_,
TransformB_> {
//
// Type Aliases
//
using DispatchPolicy = MainloopSm120TmaWarpSpecializedBlockwiseScaling<Stages, SchedulerPipelineStageCount, ClusterShape, KernelScheduleType>;
using TileShape = TileShape_;
using ElementA = ElementA_;
using StrideA = cute::remove_cvref_t<decltype(get<0>(StridePairA_{}))>;
using LayoutSFA = cute::remove_cvref_t<decltype(get<1>(StridePairA_{}))>;
using ElementB = ElementB_;
using StrideB = cute::remove_cvref_t<decltype(get<0>(StridePairB_{}))>;
using LayoutSFB = cute::remove_cvref_t<decltype(get<1>(StridePairB_{}))>;
using TiledMma = TiledMma_;
using CtaShape_MNK = decltype(shape_div(TileShape{}, ClusterShape{}));
using ElementAccumulator = typename TiledMma::ValTypeC;
using ElementSF = ElementAccumulator;
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 RuntimeDataTypeA = void*;
using RuntimeDataTypeB = void*;
static constexpr int ThreadCount = size(TiledMma{});
using MainloopPipeline = cutlass::PipelineTmaAsync<DispatchPolicy::Stages>;
using PipelineParams = typename MainloopPipeline::Params;
using PipelineState = typename cutlass::PipelineState<DispatchPolicy::Stages>;
// One threads per CTA are producers (1 for operand tile)
static constexpr int NumProducerThreadEvents = 33;
static constexpr int ScaleGranularityM = size<0,0>(LayoutSFA{});
static constexpr int ScaleGranularityN = size<0,0>(LayoutSFB{});
static constexpr int ScaleGranularityK = size<1,0>(LayoutSFB{});
static_assert(size<1, 0>(LayoutSFA{}) == size<1, 0>(LayoutSFB{}), "Vector size K must be equal for SFA and SFB");
static_assert(size<0>(TileShape{}) % ScaleGranularityM == 0, "Scale Granularity M must evenly divide the tile shape M.");
static_assert(size<1>(TileShape{}) % ScaleGranularityN == 0, "Scale Granularity N must evenly divide the tile shape N.");
static_assert(size<2>(TileShape{}) == ScaleGranularityK , "Scale Granularity K must be equal to the tile shape K.");
static constexpr int ScaleMsPerTile = size<0>(TileShape{}) / ScaleGranularityM;
static constexpr int ScaleNsPerTile = size<1>(TileShape{}) / ScaleGranularityN;
using ScaleConfig = cutlass::detail::Sm120BlockwiseScaleConfig<ScaleGranularityM,
ScaleGranularityN,
ScaleGranularityK,
size<0,1>(LayoutSFA{}.stride()) == 1 ? UMMA::Major::MN : UMMA::Major::K,
size<0,1>(LayoutSFB{}.stride()) == 1 ? UMMA::Major::MN : UMMA::Major::K>;
static constexpr int AlignmentSFA = 1;
static constexpr int AlignmentSFB = 1;
static_assert(rank(SmemLayoutAtomA{}) == 2, "SmemLayoutAtom must be rank 2 (M/N, K)");
static_assert((size<0>(TileShape{}) % size<0>(SmemLayoutAtomA{})) == 0, "SmemLayoutAtom must evenly divide tile shape.");
static_assert((size<2>(TileShape{}) % size<1>(SmemLayoutAtomA{})) == 0, "SmemLayoutAtom must evenly divide tile shape.");
static_assert(rank(SmemLayoutAtomB{}) == 2, "SmemLayoutAtom must be rank 2 (M/N, K)");
static_assert((size<1>(TileShape{}) % size<0>(SmemLayoutAtomB{})) == 0, "SmemLayoutAtom must evenly divide tile shape.");
static_assert((size<2>(TileShape{}) % size<1>(SmemLayoutAtomB{})) == 0, "SmemLayoutAtom must evenly divide tile shape.");
static_assert(not cute::is_void_v<SmemCopyAtomA>,
"SM120 mainloop must specify a copy atom for A operand smem->rmem reads.");
static_assert(not cute::is_void_v<SmemCopyAtomB>,
"SM120 mainloop must specify a copy atom for B operand smem->rmem reads.");
// Tile along modes in a way that maximizes the TMA box size.
using SmemLayoutA = decltype(tile_to_shape(
SmemLayoutAtomA{},
make_shape(shape<0>(TileShape{}), shape<2>(TileShape{}), Int<DispatchPolicy::Stages>{}),
conditional_t< ::cutlass::gemm::detail::is_major<0,StrideA>(), Step<_2,_1,_3>, Step<_1,_2,_3>>{}));
using SmemLayoutB = decltype(tile_to_shape(
SmemLayoutAtomB{},
make_shape(shape<1>(TileShape{}), shape<2>(TileShape{}), Int<DispatchPolicy::Stages>{}),
conditional_t< ::cutlass::gemm::detail::is_major<0,StrideB>(), Step<_2,_1,_3>, Step<_1,_2,_3>>{}));
// Block scaling gmem-to-smem copy atom
// we can have partial tiles in M or N, so don't vectorize those loads
using SmemBlockScalingCopyAtomA = Copy_Atom<SM80_CP_ASYNC_CACHEALWAYS<ElementSF>, ElementSF>;
using SmemBlockScalingCopyAtomB = Copy_Atom<SM80_CP_ASYNC_CACHEALWAYS<ElementSF>, ElementSF>;
// Block scaling smem layout
using SmemLayoutScaleA = Layout<Shape<Int<ScaleMsPerTile>, Int<DispatchPolicy::Stages>>>;
using SmemLayoutScaleB = Layout<Shape<Int<ScaleNsPerTile>, Int<DispatchPolicy::Stages>>>;
static_assert(rank(SmemLayoutA{}) == 3, "Smem layout must be rank 3.");
static_assert(rank(SmemLayoutB{}) == 3, "Smem layout must be rank 3.");
static_assert(DispatchPolicy::Stages >= 2, "Specialization requires Stages set to value 2 or more.");
static_assert(not cute::is_base_of<cute::GMMA::DescriptorIterator, typename TiledMma::FrgTypeA>::value &&
not cute::is_base_of<cute::GMMA::DescriptorIterator, typename TiledMma::FrgTypeB>::value,
"MMA atom must source both A and B operands from rmem for this mainloop.");
static_assert(cute::is_same_v<GmemTiledCopyA, SM90_TMA_LOAD> || cute::is_same_v<GmemTiledCopyA, SM90_TMA_LOAD_MULTICAST>,
"GmemTiledCopy - invalid SM90 TMA copy atom specified.");
static_assert(cute::is_same_v<GmemTiledCopyB, SM90_TMA_LOAD> || cute::is_same_v<GmemTiledCopyB, SM90_TMA_LOAD_MULTICAST>,
"GmemTiledCopy - invalid SM90 TMA copy atom specified.");
static constexpr bool IsF8F6F4 = detail::is_sm120_f8f6f4<TiledMma, ElementA, ElementB>();
// TMA converts f32 input to tf32 when copying from GMEM to SMEM
// For all other types, cast to size equivalent uint type to avoid any rounding by TMA.
using TmaInternalElementA = cute::conditional_t<cute::is_same_v<ElementA, float>,
cutlass::tfloat32_t,
cute::conditional_t<cute::is_same_v<ElementA, cutlass::float_e2m1_t>,
cutlass::detail::float_e2m1_unpacksmem_t,
cute::conditional_t<cute::is_same_v<ElementA, cutlass::float_e2m3_t>,
cutlass::detail::float_e2m3_unpacksmem_t,
cute::conditional_t<cute::is_same_v<ElementA, cutlass::float_e3m2_t>,
cutlass::detail::float_e3m2_unpacksmem_t,
uint_bit_t<sizeof_bits_v<ElementA>>>>>>;
using TmaInternalElementB = cute::conditional_t<cute::is_same_v<ElementB, float>,
cutlass::tfloat32_t,
cute::conditional_t<cute::is_same_v<ElementB, cutlass::float_e2m1_t>,
cutlass::detail::float_e2m1_unpacksmem_t,
cute::conditional_t<cute::is_same_v<ElementB, cutlass::float_e2m3_t>,
cutlass::detail::float_e2m3_unpacksmem_t,
cute::conditional_t<cute::is_same_v<ElementB, cutlass::float_e3m2_t>,
cutlass::detail::float_e3m2_unpacksmem_t,
uint_bit_t<sizeof_bits_v<ElementB>>>>>>;
using SmemAllocTypeA = cute::conditional_t<IsF8F6F4, uint8_t, typename TiledMma::ValTypeA>;
using SmemAllocTypeB = cute::conditional_t<IsF8F6F4, uint8_t, typename TiledMma::ValTypeB>;
// Set the bytes transferred in this TMA transaction (may involve multiple issues)
static constexpr uint32_t TmaTransactionBytesMK = static_cast<uint32_t>(
cutlass::bits_to_bytes(size(take<0,2>(SmemLayoutA{})) * sizeof_bits<ElementA>::value));
static constexpr uint32_t TmaTransactionBytesNK = static_cast<uint32_t>(
cutlass::bits_to_bytes(size(take<0,2>(SmemLayoutB{})) * sizeof_bits<ElementB>::value));
static constexpr uint32_t TmaTransactionBytes = TmaTransactionBytesMK + TmaTransactionBytesNK;
struct SharedStorage {
struct TensorStorage : cute::aligned_struct<128, _0> {
alignas(1024) cute::array_aligned<SmemAllocTypeA, cute::cosize_v<SmemLayoutA>> smem_A;
alignas(1024) cute::array_aligned<SmemAllocTypeB, cute::cosize_v<SmemLayoutB>> smem_B;
cute::array_aligned<ElementSF, cute::cosize_v<SmemLayoutScaleA>> smem_scale_A;
cute::array_aligned<ElementSF, cute::cosize_v<SmemLayoutScaleB>> smem_scale_B;
} tensors;
using PipelineStorage = typename MainloopPipeline::SharedStorage;
alignas(16) PipelineStorage pipeline_storage;
};
using TensorStorage = typename SharedStorage::TensorStorage;
using PipelineStorage = typename SharedStorage::PipelineStorage;
// Host side kernel arguments
struct Arguments {
ElementA const* ptr_A{nullptr};
StrideA dA{};
ElementB const* ptr_B{nullptr};
StrideB dB{};
ElementAccumulator const* ptr_SFA{nullptr};
LayoutSFA layout_SFA{};
ElementAccumulator const* ptr_SFB{nullptr};
LayoutSFB layout_SFB{};
};
// Device side kernel params
struct Params {
// Assumption: StrideA is congruent with Problem_MK
using TMA_A = decltype(make_tma_copy(
GmemTiledCopyA{},
make_tensor(recast_ptr<TmaInternalElementA>(nullptr), repeat_like(StrideA{}, int32_t(0)), StrideA{}),
SmemLayoutA{}(_,_,0),
make_shape(shape<0>(TileShape{}), shape<2>(TileShape{})),
size<1>(ClusterShape{}))); // mcast along N mode for this M load, if any
// Assumption: StrideB is congruent with Problem_NK
using TMA_B = decltype(make_tma_copy(
GmemTiledCopyB{},
make_tensor(recast_ptr<TmaInternalElementB>(nullptr), repeat_like(StrideB{}, int32_t(0)), StrideB{}),
SmemLayoutB{}(_,_,0),
make_shape(shape<1>(TileShape{}), shape<2>(TileShape{})),
size<0>(ClusterShape{}))); // mcast along M mode for this N load, if any
TMA_A tma_load_a;
TMA_B tma_load_b;
uint32_t tma_transaction_bytes = TmaTransactionBytes;
uint32_t tma_transaction_bytes_mk = TmaTransactionBytesMK;
uint32_t tma_transaction_bytes_nk = TmaTransactionBytesNK;
// Block scaling factors for A and B
ElementSF const* ptr_SFA;
LayoutSFA layout_SFA;
ElementSF const* ptr_SFB;
LayoutSFB layout_SFB;
};
//
// Methods
//
template <class ProblemShape>
static constexpr Params
to_underlying_arguments(ProblemShape const& problem_shape, Arguments const& args, void* workspace) {
(void) workspace;
// 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<TmaInternalElementA>(args.ptr_A);
auto ptr_B = recast_ptr<TmaInternalElementB>(args.ptr_B);
Tensor tensor_a = make_tensor(ptr_A, make_layout(make_shape(M,K,L), args.dA));
Tensor tensor_b = make_tensor(ptr_B, make_layout(make_shape(N,K,L), args.dB));
typename Params::TMA_A tma_load_a = make_tma_copy(
GmemTiledCopyA{},
tensor_a,
SmemLayoutA{}(_,_,cute::Int<0>{}),
make_shape(shape<0>(TileShape{}), shape<2>(TileShape{})),
size<1>(ClusterShape{})); // mcast along N mode for this M load, if any
typename Params::TMA_B tma_load_b = make_tma_copy(
GmemTiledCopyB{},
tensor_b,
SmemLayoutB{}(_,_,cute::Int<0>{}),
make_shape(shape<1>(TileShape{}), shape<2>(TileShape{})),
size<0>(ClusterShape{})); // mcast along M mode for this N load, if any
return {
tma_load_a,
tma_load_b,
TmaTransactionBytes,
TmaTransactionBytesMK,
TmaTransactionBytesNK,
args.ptr_SFA,
args.layout_SFA,
args.ptr_SFB,
args.layout_SFB
};
}
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;
constexpr int tma_alignment_bits_A = cutlass::detail::get_input_alignment_bits<ElementA, IsF8F6F4>();
constexpr int tma_alignment_bits_B = cutlass::detail::get_input_alignment_bits<ElementB, IsF8F6F4>();
bool implementable = true;
constexpr int min_tma_aligned_elements_A = tma_alignment_bits_A / cutlass::sizeof_bits<ElementA>::value;
implementable = implementable && cutlass::detail::check_alignment<min_tma_aligned_elements_A>(cute::make_shape(M,K,L), StrideA{});
constexpr int min_tma_aligned_elements_B = tma_alignment_bits_B / cutlass::sizeof_bits<ElementB>::value;
implementable = implementable && cutlass::detail::check_alignment<min_tma_aligned_elements_B>(cute::make_shape(N,K,L), StrideB{});
if (!implementable) {
CUTLASS_TRACE_HOST(" CAN IMPLEMENT: Problem Size doesn't meet the minimum alignment requirements for TMA.\n");
}
// Ensure complete scale blocks
implementable = implementable && (M % ScaleGranularityM == 0);
implementable = implementable && (N % ScaleGranularityN == 0);
// We expect full tiles in K
implementable = implementable && (K % size<2>(TileShape{}) == 0);
if (!implementable) {
CUTLASS_TRACE_HOST(" CAN IMPLEMENT: Problem Size doesn't meet the alignment requirements for blockwise scaling.\n");
}
return implementable;
}
/// Issue Tma Descriptor Prefetch -- ideally from a single thread for best performance
CUTLASS_DEVICE
static void prefetch_tma_descriptors(Params const& mainloop_params) {
cute::prefetch_tma_descriptor(mainloop_params.tma_load_a.get_tma_descriptor());
cute::prefetch_tma_descriptor(mainloop_params.tma_load_b.get_tma_descriptor());
}
/// Set up the data needed by this collective for load and mma.
/// Returns a tuple of tensors. The collective and the kernel layer have the contract
/// Returned tuple must contain at least two elements, with the first two elements being:
/// gA_mkl - The tma tensor, A after a local tile so it has shape (BLK_M,BLK_K,m,k,l)
/// gB_nkl - The tma tensor, B after a local tile so it has shape (BLK_N,BLK_K,n,k,l)
/// The rest of the tensors can be specified as needed by this collective.
template <class ProblemShape_MNKL>
CUTLASS_DEVICE auto
load_init(ProblemShape_MNKL const& problem_shape_MNKL, Params const& mainloop_params) const {
using X = Underscore;
// Separate out problem shape for convenience
auto [M, N, K, L] = problem_shape_MNKL;
// TMA requires special handling of strides to deal with coord codomain mapping
// Represent the full tensors -- get these from TMA
Tensor mA_mkl = mainloop_params.tma_load_a.get_tma_tensor(make_shape(M,K,L)); // (m,k,l)
Tensor mB_nkl = mainloop_params.tma_load_b.get_tma_tensor(make_shape(N,K,L)); // (n,k,l)
// Make tiled views, defer the slice
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)
Tensor mSFA_mkl = make_tensor(make_gmem_ptr(mainloop_params.ptr_SFA), filter(mainloop_params.layout_SFA)); // (Ms, Ks)
Tensor mSFB_nkl = make_tensor(make_gmem_ptr(mainloop_params.ptr_SFB), filter(mainloop_params.layout_SFB)); // (Ns, Ks)
return cute::make_tuple(gA_mkl, gB_nkl, mSFA_mkl, mSFB_nkl);
}
/// Perform a collective-scoped matrix multiply-accumulate
/// Producer Perspective
template <
class TensorA, class TensorB,
class TensorSFA, class TensorSFB,
class KTileIterator, class BlockCoord
>
CUTLASS_DEVICE void
load(
Params const& mainloop_params,
MainloopPipeline pipeline,
PipelineState smem_pipe_write,
cute::tuple<TensorA, TensorB, TensorSFA, TensorSFB> const& load_inputs,
BlockCoord const& blk_coord,
KTileIterator k_tile_iter, int k_tile_count,
int thread_idx,
uint32_t block_rank_in_cluster,
TensorStorage& shared_tensors) {
int lane_predicate = cute::elect_one_sync();
Tensor sA = make_tensor(make_smem_ptr(shared_tensors.smem_A.data()), SmemLayoutA{}); // (BLK_M,BLK_K,PIPE)
Tensor sB = make_tensor(make_smem_ptr(shared_tensors.smem_B.data()), SmemLayoutB{}); // (BLK_N,BLK_K,PIPE)
Tensor sSFA = make_tensor(make_smem_ptr(shared_tensors.smem_scale_A.data()), SmemLayoutScaleA{});
Tensor sSFB = make_tensor(make_smem_ptr(shared_tensors.smem_scale_B.data()), SmemLayoutScaleB{});
//
// Prepare the TMA loads for A and B
//
constexpr uint32_t cluster_shape_x = get<0>(typename DispatchPolicy::ClusterShape());
uint2 cluster_local_block_id = {block_rank_in_cluster % cluster_shape_x, block_rank_in_cluster / cluster_shape_x};
Tensor gA_mkl = get<0>(load_inputs);
Tensor gB_nkl = get<1>(load_inputs);
auto block_tma_a = mainloop_params.tma_load_a.get_slice(cluster_local_block_id.y);
auto block_tma_b = mainloop_params.tma_load_b.get_slice(cluster_local_block_id.x);
// Partition the inputs based on the current block coordinates.
auto [m_coord, n_coord, k_coord, l_coord] = blk_coord;
Tensor gA = gA_mkl(_,_,m_coord,_,l_coord); // (BLK_M,BLK_K,k)
Tensor gB = gB_nkl(_,_,n_coord,_,l_coord); // (BLK_N,BLK_K,k)
// Block scaling: load_scale has scaling tensors in global memory which are not tiled
Tensor mSFA_mkl = get<2>(load_inputs);
Tensor mSFB_nkl = get<3>(load_inputs);
auto scales_m = get<0>(mSFA_mkl.shape());
auto scales_n = get<0>(mSFB_nkl.shape());
Tensor cSFA_mkl = make_identity_tensor(mSFA_mkl.shape());
Tensor cSFB_nkl = make_identity_tensor(mSFB_nkl.shape());
Tensor gSFA = local_tile(
mSFA_mkl, make_tile(Int<ScaleMsPerTile>{}),
make_coord(m_coord,_,l_coord)); // (ScaleMsPerTile,k,1)
Tensor cSFA = local_tile(
cSFA_mkl, make_tile(Int<ScaleMsPerTile>{}),
make_coord(m_coord,_,l_coord));
Tensor gSFB = local_tile(
mSFB_nkl, make_tile(Int<ScaleNsPerTile>{}),
make_coord(n_coord,_,l_coord)); // (ScaleNsPerTile,k,1)
Tensor cSFB = local_tile(
cSFB_nkl, make_tile(Int<ScaleNsPerTile>{}),
make_coord(n_coord,_,l_coord));
TiledCopy scale_copy_a = make_tiled_copy(SmemBlockScalingCopyAtomA{},
Layout<Shape<_32>>{}, Layout<Shape<_1>>{});
TiledCopy scale_copy_b = make_tiled_copy(SmemBlockScalingCopyAtomB{},
Layout<Shape<_32>>{}, Layout<Shape<_1>>{});
ThrCopy thr_scale_copy_a = scale_copy_a.get_slice(thread_idx);
ThrCopy thr_scale_copy_b = scale_copy_b.get_slice(thread_idx);
Tensor tAgA_SFA = thr_scale_copy_a.partition_S(gSFA);
Tensor tAcA_SFA = thr_scale_copy_a.partition_S(cSFA);
Tensor tAsA_SFA = thr_scale_copy_a.partition_D(sSFA);
Tensor tBgB_SFB = thr_scale_copy_b.partition_S(gSFB);
Tensor tBcB_SFB = thr_scale_copy_b.partition_S(cSFB);
Tensor tBsB_SFB = thr_scale_copy_b.partition_D(sSFB);
Tensor tApA_SFA = make_tensor<bool>(shape(tAsA_SFA(_,_,0)));
Tensor tBpB_SFB = make_tensor<bool>(shape(tBsB_SFB(_,_,0)));
auto scale_m_lim = std::min(scales_m, (m_coord + 1) * ScaleMsPerTile);
auto scale_n_lim = std::min(scales_n, (n_coord + 1) * ScaleNsPerTile);
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < size(tApA_SFA); ++i)
tApA_SFA(i) = get<0>(tAcA_SFA(i)) < scale_m_lim;
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < size(tBpB_SFB); ++i)
tBpB_SFB(i) = get<0>(tBcB_SFB(i)) < scale_n_lim;
// Applies the mapping from block_tma_a
Tensor tAgA = block_tma_a.partition_S(gA); // (TMA,TMA_M,TMA_K,k)
Tensor tAsA = block_tma_a.partition_D(sA); // (TMA,TMA_M,TMA_K,PIPE)
Tensor tBgB = block_tma_b.partition_S(gB); // (TMA,TMA_N,TMA_K,k)
Tensor tBsB = block_tma_b.partition_D(sB); // (TMA,TMA_N,TMA_K,PIPE)
// TMA Multicast Masks
Layout cta_layout_mnk = make_layout(ClusterShape{});
auto cta_coord_mnk = cta_layout_mnk.get_flat_coord(block_rank_in_cluster);
uint16_t mcast_mask_a = create_tma_multicast_mask<1>(cta_layout_mnk, cta_coord_mnk);
uint16_t mcast_mask_b = create_tma_multicast_mask<0>(cta_layout_mnk, cta_coord_mnk);
// Mainloop
CUTLASS_PRAGMA_NO_UNROLL
for ( ; k_tile_count > 0; --k_tile_count) {
// LOCK smem_pipe_write for _writing_
pipeline.producer_acquire(smem_pipe_write);
//
// Copy gmem to smem for *k_tile_iter
//
int write_stage = smem_pipe_write.index();
if (lane_predicate) {
using BarrierType = typename MainloopPipeline::ProducerBarrierType;
BarrierType* tma_barrier = pipeline.producer_get_barrier(smem_pipe_write);
copy(mainloop_params.tma_load_a.with(*tma_barrier, mcast_mask_a), tAgA(_,_,_,*k_tile_iter), tAsA(_,_,_,write_stage));
copy(mainloop_params.tma_load_b.with(*tma_barrier, mcast_mask_b), tBgB(_,_,_,*k_tile_iter), tBsB(_,_,_,write_stage));
}
// Copy scale tensors
copy_if(scale_copy_a, tApA_SFA, tAgA_SFA(_,_,*k_tile_iter), tAsA_SFA(_,_,write_stage));
copy_if(scale_copy_b, tBpB_SFB, tBgB_SFB(_,_,*k_tile_iter), tBsB_SFB(_,_,write_stage));
pipeline.producer_commit(smem_pipe_write, cutlass::arch::cpasync_barrier_arrive_noinc);
++k_tile_iter;
// Advance smem_pipe_write
++smem_pipe_write;
}
}
/// Perform a Producer Epilogue to prevent early exit of blocks in a Cluster
CUTLASS_DEVICE void
load_tail(MainloopPipeline pipeline, PipelineState smem_pipe_write) {
int lane_predicate = cute::elect_one_sync();
// Issue the epilogue waits
if (lane_predicate) {
/* This helps avoid early exit of blocks 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
*/
pipeline.producer_tail(smem_pipe_write);
}
}
/// Perform a collective-scoped matrix multiply-accumulate
/// Consumer Perspective
template <
class FrgTensorC
>
CUTLASS_DEVICE void
mma(MainloopPipeline pipeline,
PipelineState smem_pipe_read,
FrgTensorC& accum,
int k_tile_count,
int thread_idx,
TensorStorage& shared_tensors,
Params const& mainloop_params) {
using namespace cute;
static_assert(is_rmem<FrgTensorC>::value, "C tensor must be rmem resident.");
FrgTensorC tmp_accum;
clear(accum);
clear(tmp_accum);
Tensor sA = make_tensor(make_smem_ptr(shared_tensors.smem_A.data()), SmemLayoutA{}); // (BLK_M,BLK_K,PIPE)
Tensor sB = make_tensor(make_smem_ptr(shared_tensors.smem_B.data()), SmemLayoutB{}); // (BLK_N,BLK_K,PIPE)
// Block scaling
Tensor sScaleAViewAsC = make_tensor(cute::make_smem_ptr(shared_tensors.smem_scale_A.data()),
Layout<
Shape<Shape<Int<ScaleGranularityM>, Int<ScaleMsPerTile>>, cute::tuple_element_t<1, TileShape>, Int<DispatchPolicy::Stages>>,
Stride<Stride<_0, _1>, _0, Int<ScaleMsPerTile>>
>{}); // ((ScaleGranularityM,ScaleMsPerTile),TileShape_N,stage)
Tensor sScaleBViewAsC = make_tensor(cute::make_smem_ptr(shared_tensors.smem_scale_B.data()),
Layout<
Shape<cute::tuple_element_t<0, TileShape>, Shape<Int<ScaleGranularityN>, Int<ScaleNsPerTile>>, Int<DispatchPolicy::Stages>>,
Stride<_0, Stride<_0, _1>, Int<ScaleNsPerTile>>
>{}); // (TileShape_M,(ScaleGranularityN,ScaleNsPerTile),stage)
//
// Define C accumulators and A/B partitioning
//
TiledMma tiled_mma;
auto thread_mma = tiled_mma.get_thread_slice(thread_idx);
// Allocate fragments and descriptors
Tensor tCrA = thread_mma.partition_fragment_A(sA(_,_,Int<0>{})); // (MMA,MMA_M,MMA_K)
Tensor tCrB = thread_mma.partition_fragment_B(sB(_,_,Int<0>{})); // (MMA,MMA_N,MMA_K)
Tensor tCsScaleAViewAsC = thread_mma.partition_C(sScaleAViewAsC); // (MMA,MMA_M,MMA_N,PIPE)
Tensor tCsScaleBViewAsC = thread_mma.partition_C(sScaleBViewAsC); // (MMA,MMA_M,MMA_N,PIPE)
//
// Copy Atom A and B retiling
//
auto smem_tiled_copy_A = make_tiled_copy_A(SmemCopyAtomA{}, tiled_mma);
auto smem_thr_copy_A = smem_tiled_copy_A.get_thread_slice(thread_idx);
Tensor tCsA = smem_thr_copy_A.partition_S(
as_position_independent_swizzle_tensor(sA)); // (CPY,CPY_M,CPY_K,PIPE)
Tensor tCrA_copy_view = smem_thr_copy_A.retile_D(tCrA); // (CPY,CPY_M,CPY_K)
auto smem_tiled_copy_B = make_tiled_copy_B(SmemCopyAtomB{}, tiled_mma);
auto smem_thr_copy_B = smem_tiled_copy_B.get_thread_slice(thread_idx);
Tensor tCsB = smem_thr_copy_B.partition_S(
as_position_independent_swizzle_tensor(sB)); // (CPY,CPY_M,CPY_K,PIPE)
Tensor tCrB_copy_view = smem_thr_copy_B.retile_D(tCrB); // (CPY,CPY_M,CPY_K)
Tensor tCrScaleAViewAsC = make_tensor_like<ElementSF>(tCsScaleAViewAsC(_,_,_,_0{})); // (MMA,MMA_M,MMA_N)
Tensor tCrScaleBViewAsC = make_tensor_like<ElementSF>(tCsScaleBViewAsC(_,_,_,_0{})); // (MMA,MMA_M,MMA_N)
CUTE_STATIC_ASSERT_V(size<1>(tCsA) == size<1>(tCrA_copy_view));
CUTE_STATIC_ASSERT_V(size<2>(tCsA) == size<2>(tCrA_copy_view));
CUTE_STATIC_ASSERT_V(size<1>(tCrA) == size<1>(accum));
CUTE_STATIC_ASSERT_V(size<1>(tCrB) == size<2>(accum));
CUTE_STATIC_ASSERT_V(size<2>(tCsA) == size<2>(tCsB));
CUTE_STATIC_ASSERT_V(size<3>(tCsA) == size<3>(tCsB));
CUTE_STATIC_ASSERT_V(Int<DispatchPolicy::Stages>{} == size<2>(sA));
CUTE_STATIC_ASSERT_V(Int<DispatchPolicy::Stages>{} == size<2>(sB));
//
// PIPELINED MAIN LOOP
//
// Size of the register pipeline
auto K_BLOCK_MAX = size<2>(tCrA);
int read_stage = smem_pipe_read.index();
auto tCsA_stage = tCsA(_,_,_,read_stage);
auto tCsB_stage = tCsB(_,_,_,read_stage);
auto copy_kblock = [&](auto k_block) {
// copy smem->rmem for A/B operand
copy(smem_tiled_copy_A, tCsA_stage(_,_,k_block), tCrA_copy_view(_,_,k_block));
copy(smem_tiled_copy_B, tCsB_stage(_,_,k_block), tCrB_copy_view(_,_,k_block));
// Left shift A,B for FP4
using MMAOp = typename TiledMma::MMA_Op;
fp4_shift_A(MMAOp{}, tCrA_copy_view(_,_,k_block));
fp4_shift_B(MMAOp{}, tCrB_copy_view(_,_,k_block));
};
auto copy_scale_s2r = [&](auto read_stage) {
copy(tCsScaleAViewAsC(_, _, _, read_stage), tCrScaleAViewAsC);
copy(tCsScaleBViewAsC(_, _, _, read_stage), tCrScaleBViewAsC);
if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile == 1) {
tCrScaleAViewAsC.data()[0] = tCrScaleAViewAsC.data()[0] * tCrScaleBViewAsC.data()[0];
}
if constexpr (ScaleMsPerTile > 1 && ScaleNsPerTile == 1) {
ElementSF scale_b = tCrScaleBViewAsC.data()[0];
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < size(tCrScaleAViewAsC); i++) {
tCrScaleAViewAsC.data()[i] = tCrScaleAViewAsC.data()[i] * scale_b;
}
}
if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile > 1) {
ElementSF scale_a = tCrScaleAViewAsC.data()[0];
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < size(tCrScaleBViewAsC); i++) {
tCrScaleBViewAsC.data()[i] = tCrScaleBViewAsC.data()[i] * scale_a;
}
}
};
auto rescale = [&]() {
// Block scale the accumulators with reg tensor `tCrScaleAViewAsC` and `tCrScaleBViewAsC`
if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile == 1) {
ElementSF scale_ab = tCrScaleAViewAsC.data()[0];
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < size(accum); ++i) {
accum(i) += tmp_accum(i) * scale_ab;
tmp_accum(i) = 0;
}
}
if constexpr (ScaleMsPerTile > 1 && ScaleNsPerTile == 1) {
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < size(accum); ++i) {
accum(i) += tmp_accum(i) * tCrScaleAViewAsC(i);
tmp_accum(i) = 0;
}
}
if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile > 1) {
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < size(accum); ++i) {
accum(i) += tmp_accum(i) * tCrScaleBViewAsC(i);
tmp_accum(i) = 0;
}
}
if constexpr (ScaleMsPerTile > 1 && ScaleNsPerTile > 1) {
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < size(accum); ++i) {
accum(i) += tmp_accum(i) * tCrScaleAViewAsC(i) * tCrScaleBViewAsC(i);
tmp_accum(i) = 0;
}
}
};
auto gemm_kblock = [&](auto k_block) {
// (V,M) x (V,N) => (V,M,N)
cute::gemm(tiled_mma, tCrA(_,_,k_block), tCrB(_,_,k_block), tmp_accum);
};
pipeline.consumer_wait(smem_pipe_read);
copy_scale_s2r(read_stage);
copy_kblock(_0{});
CUTLASS_PRAGMA_NO_UNROLL
for ( ; k_tile_count > 1; --k_tile_count) {
//
// Compute on k_tile
//
for_each(make_int_sequence<K_BLOCK_MAX>{}, [&] (auto k_block) {
auto k_block_next = ((k_block + 1) == K_BLOCK_MAX) ? 0 : (k_block + 1);
if (k_block == K_BLOCK_MAX - 1) {
cutlass::arch::NamedBarrier::sync(
thr_size(tiled_mma), cutlass::arch::ReservedNamedBarriers::Sm120MainloopBarrier);
// UNLOCK smem_pipe_read, done _computing_ on it
pipeline.consumer_release(smem_pipe_read);
++smem_pipe_read;
read_stage = smem_pipe_read.index();
tCsA_stage = tCsA(_,_,_,read_stage);
tCsB_stage = tCsB(_,_,_,read_stage);
pipeline.consumer_wait(smem_pipe_read);
}
copy_kblock(k_block_next);
gemm_kblock(k_block);
if (k_block == K_BLOCK_MAX - 1) {
rescale();
copy_scale_s2r(read_stage);
}
});
} // k_tile_count
//
// Hoist out last k_tile
//
for_each(make_int_sequence<K_BLOCK_MAX>{}, [&] (auto k_block) {
auto k_block_next = ((k_block + 1) == K_BLOCK_MAX) ? 0 : (k_block + 1);
if (k_block == K_BLOCK_MAX - 1) {
cutlass::arch::NamedBarrier::sync(
thr_size(tiled_mma), cutlass::arch::ReservedNamedBarriers::Sm120MainloopBarrier);
// UNLOCK smem_pipe_read, done _computing_ on it
pipeline.consumer_release(smem_pipe_read);
++smem_pipe_read;
}
if (k_block_next > 0) {
copy_kblock(k_block_next);
}
gemm_kblock(k_block);
});
rescale();
}
/// Perform a Consumer Epilogue to release all buffers
CUTLASS_DEVICE void
mma_tail(MainloopPipeline, PipelineState, int) {
}
};
/////////////////////////////////////////////////////////////////////////////////////////////////
} // namespace cutlass::gemm::collective
/////////////////////////////////////////////////////////////////////////////////////////////////
include/cutlass/gemm/collective/sm120_sparse_mma_tma.hpp
+17 -19
View File
@@ -45,7 +45,6 @@
#include "cute/atom/mma_atom.hpp"
#include "cute/algorithm/functional.hpp"
#include "cute/algorithm/gemm.hpp"
#include "cute/tensor_predicate.hpp"
#include "cute/numeric/arithmetic_tuple.hpp"
/////////////////////////////////////////////////////////////////////////////////////////////////
@@ -129,14 +128,14 @@ struct CollectiveMma<
using RuntimeDataTypeA = void*;
using RuntimeDataTypeB = void*;
static constexpr int ThreadCount = size(TiledMma{});
static constexpr int ElementAMmaSparsity = ElementAMma::sparsity;
static constexpr int ElementEMmaSparsity = ElementEMma::sparsity;
// Asymmetric buffering
// Tensor A/B could have different buffering, with TILEK, and STAGEs.
// It let AsymmetricKRatio equals TILEK_A / TILEK_B, to make sure A/B's
// Tensor A/B could have different buffering, with TILEK, and STAGEs.
// It let AsymmetricKRatio equals TILEK_A / TILEK_B, to make sure A/B's
// pipeline keep same steps when procude / consume data.
static constexpr int AsymmetricKRatio = DispatchPolicy::StagesA != DispatchPolicy::StagesB ? 2 : 1;
@@ -418,7 +417,7 @@ struct CollectiveMma<
make_tile(make_layout(size<1>(thr_layout_vmnk)),
make_layout(size<3>(thr_layout_vmnk))));
auto thr_tensor = zipped_divide(tv_tensor, thr_tile); // ((ThrV,(ThrM,ThrK)),(FrgV,(RestM,RestK)))
// Fragment layout
return thr_tensor;
}
@@ -694,16 +693,15 @@ struct CollectiveMma<
if constexpr (IsELoadPred) {
// Get predication based on logical element coordinates.
Tensor cE_mk = local_tile(
make_identity_tensor(Shape_MK),
make_shape(get<0>(TileShape{}), get<2>(TileShape{})),
make_identity_tensor(Shape_MK),
make_shape(get<0>(TileShape{}), get<2>(TileShape{})),
make_shape(m_coord, k_coord)); // (BLK_M, BLK_K)
Tensor tCcE = gmem_thr_copy_E.partition_S(cE_mk); // (CPY,CPY_M,CPY_K)
auto [atom, vec] = get_copy_atom_and_common_vec();
// Coordinate comparison for out of bound (OOB) predication
Tensor tZcE = zipped_divide(tCcE, vec);
auto pred_fn = [&](auto coord){ return cute::elem_less(tZcE(Int<0>{}, coord), Shape_MK); };
Tensor tZpE = cute::lazy::transform(zipped_divide(tCcE, vec), [&](auto const& c){ return cute::elem_less(c, Shape_MK); });
// Copy
cute::copy_if(atom, pred_fn, zipped_divide(tCgE, vec), zipped_divide(tCrE_copy_view, vec));
cute::copy_if(atom, tZpE, zipped_divide(tCgE, vec), zipped_divide(tCrE_copy_view, vec));
}
else {
// Copy
@@ -712,7 +710,7 @@ struct CollectiveMma<
}
return tCrE;
}
/// Perform a collective-scoped matrix multiply-accumulate
/// Consumer Perspective
template <
@@ -849,7 +847,7 @@ struct CollectiveMma<
// Copy E from SMEM to register
auto copy_E = [&](auto m_block, auto k_block) CUTLASS_LAMBDA_FUNC_INLINE {
// copy smem->rmem for E operand
copy( recast<RegisterE>(tCsE(_,m_block,k_block,smem_pipe_read_mk.index())),
copy( recast<RegisterE>(tCsE(_,m_block,k_block,smem_pipe_read_mk.index())),
recast<RegisterE>(tCrE_copy_view(_,m_block,k_block)));
};
@@ -877,8 +875,8 @@ struct CollectiveMma<
copy_E(m_block, k_block);
// Gemm
cute::gemm(tiled_mma,
make_zip_tensor(tCrA(_,m_block,k_block), tCrE(_,m_block,k_block)),
cute::gemm(tiled_mma,
make_zip_tensor(tCrA(_,m_block,k_block), tCrE(_,m_block,k_block)),
tCrB(_,n_block,k_block),
accum(_,m_block,n_block));
});
@@ -914,8 +912,8 @@ struct CollectiveMma<
copy_transform_A(m_block, k_block);
// Gemm
cute::gemm(tiled_mma,
make_zip_tensor(tCrA(_,m_block,k_block), tCrE(_,m_block,k_block)),
cute::gemm(tiled_mma,
make_zip_tensor(tCrA(_,m_block,k_block), tCrE(_,m_block,k_block)),
tCrB(_,n_block,k_block),
accum(_,m_block,n_block));
});
@@ -941,8 +939,8 @@ struct CollectiveMma<
copy_transform_A(m_block, k_block_a);
// Gemm
cute::gemm(tiled_mma,
make_zip_tensor(tCrA(_,m_block,k_block_a), tCrE(_,m_block,k_block_a)),
cute::gemm(tiled_mma,
make_zip_tensor(tCrA(_,m_block,k_block_a), tCrE(_,m_block,k_block_a)),
tCrB(_,n_block,k_block),
accum(_,m_block,n_block));
});
@@ -970,7 +968,7 @@ struct CollectiveMma<
gemm_loop_with_SmemE();
}
// Case when A/B with different stages, and keep E in GMEM.
else {
else {
gemm_loop_with_GmemE();
} // end if
include/cutlass/gemm/collective/sm70_mma_twostage.hpp
-1
View File
@@ -37,7 +37,6 @@
#include "cute/atom/mma_atom.hpp"
#include "cute/algorithm/gemm.hpp"
#include "cute/atom/mma_atom.hpp"
#include "cute/tensor_predicate.hpp"
#include "cutlass/gemm/collective/collective_mma_decl.hpp"
include/cutlass/gemm/collective/sm80_mma_multistage.hpp
+3 -4
View File
@@ -36,7 +36,6 @@
#include "cute/algorithm/functional.hpp"
#include "cute/atom/mma_atom.hpp"
#include "cute/algorithm/gemm.hpp"
#include "cute/tensor_predicate.hpp"
#include "cute/numeric/arithmetic_tuple.hpp"
@@ -100,7 +99,7 @@ struct CollectiveMma<
using TransformA = TransformA_;
using TransformB = TransformB_;
using ArchTag = typename DispatchPolicy::ArchTag;
// Follow the change in TestSmall: TileShape switch to CtaShape
// Follow the change in TestSmall: TileShape switch to CtaShape
// For sm80 arch, CtaShape should euqal to TileShape
using CtaShape_MNK = TileShape;
@@ -318,7 +317,7 @@ struct CollectiveMma<
copy(gmem_tiled_copy_A, tAgA(_,_,_,*k_tile_iter), tAsA(_,_,_,smem_pipe_write));
copy(gmem_tiled_copy_B, tBgB(_,_,_,*k_tile_iter), tBsB(_,_,_,smem_pipe_write));
cp_async_fence();
// Advance the tile
--k_tile_count;
if (k_tile_count > 0) { ++k_tile_iter; }
@@ -390,7 +389,7 @@ struct CollectiveMma<
Stages,
ClusterShape_>;
using TileShape = TileShape_;
// Follow the change in TestSmall: TileShape switch to CtaShape
// Follow the change in TestSmall: TileShape switch to CtaShape
// In legacy arch, it should be same
using CtaShape_MNK = TileShape;
using ElementA = ElementA_;
include/cutlass/gemm/collective/sm90_mma_array_tma_gmma_rs_warpspecialized_mixed_input.hpp
+71 -72
View File
@@ -43,7 +43,6 @@
#include "cute/algorithm/functional.hpp"
#include "cute/atom/mma_atom.hpp"
#include "cute/algorithm/gemm.hpp"
#include "cute/tensor_predicate.hpp"
#include "cute/numeric/arithmetic_tuple.hpp"
/////////////////////////////////////////////////////////////////////////////////////////////////
@@ -95,7 +94,7 @@ public:
ConvertAndScale,
ConvertAndScaleWithZero
};
//
// Type Aliases
//
@@ -105,10 +104,10 @@ public:
private:
template<class T> friend struct detail::MixedInputUtils;
using CollectiveType = CollectiveMma<DispatchPolicy, TileShape_,
ElementAOptionalTuple, StrideA_,
using CollectiveType = CollectiveMma<DispatchPolicy, TileShape_,
ElementAOptionalTuple, StrideA_,
ElementBOptionalTuple, StrideB_,
TiledMma_,
TiledMma_,
GmemTiledCopyA_, SmemLayoutAtomA_, SmemCopyAtomA_,
TransformA_,
GmemTiledCopyB_, SmemLayoutAtomB_, SmemCopyAtomB_,
@@ -124,9 +123,9 @@ private:
using ZeroB = detail::deduce_mixed_width_dtype_t<2, ElementBOptionalTuple>;
public:
static_assert(cute::is_tuple<ElementAOptionalTuple>::value ^ cute::is_tuple<ElementBOptionalTuple>::value,
static_assert(cute::is_tuple<ElementAOptionalTuple>::value ^ cute::is_tuple<ElementBOptionalTuple>::value,
"Either A OR B must be a tuple. It must take the from {ElementOperand, [ElementScale], [ElementZero]}. Inputs in [] are optional.");
using ElementA = detail::deduce_mixed_width_dtype_t<0, ElementAOptionalTuple>;
using ElementB = detail::deduce_mixed_width_dtype_t<0, ElementBOptionalTuple>;
static constexpr bool IsATransformed = cute::is_tuple<ElementAOptionalTuple>::value;
@@ -140,23 +139,23 @@ public:
using InternalStrideA = cute::remove_pointer_t<StrideA>;
using StrideB = StrideB_;
using InternalStrideB = cute::remove_pointer_t<StrideB>;
using StrideScale = cute::Stride<cute::Int<1>, int64_t, int64_t>;
using NonVoidStrideScale = cute::conditional_t<cute::is_void_v<StrideScale>, cute::Stride<_1, int64_t, int64_t>, StrideScale>;
static_assert(( IsATransformed && (cutlass::gemm::detail::is_k_major<StrideA>() || is_layout<StrideA>::value || is_layout<InternalStrideA>::value)) ||
static_assert(( IsATransformed && (cutlass::gemm::detail::is_k_major<StrideA>() || is_layout<StrideA>::value || is_layout<InternalStrideA>::value)) ||
(!IsATransformed && (cutlass::gemm::detail::is_k_major<StrideB>() || is_layout<StrideB>::value || is_layout<InternalStrideB>::value)),
"The transformed type must be K-major.");
static_assert(( IsATransformed && (sizeof(ElementB) == 2)) ||
(!IsATransformed && (sizeof(ElementA) == 2)) ||
((cutlass::gemm::detail::is_k_major<StrideA>() || is_layout<StrideA>::value || is_layout<InternalStrideA>::value) &&
(cutlass::gemm::detail::is_k_major<StrideB>() || is_layout<StrideB>::value || is_layout<InternalStrideB>::value)),
((cutlass::gemm::detail::is_k_major<StrideA>() || is_layout<StrideA>::value || is_layout<InternalStrideA>::value) &&
(cutlass::gemm::detail::is_k_major<StrideB>() || is_layout<StrideB>::value || is_layout<InternalStrideB>::value)),
"The unscaled element must be 2 bytes OR both inputs must be K-major");
static_assert(cutlass::gemm::detail::is_mn_major<NonVoidStrideScale>(),
static_assert(cutlass::gemm::detail::is_mn_major<NonVoidStrideScale>(),
"Scale must be MN major [Col Major if A is scaled, Row Major if B is scaled].");
using CtaShape_MNK = decltype(shape_div(TileShape{}, ClusterShape{}));
using TiledMma = TiledMma_;
using ElementAccumulator = typename TiledMma::ValTypeC;
@@ -224,10 +223,10 @@ public:
/// Tile along modes in a way that maximizes the TMA box size.
using SmemLayoutA = decltype(detail::get_smem_layout<DispatchPolicy::Stages>(SwappedSmemLayoutAtomA{}, select<0,2>(TileShape{}), InternalSwappedStrideA{}));
using SmemLayoutB = decltype(detail::get_smem_layout<DispatchPolicy::Stages>(SwappedSmemLayoutAtomB{}, select<1,2>(TileShape{}), InternalSwappedStrideB{}));
// It is assumed that the scales and zero-points share the same smem layout
using SmemLayoutScale = decltype(tile_to_shape(
SmemLayoutAtomScale{},
SmemLayoutAtomScale{},
make_shape(shape<0>(ScaleTileShape{}), shape<1>(ScaleTileShape{}), Int<Stages>{}),
cute::conditional_t< ::cutlass::gemm::detail::is_major<0,NonVoidStrideScale>(), Step<_2,_1,_3>, Step<_1,_2,_3>>{}));
@@ -245,18 +244,18 @@ public:
static_assert(size<1>(SmemLayoutAtomScale{}) == 1, "size<1>(SmemLayoutAtomScale) must be 1.");
private:
static constexpr ConversionMode
static constexpr ConversionMode
get_conversion_mode() {
if constexpr (cute::is_void_v<ElementScale>) {
return ConversionMode::DirectConvert;
}
}
else if constexpr (cute::is_void_v<ElementZero>) {
return ConversionMode::ConvertAndScale;
}
else {
return ConversionMode::ConvertAndScaleWithZero;
}
}
}
public:
static constexpr ConversionMode KernelConversionMode = get_conversion_mode();
@@ -264,7 +263,7 @@ public:
KernelConversionMode == ConversionMode::ConvertAndScaleWithZero;
static constexpr bool UseScaleLookupTable = KernelConversionMode == ConversionMode::ConvertAndScale &&
cutlass::detail::is_Array_v<ElementScale>;
static constexpr size_t SmemAlignmentA = cutlass::detail::alignment_for_swizzle(SmemLayoutA{});
static constexpr size_t SmemAlignmentA = cutlass::detail::alignment_for_swizzle(SmemLayoutA{});
static constexpr size_t SmemAlignmentB = cutlass::detail::alignment_for_swizzle(SmemLayoutB{});
static constexpr size_t SmemAlignmentScale = cute::max(SmemAlignmentA, SmemAlignmentB);
@@ -341,7 +340,7 @@ public:
SmemLayoutScale{}(_,_,cute::Int<0>{}),
ScaleTileShape{},
_1{})); // mcast along N mode for this M load, if any. Scale is ALWAYS loaded with A for RF kernel
TMA_A tma_load_a;
TMA_B tma_load_b;
uint32_t tma_transaction_bytes = TmaTransactionBytes;
@@ -415,7 +414,7 @@ public:
dA = InternalSwappedStrideA{};
if constexpr (is_layout<InternalSwappedStrideA>::value) {
dA = make_layout(
transform_leaf(dA.shape(), [](auto x){
transform_leaf(dA.shape(), [](auto x){
if constexpr (not is_static_v<decltype(x)>) {
return static_cast<decltype(x)>(1);
} else {
@@ -521,15 +520,15 @@ public:
_1{}); // mcast along N mode for this M load, if any
return SwapAB ? args_setup(args.ptr_B, args.ptr_A, scale_k, args.chunk_size, (args.chunk_size + size<2>(TileShape{}) - 1) / size<2>(TileShape{}))
: args_setup(args.ptr_A, args.ptr_B, scale_k, args.chunk_size, (args.chunk_size + size<2>(TileShape{}) - 1) / size<2>(TileShape{}));
}
}
else {
static_assert(cutlass::detail::dependent_false<KernelSchedule>, "Conversion mode not handled in to_underlying_arguments.");
}
}
}
else {
static_assert(cutlass::detail::dependent_false<KernelSchedule>, "Conversion mode not handled in to_underlying_arguments.");
}
}
}
template <class ProblemShape>
@@ -545,15 +544,15 @@ public:
if constexpr (KernelConversionMode == ConversionMode::DirectConvert) {
// Allocate gmem space for input tensormaps per each SM, A tensormap copies followed by B tensormap copies
return calculate_workspace_size(2);
}
}
else if constexpr (KernelConversionMode == ConversionMode::ConvertAndScale) {
// Allocate gmem space for input tensormaps per each SM, A tensormap copies followed by B tensormap copies, followed by scale tensormap copies
return calculate_workspace_size(3);
}
}
else if constexpr (KernelConversionMode == ConversionMode::ConvertAndScaleWithZero) {
// Allocate gmem space for input tensormaps per each SM, A tensormap copies followed by B tensormap copies, followed by scale and zeros tensormap copies
return calculate_workspace_size(4);
}
}
else {
static_assert(cutlass::detail::dependent_false<KernelSchedule>, "Conversion mode not handled in get_workspace_size.");
}
@@ -612,7 +611,7 @@ public:
constexpr int min_tma_aligned_elements_zero = tma_alignment_bits / cutlass::sizeof_bits<ElementZero>::value;
implementable = implementable && cutlass::detail::check_alignment<min_tma_aligned_elements_zero>(cute::make_shape(scale_mn,scale_k,L), StrideScale{});
implementable = implementable && (args.ptr_Z != nullptr);
}
}
else {
static_assert(cutlass::detail::dependent_false<KernelSchedule>, "Conversion mode not handled in can_implement.");
}
@@ -661,7 +660,7 @@ public:
if constexpr (KernelConversionMode == ConversionMode::DirectConvert) {
return cute::make_tuple(gA_mkl, gB_nkl);
}
}
else if constexpr (ModeHasScales) {
const int scale_mn = SwapAB ? N : M;
auto scale_k = mainloop_params.scale_k;
@@ -678,7 +677,7 @@ public:
else {
static_assert(cutlass::detail::dependent_false<KernelSchedule>, "Conversion mode not handled in load_init.");
}
}
}
else {
static_assert(cutlass::detail::dependent_false<KernelSchedule>, "Conversion mode not handled in load_init.");
}
@@ -694,7 +693,7 @@ public:
CUTLASS_DEVICE void
load(
Params const& mainloop_params,
MainloopPipeline pipeline,
MainloopPipeline pipeline,
PipelineState smem_pipe_write,
cute::tuple<Ts...> const& load_inputs,
cute::tuple<TMs...> const& input_tensormaps,
@@ -707,15 +706,15 @@ public:
if constexpr (KernelConversionMode == ConversionMode::DirectConvert) {
static_assert(sizeof... (Ts) == 2, "Direct convert needs two inputs");
static_assert(sizeof... (TMs) == 2, "Direct convert needs two tensormaps");
}
}
else if constexpr (KernelConversionMode == ConversionMode::ConvertAndScale) {
static_assert(sizeof... (Ts) == 3, "Scaled convert needs three inputs");
static_assert(sizeof... (TMs) == 3, "Scaled convert needs three tensormaps");
}
}
else if constexpr (KernelConversionMode == ConversionMode::ConvertAndScaleWithZero) {
static_assert(sizeof... (Ts) == 4, "Scaled and zero convert needs four inputs");
static_assert(sizeof... (TMs) == 4, "Scaled and zero convert needs four tensormaps");
}
}
else {
static_assert(cutlass::detail::dependent_false<KernelSchedule>, "Conversion mode not handled in TMA load.");
}
@@ -809,7 +808,7 @@ public:
if constexpr (KernelConversionMode == ConversionMode::ConvertAndScale) {
// Nothing extra to do
}
}
else if constexpr (KernelConversionMode == ConversionMode::ConvertAndScaleWithZero) {
auto tZgZ = get<2>(extra_input_partitions);
auto tZsZ = get<3>(extra_input_partitions);
@@ -819,8 +818,8 @@ public:
}
else {
static_assert(cutlass::detail::dependent_false<KernelSchedule>, "Conversion mode not handled for TMA copy op.");
}
}
}
}
else {
static_assert(cutlass::detail::dependent_false<KernelSchedule>, "Conversion mode not handled for TMA copy op.");
}
@@ -839,9 +838,9 @@ public:
// Issue the epilogue waits
if (lane_predicate) {
// This helps avoid early exit of blocks in Cluster.
// Waits for all stages to either be released (all
// Waits for all stages to either be released (all
// Consumer UNLOCKs), or if the stage was never used
// then it would just be acquired since the phase was
// then it would just be acquired since the phase was
// still inverted from make_producer_start_state.
pipeline.producer_tail(smem_pipe_write);
}
@@ -875,7 +874,7 @@ public:
int warp_idx = canonical_warp_idx_sync();
[[maybe_unused]] int warp_group_thread_idx = thread_idx % 128;
Tensor sA_ = make_tensor(make_smem_ptr(shared_tensors.smem_A.begin()), SmemLayoutA{}); // (BLK_M,BLK_K,PIPE)
Tensor sA = as_position_independent_swizzle_tensor(sA_); // (BLK_M,BLK_K,PIPE)
@@ -889,7 +888,7 @@ public:
// Layout of warp group to thread mapping
static_assert(stride<0>(typename TiledMma::BLayout{}) == 0 and
size<0>(typename TiledMma::BLayout{}) == NumThreadsPerWarpGroup,
size<0>(typename TiledMma::BLayout{}) == NumThreadsPerWarpGroup,
"Stride of the first mode must be 0 and the size of the mode must be NumThreadsPerWarpGroup");
constexpr int MmaWarpGroups = size(TiledMma{}) / NumThreadsPerWarpGroup;
@@ -938,7 +937,7 @@ public:
CUTE_STATIC_ASSERT_V(size<3>(tCsA) == size<3>(tCsB)); // PIPE
CUTE_STATIC_ASSERT_V(Int<DispatchPolicy::Stages>{} == size<2>(sA)); // PIPE
CUTE_STATIC_ASSERT_V(Int<DispatchPolicy::Stages>{} == size<2>(sB)); // PIPE
//
// PIPELINED MAIN LOOP
//
@@ -967,15 +966,15 @@ public:
// copy smem->rmem for A operand
Utils::copy_tensors_MK(smem_tiled_copy_A, tCsA, tCrA_copy_view,
Utils::copy_tensors_MK(smem_tiled_copy_A, tCsA, tCrA_copy_view,
partitioned_extra_info, copy_partitions_extra_info, 0, read_stage);
if (K_BLOCK_MAX > 1) {
Utils::copy_tensors_MK(smem_tiled_copy_A, tCsA, tCrA_copy_view,
Utils::copy_tensors_MK(smem_tiled_copy_A, tCsA, tCrA_copy_view,
partitioned_extra_info, copy_partitions_extra_info, 1, read_stage);
}
Utils::dequantize_A_kblock(tCrA_load, tCrA_mma, partitioned_extra_info, 0);
// Unroll the K mode manually to set scale D to 1
CUTLASS_PRAGMA_UNROLL
for (int k_block = 0; k_block < K_BLOCK_MAX; ++k_block) {
@@ -986,26 +985,26 @@ public:
warpgroup_commit_batch();
if (k_block < K_BLOCK_MAX - 2) {
Utils::copy_tensors_MK(smem_tiled_copy_A, tCsA, tCrA_copy_view,
Utils::copy_tensors_MK(smem_tiled_copy_A, tCsA, tCrA_copy_view,
partitioned_extra_info, copy_partitions_extra_info, k_block + 2, read_stage);
}
if (k_block < K_BLOCK_MAX - 1) {
Utils::dequantize_A_kblock(tCrA_load, tCrA_mma, partitioned_extra_info, k_block + 1);
}
}
}
--k_tile_count;
if (k_tile_count > 0) {
// Wait for K_BLOCK_MAX - 1 to be in flight to ensure that it is safe to overwrite the A registers for the first mma.
// Wait for K_BLOCK_MAX - 1 to be in flight to ensure that it is safe to overwrite the A registers for the first mma.
pipeline.consumer_wait(smem_pipe_read, barrier_token);
Utils::copy_tensors_MK(smem_tiled_copy_A, tCsA, tCrA_copy_view,
Utils::copy_tensors_MK(smem_tiled_copy_A, tCsA, tCrA_copy_view,
partitioned_extra_info, copy_partitions_extra_info, 0, smem_pipe_read.index());
Utils::copy_tensors_MK(smem_tiled_copy_A, tCsA, tCrA_copy_view,
Utils::copy_tensors_MK(smem_tiled_copy_A, tCsA, tCrA_copy_view,
partitioned_extra_info, copy_partitions_extra_info, 1, smem_pipe_read.index());
warpgroup_wait<K_WAIT_MAX>();
warpgroup_wait<K_WAIT_MAX>();
Utils::dequantize_A_kblock(tCrA_load, tCrA_mma, partitioned_extra_info, 0);
}
}
@@ -1030,7 +1029,7 @@ public:
// Unroll the K mode manually to set scale D to 1
CUTLASS_PRAGMA_UNROLL
for (int k_block = 0; k_block < K_BLOCK_MAX; ++k_block) {
warpgroup_arrive();
// (V,M) x (V,N) => (V,M,N)
cute::gemm(tiled_mma, tCrA_mma(_,_,k_block), tCrB(_,_,k_block,read_stage), accum);
@@ -1047,18 +1046,18 @@ public:
barrier_token = pipeline.consumer_try_wait(smem_pipe_read);
}
if (k_block == K_BLOCK_MAX - 1) {
if (k_block == K_BLOCK_MAX - 1) {
pipeline.consumer_wait(smem_pipe_read, barrier_token);
Utils::copy_tensors_MK(smem_tiled_copy_A, tCsA, tCrA_copy_view,
Utils::copy_tensors_MK(smem_tiled_copy_A, tCsA, tCrA_copy_view,
partitioned_extra_info, copy_partitions_extra_info, 0, smem_pipe_read.index());
Utils::copy_tensors_MK(smem_tiled_copy_A, tCsA, tCrA_copy_view,
Utils::copy_tensors_MK(smem_tiled_copy_A, tCsA, tCrA_copy_view,
partitioned_extra_info, copy_partitions_extra_info, 1, smem_pipe_read.index());
Utils::dequantize_A_kblock(tCrA_load, tCrA_mma, partitioned_extra_info, 0);
}
}
else {
if (k_block < K_BLOCK_MAX - 2) {
Utils::copy_tensors_MK(smem_tiled_copy_A, tCsA, tCrA_copy_view,
Utils::copy_tensors_MK(smem_tiled_copy_A, tCsA, tCrA_copy_view,
partitioned_extra_info, copy_partitions_extra_info, k_block + 2, read_stage);
}
Utils::dequantize_A_kblock(tCrA_load, tCrA_mma, partitioned_extra_info, k_block + 1);
@@ -1078,7 +1077,7 @@ public:
int read_stage = smem_pipe_read.index();
warpgroup_fence_operand(accum);
// Unroll the K mode manually to set scale D to 1
CUTLASS_PRAGMA_UNROLL
for (int k_block = 0; k_block < K_BLOCK_MAX; ++k_block) {
@@ -1097,7 +1096,7 @@ public:
}
if (k_block < K_BLOCK_MAX - 2) {
Utils::copy_tensors_MK(smem_tiled_copy_A, tCsA, tCrA_copy_view,
Utils::copy_tensors_MK(smem_tiled_copy_A, tCsA, tCrA_copy_view,
partitioned_extra_info, copy_partitions_extra_info, k_block + 2, read_stage);
}
if (k_block < K_BLOCK_MAX - 1) {
@@ -1117,7 +1116,7 @@ public:
k_tile_count -= prologue_mma_count;
smem_pipe_release.advance(k_tile_count);
// Wait on all GMMAs to complete
warpgroup_wait<0>();
@@ -1153,7 +1152,7 @@ public:
copy(recast<uint128_t>(pA_tensormap), recast<uint128_t>(sA_tensormap));
copy(recast<uint128_t>(pB_tensormap), recast<uint128_t>(sB_tensormap));
}
if constexpr (KernelConversionMode == ConversionMode::ConvertAndScale) {
Tensor pS_tensormap = make_tensor(mainloop_params.tma_load_scale.get_tma_descriptor(), Int<1>{}, Int<1>{});
Tensor sS_tensormap = make_tensor(make_smem_ptr(&shared_tensormaps.smem_tensormap_scale), Int<1>{}, Int<1>{});
@@ -1225,7 +1224,7 @@ public:
const uint32_t M = (SwapAB? get<1>(problem_shape_mnkl) : get<0>(problem_shape_mnkl));
const uint32_t N = (SwapAB? get<0>(problem_shape_mnkl) : get<1>(problem_shape_mnkl));
const uint32_t K = get<2>(problem_shape_mnkl);
// Replace all dims for consistency
constexpr int MaxTensorRank = 5;
cute::array<uint32_t, MaxTensorRank> prob_shape_A = {1,1,1,1,1};
@@ -1243,23 +1242,23 @@ public:
SwappedElementB const* ptr_B = nullptr;
Tensor tensor_b = make_tensor(ptr_B, detail::get_gmem_layout(make_shape(N,K,Int<1>{}), mainloop_params.ptr_dB[next_group]));
cute::detail::fill_tma_gmem_shape_stride(mainloop_params.tma_load_a, tensor_a,
cute::detail::fill_tma_gmem_shape_stride(mainloop_params.tma_load_a, tensor_a,
prob_shape_A, prob_stride_A);
cute::detail::fill_tma_gmem_shape_stride(mainloop_params.tma_load_b, tensor_b,
cute::detail::fill_tma_gmem_shape_stride(mainloop_params.tma_load_b, tensor_b,
prob_shape_B, prob_stride_B);
if constexpr (KernelConversionMode == ConversionMode::ConvertAndScale) {
NonVoidElementScale const* ptr_S = nullptr;
auto scale_k = ceil_div(K, mainloop_params.chunk_size);
Tensor tensor_scale = make_tensor(detail::get_logical_ptr(ptr_S), make_shape(M,scale_k,Int<1>{}), mainloop_params.dS[next_group]);
cute::detail::fill_tma_gmem_shape_stride(mainloop_params.tma_load_scale, tensor_scale,
cute::detail::fill_tma_gmem_shape_stride(mainloop_params.tma_load_scale, tensor_scale,
prob_shape_scale, prob_stride_scale);
}
else if constexpr (KernelConversionMode == ConversionMode::ConvertAndScaleWithZero) {
ElementZero const* ptr_Z = nullptr;
auto scale_k = ceil_div(K, mainloop_params.chunk_size);
Tensor tensor_zero = make_tensor(detail::get_logical_ptr(ptr_Z), make_shape(M,scale_k,Int<1>{}), mainloop_params.dS[next_group]);
cute::detail::fill_tma_gmem_shape_stride(mainloop_params.tma_load_zero, tensor_zero,
cute::detail::fill_tma_gmem_shape_stride(mainloop_params.tma_load_zero, tensor_zero,
prob_shape_zero, prob_stride_zero);
}
else if constexpr (KernelConversionMode != ConversionMode::DirectConvert){
@@ -1300,7 +1299,7 @@ public:
}
else if constexpr (KernelConversionMode != ConversionMode::DirectConvert){
static_assert(cutlass::detail::dependent_false<KernelSchedule>, "Conversion mode not handled in tensormaps_replace_global_tensor_properties.");
}
}
}
template <class... TMs, class ProblemShape_MNKL>
include/cutlass/gemm/collective/sm90_mma_array_tma_gmma_ss_warpspecialized.hpp
+8 -9
View File
@@ -42,7 +42,6 @@
#include "cute/algorithm/functional.hpp"
#include "cute/atom/mma_atom.hpp"
#include "cute/algorithm/gemm.hpp"
#include "cute/tensor_predicate.hpp"
#include "cute/numeric/arithmetic_tuple.hpp"
/////////////////////////////////////////////////////////////////////////////////////////////////
@@ -359,7 +358,7 @@ struct CollectiveMma<
CUTLASS_DEVICE void
load(
Params const& mainloop_params,
MainloopPipeline pipeline,
MainloopPipeline pipeline,
PipelineState smem_pipe_write,
cute::tuple<TensorA, TensorB> const& load_inputs,
cute::tuple<TensorMapA, TensorMapB> const& input_tensormaps,
@@ -451,9 +450,9 @@ struct CollectiveMma<
// Issue the epilogue waits
if (lane_predicate) {
// This helps avoid early exit of blocks in Cluster.
// Waits for all stages to either be released (all
// Waits for all stages to either be released (all
// Consumer UNLOCKs), or if the stage was never used
// then it would just be acquired since the phase was
// then it would just be acquired since the phase was
// still inverted from make_producer_start_state.
pipeline.producer_tail(smem_pipe_write);
}
@@ -489,10 +488,10 @@ struct CollectiveMma<
// Layout of warp group to thread mapping
static_assert(stride<0>(typename TiledMma::ALayout{}) == 0 and
static_assert(stride<0>(typename TiledMma::ALayout{}) == 0 and
stride<0>(typename TiledMma::BLayout{}) == 0 and
size<0>(typename TiledMma::ALayout{}) == NumThreadsPerWarpGroup and
size<0>(typename TiledMma::BLayout{}) == NumThreadsPerWarpGroup,
size<0>(typename TiledMma::BLayout{}) == NumThreadsPerWarpGroup,
"Stride of the first mode must be 0 and the size of the mode must be NumThreadsPerWarpGroup");
constexpr int MmaWarpGroups = size(TiledMma{}) / NumThreadsPerWarpGroup;
@@ -611,7 +610,7 @@ struct CollectiveMma<
k_tile_count -= prologue_mma_count;
smem_pipe_release.advance(k_tile_count);
// Wait on all GMMAs to complete
warpgroup_wait<0>();
@@ -690,9 +689,9 @@ struct CollectiveMma<
InternalElementB const* ptr_B = nullptr;
Tensor tensor_b = make_tensor(ptr_B, make_shape(N,K,Int<1>{}), mainloop_params.dB[next_group]);
cute::detail::fill_tma_gmem_shape_stride(mainloop_params.tma_load_a, tensor_a,
cute::detail::fill_tma_gmem_shape_stride(mainloop_params.tma_load_a, tensor_a,
prob_shape_A, prob_stride_A);
cute::detail::fill_tma_gmem_shape_stride(mainloop_params.tma_load_b, tensor_b,
cute::detail::fill_tma_gmem_shape_stride(mainloop_params.tma_load_b, tensor_b,
prob_shape_B, prob_stride_B);
// Convert strides to byte strides
include/cutlass/gemm/collective/sm90_mma_array_tma_gmma_ss_warpspecialized_fp8.hpp
+5 -6
View File
@@ -42,7 +42,6 @@
#include "cute/algorithm/functional.hpp"
#include "cute/atom/mma_atom.hpp"
#include "cute/algorithm/gemm.hpp"
#include "cute/tensor_predicate.hpp"
#include "cute/tensor.hpp"
#include "cute/numeric/arithmetic_tuple.hpp"
@@ -490,10 +489,10 @@ struct CollectiveMma<
// Layout of warp group to thread mapping
static_assert(stride<0>(typename TiledMma::ALayout{}) == 0 and
static_assert(stride<0>(typename TiledMma::ALayout{}) == 0 and
stride<0>(typename TiledMma::BLayout{}) == 0 and
size<0>(typename TiledMma::ALayout{}) == NumThreadsPerWarpGroup and
size<0>(typename TiledMma::BLayout{}) == NumThreadsPerWarpGroup,
size<0>(typename TiledMma::BLayout{}) == NumThreadsPerWarpGroup,
"Stride of the first mode must be 0 and the size of the mode must be NumThreadsPerWarpGroup");
constexpr int MmaWarpGroups = size(TiledMma{}) / NumThreadsPerWarpGroup;
@@ -620,7 +619,7 @@ struct CollectiveMma<
k_tile_count -= prologue_mma_count;
smem_pipe_release.advance(k_tile_count);
// Wait on all GMMAs to complete
warpgroup_wait<0>();
@@ -699,9 +698,9 @@ struct CollectiveMma<
ElementB const* ptr_B = nullptr;
Tensor tensor_b = make_tensor(ptr_B, make_shape(N,K,Int<1>{}), mainloop_params.dB[next_group]);
cute::detail::fill_tma_gmem_shape_stride(mainloop_params.tma_load_a, tensor_a,
cute::detail::fill_tma_gmem_shape_stride(mainloop_params.tma_load_a, tensor_a,
prob_shape_A, prob_stride_A);
cute::detail::fill_tma_gmem_shape_stride(mainloop_params.tma_load_b, tensor_b,
cute::detail::fill_tma_gmem_shape_stride(mainloop_params.tma_load_b, tensor_b,
prob_shape_B, prob_stride_B);
// Convert strides to byte strides
include/cutlass/gemm/collective/sm90_mma_array_tma_gmma_ss_warpspecialized_fp8_blockwise_scaling.hpp
+31 -32
View File
@@ -43,7 +43,6 @@
#include "cute/algorithm/functional.hpp"
#include "cute/atom/mma_atom.hpp"
#include "cute/algorithm/gemm.hpp"
#include "cute/tensor_predicate.hpp"
#include "cute/numeric/arithmetic_tuple.hpp"
#include "cutlass/detail/blockwise_scale_layout.hpp"
@@ -168,11 +167,11 @@ struct CollectiveMma<
make_shape(shape<1>(TileShape{}), shape<2>(TileShape{}), Int<DispatchPolicy::Stages>{}),
cute::conditional_t< ::cutlass::gemm::detail::is_major<0,StrideB>(), Step<_2,_1,_3>, Step<_1,_2,_3>>{}));
// Block scaling gmem-to-smem copy atom
// Block scaling gmem-to-smem copy atom
// we can have partial tiles in M or N, so don't vectorize those loads
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;
@@ -265,7 +264,7 @@ struct CollectiveMma<
InternalElementB const** ptr_B;
StrideB dB;
// Block scaling factors for A and B
ElementBlockScale const** ptr_SFA;
ElementBlockScale const** ptr_SFA;
LayoutSFA layout_SFA;
ElementBlockScale const** ptr_SFB;
LayoutSFB layout_SFB;
@@ -423,9 +422,9 @@ struct CollectiveMma<
// Make the tiled views of scale tensors
Tensor mSFA_mkl = make_tensor(make_gmem_ptr(ptr_SFA),
Tensor mSFA_mkl = make_tensor(make_gmem_ptr(ptr_SFA),
ScaleConfig::tile_atom_to_shape_SFA(make_shape(M, N, K, init_L))); // (scale_m,k,l)
Tensor mSFB_nkl = make_tensor(make_gmem_ptr(ptr_SFB),
Tensor mSFB_nkl = make_tensor(make_gmem_ptr(ptr_SFB),
ScaleConfig::tile_atom_to_shape_SFB(make_shape(M, N, K, init_L))); // (scale_n,k,l)
return cute::make_tuple(gA_mkl, gB_nkl, mSFA_mkl, mSFB_nkl);
@@ -443,7 +442,7 @@ struct CollectiveMma<
CUTLASS_DEVICE void
load(
Params const& mainloop_params,
MainloopPipeline pipeline,
MainloopPipeline pipeline,
PipelineState smem_pipe_write,
cute::tuple<TensorA, TensorB, TensorScaleA, TensorScaleB> const& load_inputs,
cute::tuple<TensorMapA, TensorMapB> const& input_tensormaps,
@@ -457,9 +456,9 @@ struct CollectiveMma<
if (lane_predicate) {
Tensor sA = make_tensor(make_smem_ptr(shared_tensors.smem_A.data()), SmemLayoutA{}); // (BLK_M,BLK_K,PIPE)
Tensor sB = make_tensor(make_smem_ptr(shared_tensors.smem_B.data()), SmemLayoutB{}); // (BLK_N,BLK_K,PIPE)
Tensor sSFA = make_tensor(cute::make_smem_ptr(shared_tensors.smem_SFA.data()),
Tensor sSFA = make_tensor(cute::make_smem_ptr(shared_tensors.smem_SFA.data()),
SmemLayoutSFA{}); // (BLK_M,BLK_K,P)
Tensor sSFB = make_tensor(cute::make_smem_ptr(shared_tensors.smem_SFB.data()),
Tensor sSFB = make_tensor(cute::make_smem_ptr(shared_tensors.smem_SFB.data()),
SmemLayoutSFB{}); // (BLK_N,BLK_K,P)
//
@@ -561,9 +560,9 @@ struct CollectiveMma<
// Issue the epilogue waits
if (lane_predicate) {
// This helps avoid early exit of blocks in Cluster.
// Waits for all stages to either be released (all
// Waits for all stages to either be released (all
// Consumer UNLOCKs), or if the stage was never used
// then it would just be acquired since the phase was
// then it would just be acquired since the phase was
// still inverted from make_producer_start_state.
pipeline.producer_tail(smem_pipe_write);
}
@@ -579,11 +578,11 @@ struct CollectiveMma<
CUTLASS_DEVICE void
load_auxiliary(
Params const& mainloop_params,
MainloopPipeline pipeline,
MainloopPipeline pipeline,
PipelineState smem_pipe_write,
cute::tuple<TensorA,
TensorB,
TensorSFA,
cute::tuple<TensorA,
TensorB,
TensorSFA,
TensorSFB> const& load_inputs,
BlockCoord const& blk_coord,
KTileIterator k_tile_iter, int k_tile_count,
@@ -591,9 +590,9 @@ struct CollectiveMma<
uint32_t block_rank_in_cluster,
TensorStorage& shared_tensors) {
int lane_predicate = cute::elect_one_sync();
Tensor sSFA = make_tensor(cute::make_smem_ptr(shared_tensors.smem_SFA.data()),
Tensor sSFA = make_tensor(cute::make_smem_ptr(shared_tensors.smem_SFA.data()),
SmemLayoutSFA{}); // (BLK_M,BLK_K,P)
Tensor sSFB = make_tensor(cute::make_smem_ptr(shared_tensors.smem_SFB.data()),
Tensor sSFB = make_tensor(cute::make_smem_ptr(shared_tensors.smem_SFB.data()),
SmemLayoutSFB{}); // (BLK_N,BLK_K,P)
// Partition the inputs based on the current block coordinates.
@@ -741,22 +740,22 @@ struct CollectiveMma<
// Block scaling
Tensor sSFA = make_tensor(cute::make_smem_ptr(shared_tensors.smem_SFA.data()), make_layout(
make_shape(shape<0>(SmemLayoutSFA{}),
get<1>(TileShape{}),
make_shape(shape<1>(SmemLayoutSFA{}),
make_shape(shape<0>(SmemLayoutSFA{}),
get<1>(TileShape{}),
make_shape(shape<1>(SmemLayoutSFA{}),
shape<2>(SmemLayoutSFA{}))),
make_stride(stride<0>(SmemLayoutSFA{}), _0{},
make_stride(stride<1>(SmemLayoutSFA{}),
make_stride(stride<0>(SmemLayoutSFA{}), _0{},
make_stride(stride<1>(SmemLayoutSFA{}),
stride<2>(SmemLayoutSFA{})))
)); // (BLK_M,BLK_N,(BLK_K,P))
Tensor sSFB = make_tensor(cute::make_smem_ptr(shared_tensors.smem_SFB.data()), make_layout(
make_shape(get<0>(TileShape{}),
shape<0>(SmemLayoutSFB{}),
make_shape(shape<1>(SmemLayoutSFB{}),
make_shape(get<0>(TileShape{}),
shape<0>(SmemLayoutSFB{}),
make_shape(shape<1>(SmemLayoutSFB{}),
shape<2>(SmemLayoutSFB{}))),
make_stride(_0{},
stride<0>(SmemLayoutSFB{}),
make_stride(stride<1>(SmemLayoutSFB{}),
make_stride(_0{},
stride<0>(SmemLayoutSFB{}),
make_stride(stride<1>(SmemLayoutSFB{}),
stride<2>(SmemLayoutSFB{})))
)); // (BLK_M,BLK_N,(BLK_K,P))
@@ -767,10 +766,10 @@ struct CollectiveMma<
// Layout of warp group to thread mapping
static_assert(stride<0>(typename TiledMma::ALayout{}) == 0 and
static_assert(stride<0>(typename TiledMma::ALayout{}) == 0 and
stride<0>(typename TiledMma::BLayout{}) == 0 and
size<0>(typename TiledMma::ALayout{}) == NumThreadsPerWarpGroup and
size<0>(typename TiledMma::BLayout{}) == NumThreadsPerWarpGroup,
size<0>(typename TiledMma::BLayout{}) == NumThreadsPerWarpGroup,
"Stride of the first mode must be 0 and the size of the mode must be NumThreadsPerWarpGroup");
constexpr int MmaWarpGroups = size(TiledMma{}) / NumThreadsPerWarpGroup;
@@ -1139,9 +1138,9 @@ struct CollectiveMma<
InternalElementB const* ptr_B = nullptr;
Tensor tensor_b = make_tensor(ptr_B, make_shape(N,K,Int<1>{}), mainloop_params.dB[next_group]);
cute::detail::fill_tma_gmem_shape_stride(mainloop_params.tma_load_a, tensor_a,
cute::detail::fill_tma_gmem_shape_stride(mainloop_params.tma_load_a, tensor_a,
prob_shape_A, prob_stride_A);
cute::detail::fill_tma_gmem_shape_stride(mainloop_params.tma_load_b, tensor_b,
cute::detail::fill_tma_gmem_shape_stride(mainloop_params.tma_load_b, tensor_b,
prob_shape_B, prob_stride_B);
// Convert strides to byte strides
include/cutlass/gemm/collective/sm90_mma_multistage_gmma_rs_warpspecialized.hpp
+18 -19
View File
@@ -42,7 +42,6 @@
#include "cute/algorithm/functional.hpp"
#include "cute/atom/mma_atom.hpp"
#include "cute/algorithm/gemm.hpp"
#include "cute/tensor_predicate.hpp"
#include "cute/numeric/arithmetic_tuple.hpp"
/////////////////////////////////////////////////////////////////////////////////////////////////
@@ -162,7 +161,7 @@ struct CollectiveMma<
static constexpr bool TransposeB = !IsInputSizeTwoBytes && IsLayoutAmnBmn;
using TransposeOperandB = decltype(cutlass::transform::collective::detail::make_transpose_operand_b(
0, 0, TiledMma{}, SmemLayoutB{}, InternalSmemLayoutAtomB{},
InternalElementB{}, cute::bool_constant<TransposeB>{}));
InternalElementB{}, cute::bool_constant<TransposeB>{}));
static_assert(DispatchPolicy::Stages >= 2, "Specialization requires Stages set to value 2 or more.");
static_assert(not cute::is_base_of<cute::GMMA::DescriptorIterator, typename TiledMma::FrgTypeA>::value &&
@@ -187,7 +186,7 @@ struct CollectiveMma<
struct SharedStorage
{
struct TensorStorage : cute::aligned_struct<256, _0> {
struct TensorStorage : cute::aligned_struct<256, _0> {
cute::array_aligned<typename TiledMma::ValTypeA, cute::cosize_v<SmemLayoutA>, 256> smem_A;
cute::array_aligned<typename TiledMma::ValTypeB, cute::cosize_v<SmemLayoutB>, 256> smem_B;
} tensors;
@@ -264,7 +263,7 @@ struct CollectiveMma<
static constexpr int K_PIPE_MAX = DispatchPolicy::Stages;
static constexpr int K_PIPE_MMAS = 1;
/// Perform a collective-scoped matrix multiply-accumulate
/// Producer Perspective
template <
@@ -275,7 +274,7 @@ struct CollectiveMma<
>
CUTLASS_DEVICE void
load(
MainloopPipeline pipeline,
MainloopPipeline pipeline,
PipelineState smem_pipe_write,
TensorA const& gA_in,
TensorB const& gB_in,
@@ -358,7 +357,7 @@ struct CollectiveMma<
clear(tBsB(_,_,k,write_stage));
}
}
++k_tile_iter;
--k_tile_count;
@@ -392,13 +391,13 @@ struct CollectiveMma<
/// Perform a Producer Epilogue to prevent early exit of blocks in a Cluster
CUTLASS_DEVICE void
load_tail(
MainloopPipeline pipeline,
MainloopPipeline pipeline,
PipelineState smem_pipe_write) {
// Issue the epilogue waits
/* This helps avoid early exit of blocks in Cluster
* Waits for all stages to either be released (all
* 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
* then would just be acquired since the phase was
* still inverted from make_producer_start_state
*/
pipeline.producer_tail(smem_pipe_write);
@@ -432,7 +431,7 @@ struct CollectiveMma<
// Obtain warp index
int warp_idx = canonical_warp_idx_sync();
[[maybe_unused]] int warp_group_thread_idx = thread_idx % 128;
Tensor sA_ = make_tensor(make_smem_ptr(shared_tensors.smem_A.data()), SmemLayoutA{}); // (BLK_M,BLK_K,PIPE)
Tensor sA = as_position_independent_swizzle_tensor(sA_); // (BLK_M,BLK_K,PIPE)
Tensor sB_ = make_tensor(make_smem_ptr(shared_tensors.smem_B.data()), SmemLayoutB{}); // (BLK_N,BLK_K,PIPE)
@@ -448,11 +447,11 @@ struct CollectiveMma<
// Layout of warp group to thread mapping
static_assert(stride<0>(typename TiledMma::BLayout{}) == 0 and
size<0>(typename TiledMma::BLayout{}) == NumThreadsPerWarpGroup,
size<0>(typename TiledMma::BLayout{}) == NumThreadsPerWarpGroup,
"Stride of the first mode must be 0 and the size of the mode must be NumThreadsPerWarpGroup");
constexpr int MmaWarpGroups = size(TiledMma{}) / NumThreadsPerWarpGroup;
Layout warp_group_thread_layout = make_layout(Int<MmaWarpGroups>{},
Layout warp_group_thread_layout = make_layout(Int<MmaWarpGroups>{},
Int<NumThreadsPerWarpGroup>{});
int warp_group_idx = __shfl_sync(0xFFFFFFFF, thread_idx / NumThreadsPerWarpGroup, 0);
@@ -499,8 +498,8 @@ struct CollectiveMma<
tiled_mma.accumulate_ = GMMA::ScaleOut::Zero;
TransposeOperandB transpose = cutlass::transform::collective::detail::make_transpose_operand_b(
warp_idx, warp_group_thread_idx, tiled_mma, SmemLayoutB{},
InternalSmemLayoutAtomB{}, InternalElementB{},
warp_idx, warp_group_thread_idx, tiled_mma, SmemLayoutB{},
InternalSmemLayoutAtomB{}, InternalElementB{},
cute::bool_constant<TransposeB>{});
warpgroup_fence_operand(accum);
@@ -535,8 +534,8 @@ struct CollectiveMma<
}
warpgroup_wait<2>();
if (k_tile_count - 1 > 0) {
if (!skip_wait) {
pipeline.consumer_wait(smem_pipe_read);
@@ -589,7 +588,7 @@ struct CollectiveMma<
transpose(sB, gmma_sB, read_stage, 1);
}
}
warpgroup_arrive();
// (V,M) x (V,N) => (V,M,N)
cute::gemm(tiled_mma, tCrA(_,_,k_block), tCrB(_,_,k_block,read_stage), accum);
@@ -638,7 +637,7 @@ struct CollectiveMma<
++smem_pipe_release;
}
}
warpgroup_arrive();
// (V,M) x (V,N) => (V,M,N)
cute::gemm(tiled_mma, tCrA(_,_,size<2>(tCrA) - 1), tCrB(_,_,size<2>(tCrA) - 1,read_stage), accum);
@@ -659,7 +658,7 @@ struct CollectiveMma<
k_tile_count -= prologue_mma_count;
smem_pipe_release.advance(k_tile_count);
// Wait on all GMMAs to complete
warpgroup_wait<0>();
include/cutlass/gemm/collective/sm90_mma_multistage_gmma_ss_warpspecialized.hpp
+10 -11
View File
@@ -38,7 +38,6 @@
#include "cute/algorithm/functional.hpp"
#include "cute/atom/mma_atom.hpp"
#include "cute/algorithm/gemm.hpp"
#include "cute/tensor_predicate.hpp"
#include "cute/numeric/arithmetic_tuple.hpp"
#include "cutlass/pipeline/pipeline.hpp"
#include "cutlass/trace.h"
@@ -191,7 +190,7 @@ struct CollectiveMma<
static constexpr int K_PIPE_MAX = DispatchPolicy::Stages;
static constexpr int K_PIPE_MMAS = 1;
/// Perform a collective-scoped matrix multiply-accumulate
/// Producer Perspective
template <
@@ -202,7 +201,7 @@ struct CollectiveMma<
>
CUTLASS_DEVICE void
load(
MainloopPipeline pipeline,
MainloopPipeline pipeline,
PipelineState smem_pipe_write,
TensorA const& gA_in,
TensorB const& gB_in,
@@ -318,13 +317,13 @@ struct CollectiveMma<
/// Perform a Producer Epilogue to prevent early exit of blocks in a Cluster
CUTLASS_DEVICE void
load_tail(
MainloopPipeline pipeline,
MainloopPipeline pipeline,
PipelineState smem_pipe_write) {
// Issue the epilogue waits
/* This helps avoid early exit of blocks in Cluster
* Waits for all stages to either be released (all
* 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
* then would just be acquired since the phase was
* still inverted from make_producer_start_state
*/
pipeline.producer_tail(smem_pipe_write);
@@ -363,14 +362,14 @@ struct CollectiveMma<
// Layout of warp group to thread mapping
static_assert(stride<0>(typename TiledMma::ALayout{}) == 0 and
static_assert(stride<0>(typename TiledMma::ALayout{}) == 0 and
stride<0>(typename TiledMma::BLayout{}) == 0 and
size<0>(typename TiledMma::ALayout{}) == NumThreadsPerWarpGroup and
size<0>(typename TiledMma::BLayout{}) == NumThreadsPerWarpGroup,
size<0>(typename TiledMma::BLayout{}) == NumThreadsPerWarpGroup,
"Stride of the first mode must be 0 and the size of the mode must be NumThreadsPerWarpGroup");
constexpr int MmaWarpGroups = size(TiledMma{}) / NumThreadsPerWarpGroup;
Layout warp_group_thread_layout = make_layout(Int<MmaWarpGroups>{},
Layout warp_group_thread_layout = make_layout(Int<MmaWarpGroups>{},
Int<NumThreadsPerWarpGroup>{});
int warp_group_idx = __shfl_sync(0xFFFFFFFF, thread_idx / NumThreadsPerWarpGroup, 0);
@@ -461,7 +460,7 @@ struct CollectiveMma<
pipeline.consumer_wait(smem_pipe_read, barrier_token);
int read_stage = smem_pipe_read.index();
warpgroup_fence_operand(accum);
warpgroup_arrive();
// (V,M,K) x (V,N,K) => (V,M,N)
@@ -491,7 +490,7 @@ struct CollectiveMma<
k_tile_count -= prologue_mma_count;
smem_pipe_release.advance(k_tile_count);
// Wait on all GMMAs to complete
warpgroup_wait<0>();
include/cutlass/gemm/collective/sm90_mma_tma_gmma_rs_warpspecialized.hpp
+24 -25
View File
@@ -45,7 +45,6 @@
#include "cute/algorithm/functional.hpp"
#include "cute/atom/mma_atom.hpp"
#include "cute/algorithm/gemm.hpp"
#include "cute/tensor_predicate.hpp"
#include "cute/numeric/arithmetic_tuple.hpp"
/////////////////////////////////////////////////////////////////////////////////////////////////
@@ -170,7 +169,7 @@ struct CollectiveMma<
static constexpr bool TransposeB = !IsInputSizeTwoBytes && IsLayoutAmnBmn;
using TransposeOperandB = decltype(cutlass::transform::collective::detail::make_transpose_operand_b(
0, 0, TiledMma{}, SmemLayoutB{}, InternalSmemLayoutAtomB{},
InternalElementB{}, cute::bool_constant<TransposeB>{}));
InternalElementB{}, cute::bool_constant<TransposeB>{}));
static_assert(DispatchPolicy::Stages >= 2, "Specialization requires Stages set to value 2 or more.");
static_assert(not cute::is_base_of<cute::GMMA::DescriptorIterator, typename TiledMma::FrgTypeA>::value &&
@@ -207,8 +206,8 @@ struct CollectiveMma<
static_assert(!uses_universal_transposition(),
"Warp specialized ARF kernels have not supported universal B transposition yet.");
static constexpr size_t SmemAlignmentA = cutlass::detail::alignment_for_swizzle(SmemLayoutA{});
static constexpr size_t SmemAlignmentA = cutlass::detail::alignment_for_swizzle(SmemLayoutA{});
static constexpr size_t SmemAlignmentB = cutlass::detail::alignment_for_swizzle(SmemLayoutB{});
@@ -216,7 +215,7 @@ struct CollectiveMma<
struct SharedStorage
{
struct TensorStorage : cute::aligned_struct<cute::max(SmemAlignmentA, SmemAlignmentB), _0> {
struct TensorStorage : cute::aligned_struct<cute::max(SmemAlignmentA, SmemAlignmentB), _0> {
cute::array_aligned<typename TiledMma::ValTypeA, cute::cosize_v<SmemLayoutA>, SmemAlignmentA> smem_A;
cute::array_aligned<typename TiledMma::ValTypeB, cute::cosize_v<SmemLayoutB>, SmemAlignmentB> smem_B;
} tensors;
@@ -330,7 +329,7 @@ struct CollectiveMma<
constexpr int tma_alignment_bits = 128;
auto problem_shape_MNKL = append<4>(problem_shape, 1);
auto [M,N,K,L] = problem_shape_MNKL;
bool implementable = true;
constexpr int min_tma_aligned_elements_A = tma_alignment_bits / cutlass::sizeof_bits<ElementA>::value;
implementable = implementable && cutlass::detail::check_alignment<min_tma_aligned_elements_A>(cute::make_shape(M,K,L), StrideA{});
@@ -410,7 +409,7 @@ struct CollectiveMma<
//
// Prepare the TMA loads for A and B
//
constexpr uint32_t cluster_shape_x = get<0>(ClusterShape());
uint2 cluster_local_block_id = {block_rank_in_cluster % cluster_shape_x, block_rank_in_cluster / cluster_shape_x};
@@ -483,9 +482,9 @@ struct CollectiveMma<
// Issue the epilogue waits
if (lane_predicate) {
/* This helps avoid early exit of blocks in Cluster
* Waits for all stages to either be released (all
* 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
* then would just be acquired since the phase was
* still inverted from make_producer_start_state
*/
pipeline.producer_tail(smem_pipe_write);
@@ -518,15 +517,15 @@ struct CollectiveMma<
// Obtain warp index
int warp_idx = canonical_warp_idx_sync();
[[maybe_unused]] int warp_group_thread_idx = thread_idx % 128;
Tensor sA_ = make_tensor(make_smem_ptr(shared_tensors.smem_A.data()), SmemLayoutA{}); // (BLK_M,BLK_K,PIPE)
Tensor sA = as_position_independent_swizzle_tensor(sA_); // (BLK_M,BLK_K,PIPE)
Tensor sB_ = make_tensor(make_smem_ptr(shared_tensors.smem_B.data()), SmemLayoutB{}); // (BLK_N,BLK_K,PIPE)
Tensor sB = as_position_independent_swizzle_tensor(sB_); // (BLK_M,BLK_K,PIPE)
// If TransposeB, GMMA will read from transposed B layout SMEM
Tensor gmma_sB_position_dependent = make_tensor(make_smem_ptr(shared_tensors.smem_B.data()),
Tensor gmma_sB_position_dependent = make_tensor(make_smem_ptr(shared_tensors.smem_B.data()),
GmmaSmemLayoutB{}); // (BLK_N,BLK_K,PIPE)
Tensor gmma_sB = as_position_independent_swizzle_tensor(gmma_sB_position_dependent); // (BLK_N,BLK_K,PIPE)
@@ -537,11 +536,11 @@ struct CollectiveMma<
// Layout of warp group to thread mapping
static_assert(stride<0>(typename TiledMma::BLayout{}) == 0 and
size<0>(typename TiledMma::BLayout{}) == NumThreadsPerWarpGroup,
size<0>(typename TiledMma::BLayout{}) == NumThreadsPerWarpGroup,
"Stride of the first mode must be 0 and the size of the mode must be NumThreadsPerWarpGroup");
constexpr int MmaWarpGroups = size(TiledMma{}) / NumThreadsPerWarpGroup;
Layout warp_group_thread_layout = make_layout(Int<MmaWarpGroups>{},
Layout warp_group_thread_layout = make_layout(Int<MmaWarpGroups>{},
Int<NumThreadsPerWarpGroup>{});
int warp_group_idx = __shfl_sync(0xFFFFFFFF, thread_idx / NumThreadsPerWarpGroup, 0);
@@ -590,12 +589,12 @@ struct CollectiveMma<
tiled_mma.accumulate_ = GMMA::ScaleOut::Zero;
TransposeOperandB transpose = cutlass::transform::collective::detail::make_transpose_operand_b(
warp_idx, warp_group_thread_idx, tiled_mma, SmemLayoutB{},
InternalSmemLayoutAtomB{}, InternalElementB{},
warp_idx, warp_group_thread_idx, tiled_mma, SmemLayoutB{},
InternalSmemLayoutAtomB{}, InternalElementB{},
cute::bool_constant<TransposeB>{});
warpgroup_fence_operand(accum);
ConsumerToken barrier_token = {BarrierStatus::WaitAgain};
// first k tile
{
@@ -611,7 +610,7 @@ struct CollectiveMma<
copy(smem_tiled_copy_A, tCsA_copy_view(_,_,0,read_stage), tCrA_copy_view(_,_,0));
// transpose B operand in SMEM
transpose(sB, gmma_sB, read_stage, 0);
// Unroll the K mode manually to set scale D to 1
CUTLASS_PRAGMA_UNROLL
for (int k_block = 0; k_block < size<2>(tCrA) - 1; ++k_block) {
@@ -628,7 +627,7 @@ struct CollectiveMma<
}
warpgroup_wait<2>();
warpgroup_arrive();
// (V,M) x (V,N) => (V,M,N)
cute::gemm(tiled_mma, tCrA(_,_,size<2>(tCrA) - 1), tCrB(_,_,size<2>(tCrA) - 1,read_stage), accum);
@@ -667,14 +666,14 @@ struct CollectiveMma<
copy(smem_tiled_copy_A, tCsA_copy_view(_,_,0,smem_pipe_read.index()), tCrA_copy_view(_,_,0));
// transpose B operand in SMEM
transpose(sB, gmma_sB, smem_pipe_read.index(), 0);
}
}
else {
copy(smem_tiled_copy_A, tCsA_copy_view(_,_,k_block + 1,read_stage), tCrA_copy_view(_,_,k_block + 1));
// transpose B operand in SMEM
transpose.synchronize(k_block); // make transpose of k_block available
transpose(sB, gmma_sB, read_stage, k_block + 1);
}
warpgroup_arrive();
// (V,M) x (V,N) => (V,M,N)
cute::gemm(tiled_mma, tCrA(_,_,k_block), tCrB(_,_,k_block,read_stage), accum);
@@ -700,7 +699,7 @@ struct CollectiveMma<
int read_stage = smem_pipe_read.index();
warpgroup_fence_operand(accum);
// Unroll the K mode manually to set scale D to 1
CUTLASS_PRAGMA_UNROLL
for (int k_block = 0; k_block < size<2>(tCrA) - 1; ++k_block) {
@@ -719,7 +718,7 @@ struct CollectiveMma<
++smem_pipe_release;
}
}
warpgroup_arrive();
// (V,M) x (V,N) => (V,M,N)
cute::gemm(tiled_mma, tCrA(_,_,size<2>(tCrA) - 1), tCrB(_,_,size<2>(tCrA) - 1,read_stage), accum);
@@ -728,7 +727,7 @@ struct CollectiveMma<
warpgroup_fence_operand(accum);
}
/// Perform a Consumer Epilogue to release all buffers
CUTLASS_DEVICE void
mma_tail(MainloopPipeline pipeline, PipelineState smem_pipe_release, int k_tile_count) {
@@ -737,7 +736,7 @@ struct CollectiveMma<
k_tile_count -= prologue_mma_count;
smem_pipe_release.advance(k_tile_count);
// Wait on all GMMAs to complete
warpgroup_wait<0>();
include/cutlass/gemm/collective/sm90_mma_tma_gmma_rs_warpspecialized_mixed_input.hpp
+57 -58
View File
@@ -50,7 +50,6 @@
#include "cute/atom/mma_atom.hpp"
#include "cute/atom/copy_traits_sm90_tma.hpp"
#include "cute/algorithm/gemm.hpp"
#include "cute/tensor_predicate.hpp"
#include "cute/numeric/arithmetic_tuple.hpp"
/////////////////////////////////////////////////////////////////////////////////////////////////
@@ -102,7 +101,7 @@ public:
ConvertAndScale,
ConvertAndScaleWithZero
};
//
// Type Aliases
//
@@ -112,10 +111,10 @@ public:
private:
template<class T> friend struct detail::MixedInputUtils;
using CollectiveType = CollectiveMma<DispatchPolicy, TileShape_,
ElementAOptionalTuple, StrideA_,
using CollectiveType = CollectiveMma<DispatchPolicy, TileShape_,
ElementAOptionalTuple, StrideA_,
ElementBOptionalTuple, StrideB_,
TiledMma_,
TiledMma_,
GmemTiledCopyA_, SmemLayoutAtomA_, SmemCopyAtomA_,
TransformA_,
GmemTiledCopyB_, SmemLayoutAtomB_, SmemCopyAtomB_,
@@ -128,7 +127,7 @@ private:
using ZeroB = detail::deduce_mixed_width_dtype_t<2, ElementBOptionalTuple>;
public:
static_assert(cute::is_tuple<ElementAOptionalTuple>::value ^ cute::is_tuple<ElementBOptionalTuple>::value,
static_assert(cute::is_tuple<ElementAOptionalTuple>::value ^ cute::is_tuple<ElementBOptionalTuple>::value,
"Either A OR B must be a tuple. It must take the from {ElementOperand, [ElementScale],"
"[ElementZero]}. Inputs in [] are optional.");
@@ -149,17 +148,17 @@ public:
using NonVoidStrideScale = cute::conditional_t<
cute::is_void_v<StrideScale>, cute::Stride<_1, int64_t, int64_t>, StrideScale>;
static_assert(( IsATransformed && (cutlass::gemm::detail::is_k_major<StrideA>() || is_layout<StrideA>::value)) ||
static_assert(( IsATransformed && (cutlass::gemm::detail::is_k_major<StrideA>() || is_layout<StrideA>::value)) ||
(!IsATransformed && (cutlass::gemm::detail::is_k_major<StrideB>() || is_layout<StrideB>::value)),
"The transformed type must be K-major.");
static_assert(( IsATransformed && (sizeof(ElementB) == 2)) ||
(!IsATransformed && (sizeof(ElementA) == 2)) ||
((cutlass::gemm::detail::is_k_major<StrideA>() || is_layout<StrideA>::value) &&
(cutlass::gemm::detail::is_k_major<StrideB>() || is_layout<StrideB>::value)),
((cutlass::gemm::detail::is_k_major<StrideA>() || is_layout<StrideA>::value) &&
(cutlass::gemm::detail::is_k_major<StrideB>() || is_layout<StrideB>::value)),
"The unscaled element must be 2 bytes OR both inputs must be K-major");
static_assert(cutlass::gemm::detail::is_mn_major<NonVoidStrideScale>(),
static_assert(cutlass::gemm::detail::is_mn_major<NonVoidStrideScale>(),
"Scale must be MN major [Col Major if A is scaled, Row Major if B is scaled].");
using CtaShape_MNK = decltype(shape_div(TileShape{}, ClusterShape{}));
@@ -185,7 +184,7 @@ public:
using SwappedSmemLayoutAtomB = cute::conditional_t<!SwapAB, SmemLayoutAtomB, SmemLayoutAtomA>;
using SwappedSmemCopyAtomA = cute::conditional_t<!SwapAB, SmemCopyAtomA, SmemCopyAtomB>;
using SwappedSmemCopyAtomB = cute::conditional_t<!SwapAB, SmemCopyAtomB, SmemCopyAtomA>;
// TMA converts f32 input to tf32 when copying from GMEM to SMEM
// For all other types, cast to size equivalent uint type to avoid any rounding by TMA.
static constexpr bool ConvertF32toTF32A = cute::is_same_v<float, ElementA>;
@@ -238,10 +237,10 @@ public:
using SmemLayoutA = decltype(detail::get_smem_layout<DispatchPolicy::Stages>(SwappedSmemLayoutAtomA{}, select<0,2>(TileShape{}), SwappedStrideA{}));
using SmemLayoutB = decltype(detail::get_smem_layout<DispatchPolicy::Stages>(SwappedSmemLayoutAtomB{}, select<1,2>(TileShape{}), SwappedStrideB{}));
// It is assumed that the scales and zero-points share the same smem layout
using SmemLayoutScale = decltype(tile_to_shape(
SmemLayoutAtomScale{},
SmemLayoutAtomScale{},
make_shape(shape<0>(ScaleTileShape{}), shape<1>(ScaleTileShape{}), Int<Stages>{}),
cute::conditional_t< ::cutlass::gemm::detail::is_major<0,NonVoidStrideScale>(), Step<_2,_1,_3>, Step<_1,_2,_3>>{}));
@@ -260,11 +259,11 @@ public:
static_assert(size<1>(SmemLayoutAtomScale{}) == 1, "size<1>(SmemLayoutAtomScale) must be 1.");
private:
static constexpr ConversionMode
static constexpr ConversionMode
get_conversion_mode() {
if constexpr (cute::is_void_v<ElementScale>) {
return ConversionMode::DirectConvert;
}
}
else if constexpr (cute::is_void_v<ElementZero>) {
return ConversionMode::ConvertAndScale;
}
@@ -279,7 +278,7 @@ public:
KernelConversionMode == ConversionMode::ConvertAndScaleWithZero;
static constexpr bool UseScaleLookupTable = KernelConversionMode == ConversionMode::ConvertAndScale &&
cutlass::detail::is_Array_v<ElementScale>;
static constexpr size_t SmemAlignmentA = cutlass::detail::alignment_for_swizzle(SmemLayoutA{});
static constexpr size_t SmemAlignmentA = cutlass::detail::alignment_for_swizzle(SmemLayoutA{});
static constexpr size_t SmemAlignmentB = cutlass::detail::alignment_for_swizzle(SmemLayoutB{});
@@ -424,7 +423,7 @@ public:
uint32_t tma_transaction_bytes = TmaTransactionBytesMK + TmaTransactionBytesNK;
if constexpr (KernelConversionMode == ConversionMode::DirectConvert) {
return { tma_load_a, tma_load_b, tma_load_scale, tma_load_zero, 0, 0, tma_transaction_bytes, 1, dA, dB };
}
}
else if constexpr (ModeHasScales) {
auto scale_k = ceil_div(K, args.group_size);
ElementScale const* ptr_S = args.ptr_S;
@@ -452,7 +451,7 @@ public:
} else {
static_assert(cutlass::detail::dependent_false<KernelSchedule>, "Conversion mode not handled in to_underlying_arguments.");
}
}
}
else {
static_assert(cutlass::detail::dependent_false<KernelSchedule>, "Conversion mode not handled in to_underlying_arguments.");
}
@@ -480,7 +479,7 @@ public:
if constexpr (KernelConversionMode == ConversionMode::DirectConvert) {
check_mode_args = check_mode_args && (args.ptr_S == nullptr);
check_mode_args = check_mode_args && (args.ptr_Z == nullptr);
}
}
else if constexpr (ModeHasScales) {
const int scale_mn = SwapAB ? N : M;
const int scale_k = ceil_div(K, args.group_size);
@@ -497,7 +496,7 @@ public:
constexpr int min_tma_aligned_elements_zero = tma_alignment_bits / cutlass::sizeof_bits<ElementZero>::value;
check_aligned_Z = cutlass::detail::check_alignment<min_tma_aligned_elements_zero>(cute::make_shape(scale_mn,scale_k,L), args.dS);
check_mode_args = check_mode_args && (args.ptr_Z != nullptr);
}
}
else {
static_assert(cutlass::detail::dependent_false<KernelSchedule>, "Conversion mode not handled in can_implement.");
}
@@ -539,18 +538,18 @@ public:
if constexpr (KernelConversionMode == ConversionMode::DirectConvert) {
// Nothing extra to do
}
}
else if constexpr (KernelConversionMode == ConversionMode::ConvertAndScale) {
cute::prefetch_tma_descriptor(mainloop_params.tma_load_scale.get_tma_descriptor());
}
else if constexpr (KernelConversionMode == ConversionMode::ConvertAndScaleWithZero) {
cute::prefetch_tma_descriptor(mainloop_params.tma_load_scale.get_tma_descriptor());
cute::prefetch_tma_descriptor(mainloop_params.tma_load_zero.get_tma_descriptor());
}
}
else {
static_assert(cutlass::detail::dependent_false<KernelSchedule>, "Conversion mode not handled in TMA prefetch.");
}
}
/// Set up the data needed by this collective for load and mma.
@@ -577,7 +576,7 @@ public:
if constexpr (KernelConversionMode == ConversionMode::DirectConvert) {
return cute::make_tuple(gA_mkl, gB_nkl);
}
}
else if constexpr (ModeHasScales) {
auto scale_k = mainloop_params.scale_k;
Tensor mS_mkl = mainloop_params.tma_load_scale.get_tma_tensor(make_shape(M,scale_k,L)); // (m,scale_k,l)
@@ -593,11 +592,11 @@ public:
else {
static_assert(cutlass::detail::dependent_false<KernelSchedule>, "Conversion mode not handled in load_init.");
}
}
}
else {
static_assert(cutlass::detail::dependent_false<KernelSchedule>, "Conversion mode not handled in load_init.");
}
}
}
/// Perform a collective-scoped matrix multiply-accumulate
/// Producer Perspective
@@ -609,7 +608,7 @@ public:
CUTLASS_DEVICE void
load(
Params const& mainloop_params,
MainloopPipeline pipeline,
MainloopPipeline pipeline,
PipelineState smem_pipe_write,
cute::tuple<Ts...> const& load_inputs,
BlockCoord const& blk_coord,
@@ -619,13 +618,13 @@ public:
TensorStorage& shared_tensors) {
if constexpr (KernelConversionMode == ConversionMode::DirectConvert) {
static_assert(sizeof... (Ts) == 2, "Direct convert needs two inputs");
}
}
else if constexpr (KernelConversionMode == ConversionMode::ConvertAndScale) {
static_assert(sizeof... (Ts) == 3, "Scaled convert needs three inputs");
}
}
else if constexpr (KernelConversionMode == ConversionMode::ConvertAndScaleWithZero) {
static_assert(sizeof... (Ts) == 4, "Scaled and zero convert needs four inputs");
}
}
else {
static_assert(cutlass::detail::dependent_false<KernelSchedule>, "Conversion mode not handled in TMA load.");
}
@@ -638,7 +637,7 @@ public:
//
// Prepare the TMA loads for A, B and Scales
//
constexpr uint32_t cluster_shape_x = get<0>(ClusterShape());
uint2 cluster_local_block_id = {block_rank_in_cluster % cluster_shape_x, block_rank_in_cluster / cluster_shape_x};
@@ -717,7 +716,7 @@ public:
if constexpr (KernelConversionMode == ConversionMode::ConvertAndScale) {
// Nothing extra to do
}
}
else if constexpr (KernelConversionMode == ConversionMode::ConvertAndScaleWithZero) {
auto tZgZ = get<2>(extra_input_partitions);
auto tZsZ = get<3>(extra_input_partitions);
@@ -725,8 +724,8 @@ public:
}
else {
static_assert(cutlass::detail::dependent_false<KernelSchedule>, "Conversion mode not handled for TMA copy op.");
}
}
}
}
else {
static_assert(cutlass::detail::dependent_false<KernelSchedule>, "Conversion mode not handled for TMA copy op.");
}
@@ -744,9 +743,9 @@ public:
// Issue the epilogue waits
if (cute::elect_one_sync()) {
/* This helps avoid early exit of blocks in Cluster
* Waits for all stages to either be released (all
* 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
* then would just be acquired since the phase was
* still inverted from make_producer_start_state
*/
pipeline.producer_tail(smem_pipe_write);
@@ -779,10 +778,10 @@ public:
// Obtain warp index
int warp_idx = canonical_warp_idx_sync();
[[maybe_unused]] int warp_group_thread_idx = thread_idx % 128;
Tensor sA_ = make_tensor(make_smem_ptr(shared_tensors.smem_A.begin()), SmemLayoutA{}); // (BLK_M,BLK_K,PIPE)
Tensor sA = as_position_independent_swizzle_tensor(sA_); // (BLK_M,BLK_K,PIPE)
Tensor sB = make_tensor(make_smem_ptr(shared_tensors.smem_B.begin()), SmemLayoutB{}); // (BLK_N,BLK_K,PIPE)
//
@@ -792,11 +791,11 @@ public:
// Layout of warp group to thread mapping
static_assert(stride<0>(typename TiledMma::BLayout{}) == 0 and
size<0>(typename TiledMma::BLayout{}) == NumThreadsPerWarpGroup,
size<0>(typename TiledMma::BLayout{}) == NumThreadsPerWarpGroup,
"Stride of the first mode must be 0 and the size of the mode must be NumThreadsPerWarpGroup");
constexpr int MmaWarpGroups = size(TiledMma{}) / NumThreadsPerWarpGroup;
Layout warp_group_thread_layout = make_layout(Int<MmaWarpGroups>{},
Layout warp_group_thread_layout = make_layout(Int<MmaWarpGroups>{},
Int<NumThreadsPerWarpGroup>{});
int warp_group_idx = __shfl_sync(0xFFFFFFFF, thread_idx / NumThreadsPerWarpGroup, 0);
@@ -819,7 +818,7 @@ public:
return make_tensor_like<RealSwappedElementA>(tCsA(_,_,_,Int<0>{}));
}
}();
Tensor tCsB = mma_warpgroup_slice.partition_B(sB); // (MMA,MMA_N,MMA_K,PIPE)
Tensor tCrB = mma_warpgroup_slice.make_fragment_B(tCsB); // (MMA,MMA_N,MMA_K,PIPE)
@@ -871,14 +870,14 @@ public:
barrier_token = pipeline.consumer_try_wait(smem_pipe_read);
// copy smem->rmem for A operand
Utils::copy_tensors_MK(smem_tiled_copy_A, tCsA, tCrA_copy_view,
Utils::copy_tensors_MK(smem_tiled_copy_A, tCsA, tCrA_copy_view,
partitioned_extra_info, copy_partitions_extra_info, 0, read_stage);
if (K_BLOCK_MAX > 1) { // prefetch next block
Utils::copy_tensors_MK(smem_tiled_copy_A, tCsA, tCrA_copy_view,
Utils::copy_tensors_MK(smem_tiled_copy_A, tCsA, tCrA_copy_view,
partitioned_extra_info, copy_partitions_extra_info, 1, read_stage);
}
Utils::dequantize_A_kblock(tCrA_load, tCrA_mma, partitioned_extra_info, 0);
// Unroll the K mode manually to set scale D to 1
CUTLASS_PRAGMA_UNROLL
for (int k_block = 0; k_block < K_BLOCK_MAX; ++k_block) {
@@ -889,25 +888,25 @@ public:
warpgroup_commit_batch();
if (k_block < K_BLOCK_MAX - 2) { // prefetch next block
Utils::copy_tensors_MK(smem_tiled_copy_A, tCsA, tCrA_copy_view,
Utils::copy_tensors_MK(smem_tiled_copy_A, tCsA, tCrA_copy_view,
partitioned_extra_info, copy_partitions_extra_info, k_block + 2, read_stage);
}
if (k_block < K_BLOCK_MAX - 1) {
Utils::dequantize_A_kblock(tCrA_load, tCrA_mma, partitioned_extra_info, k_block + 1);
}
}
}
--k_tile_count;
if (k_tile_count > 0) {
// Wait for K_BLOCK_MAX - 1 to be in flight to ensure that it is safe to overwrite the A registers for the first mma.
pipeline.consumer_wait(smem_pipe_read, barrier_token);
Utils::copy_tensors_MK(smem_tiled_copy_A, tCsA, tCrA_copy_view,
Utils::copy_tensors_MK(smem_tiled_copy_A, tCsA, tCrA_copy_view,
partitioned_extra_info, copy_partitions_extra_info, 0, smem_pipe_read.index());
if (K_BLOCK_MAX > 1) { // prefetch next block
Utils::copy_tensors_MK(smem_tiled_copy_A, tCsA, tCrA_copy_view,
Utils::copy_tensors_MK(smem_tiled_copy_A, tCsA, tCrA_copy_view,
partitioned_extra_info, copy_partitions_extra_info, 1, smem_pipe_read.index());
}
warpgroup_wait<K_WAIT_MAX>();
warpgroup_wait<K_WAIT_MAX>();
Utils::dequantize_A_kblock(tCrA_load, tCrA_mma, partitioned_extra_info, 0);
}
}
@@ -932,7 +931,7 @@ public:
// Unroll the K mode manually to set scale D to 1
CUTLASS_PRAGMA_UNROLL
for (int k_block = 0; k_block < K_BLOCK_MAX; ++k_block) {
warpgroup_arrive();
// (V,M) x (V,N) => (V,M,N)
cute::gemm(tiled_mma, tCrA_mma(_,_,k_block), tCrB(_,_,k_block,read_stage), accum);
@@ -949,19 +948,19 @@ public:
barrier_token = pipeline.consumer_try_wait(smem_pipe_read);
}
if (k_block == K_BLOCK_MAX - 1) {
if (k_block == K_BLOCK_MAX - 1) {
pipeline.consumer_wait(smem_pipe_read, barrier_token);
Utils::copy_tensors_MK(smem_tiled_copy_A, tCsA, tCrA_copy_view,
Utils::copy_tensors_MK(smem_tiled_copy_A, tCsA, tCrA_copy_view,
partitioned_extra_info, copy_partitions_extra_info, 0, smem_pipe_read.index());
if (K_BLOCK_MAX > 1) { // prefetch next block
Utils::copy_tensors_MK(smem_tiled_copy_A, tCsA, tCrA_copy_view,
Utils::copy_tensors_MK(smem_tiled_copy_A, tCsA, tCrA_copy_view,
partitioned_extra_info, copy_partitions_extra_info, 1, smem_pipe_read.index());
}
Utils::dequantize_A_kblock(tCrA_load, tCrA_mma, partitioned_extra_info, 0);
}
}
else {
if (k_block < K_BLOCK_MAX - 2) { // prefetch next block
Utils::copy_tensors_MK(smem_tiled_copy_A, tCsA, tCrA_copy_view,
Utils::copy_tensors_MK(smem_tiled_copy_A, tCsA, tCrA_copy_view,
partitioned_extra_info, copy_partitions_extra_info, k_block + 2, read_stage);
}
Utils::dequantize_A_kblock(tCrA_load, tCrA_mma, partitioned_extra_info, k_block + 1);
@@ -981,7 +980,7 @@ public:
int read_stage = smem_pipe_read.index();
warpgroup_fence_operand(accum);
// Unroll the K mode manually to set scale D to 1
CUTLASS_PRAGMA_UNROLL
for (int k_block = 0; k_block < K_BLOCK_MAX; ++k_block) {
@@ -999,7 +998,7 @@ public:
}
if (k_block < K_BLOCK_MAX - 2) { // prefetch next block
Utils::copy_tensors_MK(smem_tiled_copy_A, tCsA, tCrA_copy_view,
Utils::copy_tensors_MK(smem_tiled_copy_A, tCsA, tCrA_copy_view,
partitioned_extra_info, copy_partitions_extra_info, k_block + 2, read_stage);
}
if (k_block < K_BLOCK_MAX - 1) {
@@ -1019,7 +1018,7 @@ public:
k_tile_count -= prologue_mma_count;
smem_pipe_release.advance(k_tile_count);
// Wait on all GMMAs to complete
warpgroup_wait<0>();
include/cutlass/gemm/collective/sm90_mma_tma_gmma_ss.hpp
+5 -6
View File
@@ -41,7 +41,6 @@
#include "cute/algorithm/functional.hpp"
#include "cute/atom/mma_atom.hpp"
#include "cute/algorithm/gemm.hpp"
#include "cute/tensor_predicate.hpp"
#include "cute/numeric/arithmetic_tuple.hpp"
/////////////////////////////////////////////////////////////////////////////////////////////////
@@ -230,7 +229,7 @@ struct CollectiveMma<
constexpr int tma_alignment_bits = 128;
auto problem_shape_MNKL = append<4>(problem_shape, 1);
auto [M,N,K,L] = problem_shape_MNKL;
bool implementable = true;
constexpr int min_tma_aligned_elements_A = tma_alignment_bits / cutlass::sizeof_bits<ElementA>::value;
implementable = implementable && cutlass::detail::check_alignment<min_tma_aligned_elements_A>(cute::make_shape(M,K,L), StrideA{});
@@ -401,14 +400,14 @@ struct CollectiveMma<
// Layout of warp group to thread mapping
static_assert(stride<0>(typename TiledMma::ALayout{}) == 0 and
static_assert(stride<0>(typename TiledMma::ALayout{}) == 0 and
stride<0>(typename TiledMma::BLayout{}) == 0 and
size<0>(typename TiledMma::ALayout{}) == NumThreadsPerWarpGroup and
size<0>(typename TiledMma::BLayout{}) == NumThreadsPerWarpGroup,
size<0>(typename TiledMma::BLayout{}) == NumThreadsPerWarpGroup,
"Stride of the first mode must be 0 and the size of the mode must be NumThreadsPerWarpGroup");
constexpr int MmaWarpGroups = size(TiledMma{}) / NumThreadsPerWarpGroup;
Layout warp_group_thread_layout = make_layout(Int<MmaWarpGroups>{},
Layout warp_group_thread_layout = make_layout(Int<MmaWarpGroups>{},
Int<NumThreadsPerWarpGroup>{});
int warp_group_idx = __shfl_sync(0xFFFFFFFF, thread_idx / NumThreadsPerWarpGroup, 0);
@@ -457,7 +456,7 @@ struct CollectiveMma<
}
CUTLASS_PRAGMA_UNROLL
for (int prologue_mma_count = min(K_PIPE_MMAS, k_tile_count) - 1;
for (int prologue_mma_count = min(K_PIPE_MMAS, k_tile_count) - 1;
prologue_mma_count > 0; --prologue_mma_count)
{
// WAIT on smem_pipe_read until it's data is available
include/cutlass/gemm/collective/sm90_mma_tma_gmma_ss_warpspecialized.hpp
+8 -9
View File
@@ -41,7 +41,6 @@
#include "cute/algorithm/functional.hpp"
#include "cute/atom/mma_atom.hpp"
#include "cute/algorithm/gemm.hpp"
#include "cute/tensor_predicate.hpp"
#include "cute/numeric/arithmetic_tuple.hpp"
/////////////////////////////////////////////////////////////////////////////////////////////////
@@ -115,7 +114,7 @@ struct CollectiveMma<
using PipelineParams = typename MainloopPipeline::Params;
// One threads per CTA are producers (1 for operand tile)
static constexpr int NumProducerThreadEvents = 1;
static constexpr int NumProducerThreadEvents = 1;
static_assert(cute::rank(SmemLayoutAtomA{}) == 2, "SmemLayoutAtom must be rank 2 (M/N, K)");
static_assert((size<0>(TileShape{}) % size<0>(SmemLayoutAtomA{})) == 0, "SmemLayoutAtom must evenly divide tile shape.");
@@ -248,7 +247,7 @@ struct CollectiveMma<
constexpr int tma_alignment_bits = 128;
auto problem_shape_MNKL = append<4>(problem_shape, 1);
auto [M,N,K,L] = problem_shape_MNKL;
bool implementable = true;
constexpr int min_tma_aligned_elements_A = tma_alignment_bits / cutlass::sizeof_bits<ElementA>::value;
implementable = implementable && cutlass::detail::check_alignment<min_tma_aligned_elements_A>(cute::make_shape(M,K,L), StrideA{});
@@ -400,9 +399,9 @@ struct CollectiveMma<
// Issue the epilogue waits
if (lane_predicate) {
/* This helps avoid early exit of blocks in Cluster
* Waits for all stages to either be released (all
* 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
* then would just be acquired since the phase was
* still inverted from make_producer_start_state
*/
pipeline.producer_tail(smem_pipe_write);
@@ -439,14 +438,14 @@ struct CollectiveMma<
// Layout of warp group to thread mapping
static_assert(stride<0>(typename TiledMma::ALayout{}) == 0 and
static_assert(stride<0>(typename TiledMma::ALayout{}) == 0 and
stride<0>(typename TiledMma::BLayout{}) == 0 and
size<0>(typename TiledMma::ALayout{}) == NumThreadsPerWarpGroup and
size<0>(typename TiledMma::BLayout{}) == NumThreadsPerWarpGroup,
size<0>(typename TiledMma::BLayout{}) == NumThreadsPerWarpGroup,
"Stride of the first mode must be 0 and the size of the mode must be NumThreadsPerWarpGroup");
constexpr int MmaWarpGroups = size(TiledMma{}) / NumThreadsPerWarpGroup;
Layout warp_group_thread_layout = make_layout(Int<MmaWarpGroups>{},
Layout warp_group_thread_layout = make_layout(Int<MmaWarpGroups>{},
Int<NumThreadsPerWarpGroup>{});
int warp_group_idx = __shfl_sync(0xFFFFFFFF, thread_idx / NumThreadsPerWarpGroup, 0);
@@ -567,7 +566,7 @@ struct CollectiveMma<
k_tile_count -= prologue_mma_count;
smem_pipe_release.advance(k_tile_count);
// Wait on all GMMAs to complete
warpgroup_wait<0>();
include/cutlass/gemm/collective/sm90_mma_tma_gmma_ss_warpspecialized_fp8.hpp
+5 -6
View File
@@ -42,7 +42,6 @@
#include "cute/algorithm/functional.hpp"
#include "cute/atom/mma_atom.hpp"
#include "cute/algorithm/gemm.hpp"
#include "cute/tensor_predicate.hpp"
#include "cute/tensor.hpp"
#include "cute/numeric/arithmetic_tuple.hpp"
@@ -244,7 +243,7 @@ struct CollectiveMma<
constexpr int tma_alignment_bits = 128;
auto problem_shape_MNKL = append<4>(problem_shape, 1);
auto [M,N,K,L] = problem_shape_MNKL;
bool implementable = true;
constexpr int min_tma_aligned_elements_A = tma_alignment_bits / cutlass::sizeof_bits<ElementA>::value;
implementable = implementable && cutlass::detail::check_alignment<min_tma_aligned_elements_A>(cute::make_shape(M,K,L), StrideA{});
@@ -437,17 +436,17 @@ struct CollectiveMma<
//
// Define C accumulators and A/B partitioning
//
// Layout of warp group to thread mapping
static_assert(stride<0>(typename TiledMma::ALayout{}) == 0 and
static_assert(stride<0>(typename TiledMma::ALayout{}) == 0 and
stride<0>(typename TiledMma::BLayout{}) == 0 and
size<0>(typename TiledMma::ALayout{}) == NumThreadsPerWarpGroup and
size<0>(typename TiledMma::BLayout{}) == NumThreadsPerWarpGroup,
size<0>(typename TiledMma::BLayout{}) == NumThreadsPerWarpGroup,
"Stride of the first mode must be 0 and the size of the mode must be NumThreadsPerWarpGroup");
constexpr int MmaWarpGroups = size(TiledMma{}) / NumThreadsPerWarpGroup;
Layout warp_group_thread_layout = make_layout(Int<MmaWarpGroups>{},
Layout warp_group_thread_layout = make_layout(Int<MmaWarpGroups>{},
Int<NumThreadsPerWarpGroup>{});
int warp_group_idx = __shfl_sync(0xFFFFFFFF, thread_idx / NumThreadsPerWarpGroup, 0);
include/cutlass/gemm/collective/sm90_mma_tma_gmma_ss_warpspecialized_fp8_blockwise_scaling.hpp
+17 -18
View File
@@ -42,7 +42,6 @@
#include "cute/algorithm/functional.hpp"
#include "cute/atom/mma_atom.hpp"
#include "cute/algorithm/gemm.hpp"
#include "cute/tensor_predicate.hpp"
#include "cute/numeric/arithmetic_tuple.hpp"
#include "cutlass/detail/blockwise_scale_layout.hpp"
@@ -166,7 +165,7 @@ struct CollectiveMma<
make_shape(shape<1>(TileShape{}), shape<2>(TileShape{}), Int<DispatchPolicy::Stages>{}),
cute::conditional_t< ::cutlass::gemm::detail::is_major<0,StrideB>(), Step<_2,_1,_3>, Step<_1,_2,_3>>{}));
// Block scaling gmem-to-smem copy atom
// Block scaling gmem-to-smem copy atom
// we can have partial tiles in M or N, so don't vectorize those loads
using CopyAtomSFA = Copy_Atom<SM80_CP_ASYNC_CACHEALWAYS<ElementBlockScale>, ElementBlockScale>;
using CopyAtomSFB = Copy_Atom<SM80_CP_ASYNC_CACHEALWAYS<ElementBlockScale>, ElementBlockScale>;
@@ -217,7 +216,7 @@ struct CollectiveMma<
StrideA dA;
ElementB const* ptr_B;
StrideB dB;
ElementBlockScale const* ptr_SFA;
ElementBlockScale const* ptr_SFA;
LayoutSFA layout_SFA;
ElementBlockScale const* ptr_SFB;
LayoutSFB layout_SFB;
@@ -607,7 +606,7 @@ struct CollectiveMma<
CUTLASS_DEVICE void
load_auxiliary(
Params const& mainloop_params,
MainloopPipeline pipeline,
MainloopPipeline pipeline,
PipelineState smem_pipe_write,
cute::tuple<TensorA, TensorB, TensorScaleA, TensorScaleB> const& load_inputs,
BlockCoord const& blk_coord,
@@ -639,7 +638,7 @@ struct CollectiveMma<
TiledCopy scale_copy_a = make_tiled_copy(CopyAtomSFA{},
Layout<Shape<_32>>{}, Layout<Shape<_1>>{});
TiledCopy scale_copy_b = make_tiled_copy(CopyAtomSFB{},
TiledCopy scale_copy_b = make_tiled_copy(CopyAtomSFB{},
Layout<Shape<_32>>{}, Layout<Shape<_1>>{});
ThrCopy thr_scale_copy_a = scale_copy_a.get_slice(thread_idx);
ThrCopy thr_scale_copy_b = scale_copy_b.get_slice(thread_idx);
@@ -778,21 +777,21 @@ struct CollectiveMma<
// Block scaling
Tensor sSFA = make_tensor(cute::make_smem_ptr(shared_tensors.smem_SFA.data()), make_layout(
make_shape(get<0>(shape(SmemLayoutSFA{})),
get<1>(TileShape{}),
make_shape(get<1>(shape(SmemLayoutSFA{})),
make_shape(get<0>(shape(SmemLayoutSFA{})),
get<1>(TileShape{}),
make_shape(get<1>(shape(SmemLayoutSFA{})),
get<2>(shape(SmemLayoutSFA{})))),
make_stride(get<0>(stride(SmemLayoutSFA{})), _0{},
make_stride(get<0>(stride(SmemLayoutSFA{})), _0{},
make_stride(get<1>(stride(SmemLayoutSFA{})), get<2>(stride(SmemLayoutSFA{}))))
)); // (BLK_M,BLK_N,(BLK_K,P))
Tensor sSFB = make_tensor(cute::make_smem_ptr(shared_tensors.smem_SFB.data()), make_layout(
make_shape(get<0>(TileShape{}),
get<0>(shape(SmemLayoutSFB{})),
make_shape(get<1>(shape(SmemLayoutSFB{})),
make_shape(get<0>(TileShape{}),
get<0>(shape(SmemLayoutSFB{})),
make_shape(get<1>(shape(SmemLayoutSFB{})),
get<2>(shape(SmemLayoutSFB{})))),
make_stride(_0{},
get<0>(stride(SmemLayoutSFB{})),
make_stride(get<1>(stride(SmemLayoutSFB{})),
make_stride(_0{},
get<0>(stride(SmemLayoutSFB{})),
make_stride(get<1>(stride(SmemLayoutSFB{})),
get<2>(stride(SmemLayoutSFB{}))))
)); // (BLK_M,BLK_N,(BLK_K,P))
@@ -802,14 +801,14 @@ struct CollectiveMma<
// Layout of warp group to thread mapping
static_assert(stride<0>(typename TiledMma::ALayout{}) == 0 and
static_assert(stride<0>(typename TiledMma::ALayout{}) == 0 and
stride<0>(typename TiledMma::BLayout{}) == 0 and
size<0>(typename TiledMma::ALayout{}) == NumThreadsPerWarpGroup and
size<0>(typename TiledMma::BLayout{}) == NumThreadsPerWarpGroup,
size<0>(typename TiledMma::BLayout{}) == NumThreadsPerWarpGroup,
"Stride of the first mode must be 0 and the size of the mode must be NumThreadsPerWarpGroup");
constexpr int MmaWarpGroups = size(TiledMma{}) / NumThreadsPerWarpGroup;
Layout warp_group_thread_layout = make_layout(Int<MmaWarpGroups>{},
Layout warp_group_thread_layout = make_layout(Int<MmaWarpGroups>{},
Int<NumThreadsPerWarpGroup>{});
int warp_group_idx = __shfl_sync(0xFFFFFFFF, thread_idx / NumThreadsPerWarpGroup, 0);
include/cutlass/gemm/collective/sm90_sparse_mma_tma_gmma_ss_warpspecialized.hpp
+10 -11
View File
@@ -42,7 +42,6 @@
#include "cute/algorithm/functional.hpp"
#include "cute/atom/mma_atom.hpp"
#include "cute/algorithm/gemm.hpp"
#include "cute/tensor_predicate.hpp"
#include "cute/numeric/arithmetic_tuple.hpp"
/////////////////////////////////////////////////////////////////////////////////////////////////
@@ -194,7 +193,7 @@ struct CollectiveMma<
cute::conditional_t<cute::is_same_v<ElementA, float>,
cutlass::tfloat32_t,
uint_bit_t<sizeof_bits_v<ElementAMmaRaw>>>>;
using TmaInternalElementB = cute::conditional_t<cute::is_same_v<float, ElementB>,
using TmaInternalElementB = cute::conditional_t<cute::is_same_v<float, ElementB>,
tfloat32_t,
uint_bit_t<sizeof_bits_v<ElementBMma>>>;
@@ -215,7 +214,7 @@ struct CollectiveMma<
static constexpr int K_PIPE_MAX = DispatchPolicy::Stages;
static constexpr int K_PIPE_MMAS = 0;
static constexpr uint32_t TmaTransactionBytesMK =
static constexpr uint32_t TmaTransactionBytesMK =
cutlass::bits_to_bytes(cosize(take<0,2>(SmemLayoutA{})) * cute::sizeof_bits_v<ElementAMma>) +
cutlass::bits_to_bytes(cosize(take<0,2>(SmemLayoutE{})) * cute::sizeof_bits_v<ElementEMma>);
@@ -332,7 +331,7 @@ struct CollectiveMma<
constexpr int min_tma_aligned_elements_B = tma_alignment_bits / cutlass::sizeof_bits<ElementB>::value;
auto problem_shape_MNKL = append<4>(problem_shape, 1);
auto [M,N,K,L] = problem_shape_MNKL;
bool size_check = true;
// Check Alignment A
if constexpr (is_A_mn_major) {
@@ -405,7 +404,7 @@ struct CollectiveMma<
CUTLASS_DEVICE void
load(
Params const& mainloop_params,
MainloopPipeline pipeline,
MainloopPipeline pipeline,
PipelineState smem_pipe_write,
cute::tuple<TensorA, TensorB, TensorE> const& load_inputs,
BlockCoord const& blk_coord,
@@ -485,9 +484,9 @@ struct CollectiveMma<
// Issue the epilogue waits
if (lane_predicate) {
/* This helps avoid early exit of blocks in Cluster
* Waits for all stages to either be released (all
* 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
* then would just be acquired since the phase was
* still inverted from make_producer_start_state
*/
pipeline.producer_tail(smem_pipe_write);
@@ -523,14 +522,14 @@ struct CollectiveMma<
// Layout of warp group to thread mapping
static_assert(stride<0>(typename TiledMma::ALayout{}) == 0 and
static_assert(stride<0>(typename TiledMma::ALayout{}) == 0 and
stride<0>(typename TiledMma::BLayout{}) == 0 and
size<0>(typename TiledMma::ALayout{}) == NumThreadsPerWarpGroup and
size<0>(typename TiledMma::BLayout{}) == NumThreadsPerWarpGroup,
size<0>(typename TiledMma::BLayout{}) == NumThreadsPerWarpGroup,
"Stride of the first mode must be 0 and the size of the mode must be NumThreadsPerWarpGroup");
constexpr int MmaWarpGroups = size(TiledMma{}) / NumThreadsPerWarpGroup;
Layout warp_group_thread_layout = make_layout(Int<MmaWarpGroups>{},
Layout warp_group_thread_layout = make_layout(Int<MmaWarpGroups>{},
Int<NumThreadsPerWarpGroup>{});
int warp_group_idx = __shfl_sync(0xFFFFFFFF, thread_idx / NumThreadsPerWarpGroup, 0);
@@ -650,7 +649,7 @@ struct CollectiveMma<
k_tile_count -= prologue_mma_count;
smem_pipe_release.advance(k_tile_count);
// Wait on all GMMAs to complete
warpgroup_wait<0>();
include/cutlass/gemm/collective/sm90_sparse_mma_tma_gmma_ss_warpspecialized_fp8.hpp
+10 -11
View File
@@ -43,7 +43,6 @@
#include "cute/algorithm/functional.hpp"
#include "cute/atom/mma_atom.hpp"
#include "cute/algorithm/gemm.hpp"
#include "cute/tensor_predicate.hpp"
#include "cute/numeric/arithmetic_tuple.hpp"
/////////////////////////////////////////////////////////////////////////////////////////////////
@@ -195,7 +194,7 @@ struct CollectiveMma<
cute::conditional_t<cute::is_same_v<ElementA, float>,
cutlass::tfloat32_t,
uint_bit_t<sizeof_bits_v<ElementAMmaRaw>>>>;
using TmaInternalElementB = cute::conditional_t<cute::is_same_v<float, ElementB>,
using TmaInternalElementB = cute::conditional_t<cute::is_same_v<float, ElementB>,
tfloat32_t,
uint_bit_t<sizeof_bits_v<ElementBMma>>>;
@@ -216,7 +215,7 @@ struct CollectiveMma<
static constexpr int K_PIPE_MAX = DispatchPolicy::Stages;
static constexpr int K_PIPE_MMAS = 0;
static constexpr uint32_t TmaTransactionBytesMK =
static constexpr uint32_t TmaTransactionBytesMK =
cutlass::bits_to_bytes(cosize(take<0,2>(SmemLayoutA{})) * cute::sizeof_bits_v<ElementAMma>) +
cutlass::bits_to_bytes(cosize(take<0,2>(SmemLayoutE{})) * cute::sizeof_bits_v<ElementEMma>);
@@ -336,7 +335,7 @@ struct CollectiveMma<
constexpr int min_tma_aligned_elements_B = tma_alignment_bits / cutlass::sizeof_bits<ElementB>::value;
auto problem_shape_MNKL = append<4>(problem_shape, 1);
auto [M,N,K,L] = problem_shape_MNKL;
bool size_check = true;
// Check Alignment A
if constexpr (is_A_mn_major) {
@@ -416,7 +415,7 @@ struct CollectiveMma<
CUTLASS_DEVICE void
load(
Params const& mainloop_params,
MainloopPipeline pipeline,
MainloopPipeline pipeline,
PipelineState smem_pipe_write,
cute::tuple<TensorA, TensorB, TensorE> const& load_inputs,
BlockCoord const& blk_coord,
@@ -496,9 +495,9 @@ struct CollectiveMma<
// Issue the epilogue waits
if (lane_predicate) {
/* This helps avoid early exit of blocks in Cluster
* Waits for all stages to either be released (all
* 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
* then would just be acquired since the phase was
* still inverted from make_producer_start_state
*/
pipeline.producer_tail(smem_pipe_write);
@@ -534,14 +533,14 @@ struct CollectiveMma<
// Layout of warp group to thread mapping
static_assert(stride<0>(typename TiledMma::ALayout{}) == 0 and
static_assert(stride<0>(typename TiledMma::ALayout{}) == 0 and
stride<0>(typename TiledMma::BLayout{}) == 0 and
size<0>(typename TiledMma::ALayout{}) == NumThreadsPerWarpGroup and
size<0>(typename TiledMma::BLayout{}) == NumThreadsPerWarpGroup,
size<0>(typename TiledMma::BLayout{}) == NumThreadsPerWarpGroup,
"Stride of the first mode must be 0 and the size of the mode must be NumThreadsPerWarpGroup");
constexpr int MmaWarpGroups = size(TiledMma{}) / NumThreadsPerWarpGroup;
Layout warp_group_thread_layout = make_layout(Int<MmaWarpGroups>{},
Layout warp_group_thread_layout = make_layout(Int<MmaWarpGroups>{},
Int<NumThreadsPerWarpGroup>{});
int warp_group_idx = __shfl_sync(0xFFFFFFFF, thread_idx / NumThreadsPerWarpGroup, 0);
@@ -676,7 +675,7 @@ struct CollectiveMma<
k_tile_count -= prologue_mma_count;
smem_pipe_release.advance(k_tile_count);
// Wait on all GMMAs to complete
warpgroup_wait<0>();
include/cutlass/gemm/dispatch_policy.hpp
+74 -10
View File
@@ -511,61 +511,82 @@ struct KernelPtrArrayTmaWarpSpecializedInputTransformSm100 final {
// SM120 kernel schedules
template< int SchedulerPipelineStageCount_>
template<int SchedulerPipelineStageCount_>
struct KernelTmaWarpSpecializedCooperativeSm120 : KernelTmaWarpSpecializedCooperative {
static constexpr int SchedulerPipelineStageCount = SchedulerPipelineStageCount_;
};
template< int SchedulerPipelineStageCount_>
template<int SchedulerPipelineStageCount_>
struct KernelTmaWarpSpecializedPingpongSm120 : KernelTmaWarpSpecializedPingpong {
static constexpr int SchedulerPipelineStageCount = SchedulerPipelineStageCount_;
};
template< int SchedulerPipelineStageCount_>
template<int SchedulerPipelineStageCount_>
struct KernelTmaWarpSpecializedCooperativeBlockScaledSm120 : KernelTmaWarpSpecializedCooperative {
static constexpr int SchedulerPipelineStageCount = SchedulerPipelineStageCount_;
};
template< int SchedulerPipelineStageCount_>
template<int SchedulerPipelineStageCount_>
struct KernelTmaWarpSpecializedPingpongBlockScaledSm120 : KernelTmaWarpSpecializedPingpong {
static constexpr int SchedulerPipelineStageCount = SchedulerPipelineStageCount_;
};
// SM120 dense Ptr-array kernel schedules
template< int SchedulerPipelineStageCount_>
template<int SchedulerPipelineStageCount_>
struct KernelPtrArrayTmaWarpSpecializedCooperativeSm120 : KernelPtrArrayTmaWarpSpecializedCooperative {
static constexpr int SchedulerPipelineStageCount = SchedulerPipelineStageCount_;
};
template< int SchedulerPipelineStageCount_>
template<int SchedulerPipelineStageCount_>
struct KernelPtrArrayTmaWarpSpecializedPingpongSm120 : KernelPtrArrayTmaWarpSpecializedPingpong {
static constexpr int SchedulerPipelineStageCount = SchedulerPipelineStageCount_;
};
template< int SchedulerPipelineStageCount_>
template<int SchedulerPipelineStageCount_>
struct KernelPtrArrayTmaWarpSpecializedCooperativeBlockScaledSm120 : KernelPtrArrayTmaWarpSpecializedCooperative {
static constexpr int SchedulerPipelineStageCount = SchedulerPipelineStageCount_;
};
template< int SchedulerPipelineStageCount_>
template<int SchedulerPipelineStageCount_>
struct KernelPtrArrayTmaWarpSpecializedPingpongBlockScaledSm120 : KernelPtrArrayTmaWarpSpecializedPingpong {
static constexpr int SchedulerPipelineStageCount = SchedulerPipelineStageCount_;
};
// SM120 sparse kernel schedules
template< int SchedulerPipelineStageCount_, bool isAsymmetric_>
template<int SchedulerPipelineStageCount_, bool isAsymmetric_>
struct KernelTmaWarpSpecializedCooperativeSparseSm120 {
static constexpr int SchedulerPipelineStageCount = SchedulerPipelineStageCount_;
static constexpr bool isAsymmetric = isAsymmetric_;
};
template< int SchedulerPipelineStageCount_, bool isAsymmetric_>
template<int SchedulerPipelineStageCount_, bool isAsymmetric_>
struct KernelTmaWarpSpecializedCooperativeSparseBlockScaledSm120 {
static constexpr int SchedulerPipelineStageCount = SchedulerPipelineStageCount_;
static constexpr bool isAsymmetric = isAsymmetric_;
};
// SM120 blockwise kernel schedules
template <int SchedulerPipelineStageCount_>
struct KernelTmaWarpSpecializedCooperativeBlockwiseScalingSm120 : KernelTmaWarpSpecializedCooperative {
static constexpr int SchedulerPipelineStageCount = SchedulerPipelineStageCount_;
};
template <int SchedulerPipelineStageCount_>
struct KernelTmaWarpSpecializedPingpongBlockwiseScalingSm120 : KernelTmaWarpSpecializedPingpong {
static constexpr int SchedulerPipelineStageCount = SchedulerPipelineStageCount_;
};
template <int SchedulerPipelineStageCount_>
struct KernelPtrArrayTmaWarpSpecializedCooperativeBlockwiseScalingSm120 : KernelPtrArrayTmaWarpSpecializedCooperative {
static constexpr int SchedulerPipelineStageCount = SchedulerPipelineStageCount_;
};
template <int SchedulerPipelineStageCount_>
struct KernelPtrArrayTmaWarpSpecializedPingpongBlockwiseScalingSm120 : KernelPtrArrayTmaWarpSpecializedPingpong {
static constexpr int SchedulerPipelineStageCount = SchedulerPipelineStageCount_;
};
// Auxiliary Load Tag.
namespace kernel::detail {
@@ -776,6 +797,12 @@ struct KernelTmaWarpSpecializedMxf4Sm120 final : KernelScheduleMxNvf
struct KernelTmaWarpSpecializedPingpongMxf4Sm120 final : KernelScheduleMxNvf4Sm120, KernelTmaWarpSpecializedPingpong { };
struct KernelTmaWarpSpecializedMxf8f6f4Sm120 final : KernelScheduleMxf8f6f4Sm120, KernelTmaWarpSpecializedCooperative { };
struct KernelTmaWarpSpecializedPingpongMxf8f6f4Sm120 final : KernelScheduleMxf8f6f4Sm120, KernelTmaWarpSpecializedPingpong { };
// Blockwise Scaled GEMM
struct KernelScheduleSm120Blockwise: KernelScheduleSm120 { };
struct KernelTmaWarpSpecializedBlockwiseCooperativeSm120 final : KernelScheduleSm120Blockwise, KernelTmaWarpSpecializedCooperative { };
struct KernelTmaWarpSpecializedBlockwisePingpongSm120 final : KernelScheduleSm120Blockwise, KernelTmaWarpSpecializedPingpong { };
///////////////////////////////////////////////////////////////////////////////////////////////////////
// SM120 Sparse GEMM Dispatch Policies
///////////////////////////////////////////////////////////////////////////////////////////////////////
@@ -1120,6 +1147,43 @@ struct MainloopSm120TmaWarpSpecializedSparseBlockScaled {
using Schedule = KernelTmaWarpSpecializedCooperativeSparseBlockScaledSm120<SchedulerPipelineStageCount_, isAsymmetric>;
};
template <
int Stages_,
int SchedulerPipelineStageCount_,
class ClusterShape_,
class KernelSchedule_
>
struct MainloopSm120TmaWarpSpecializedBlockwiseScaling {
constexpr static int Stages = Stages_;
constexpr static int SchedulerPipelineStageCount = SchedulerPipelineStageCount_;
using ClusterShape = ClusterShape_;
using Schedule = KernelSchedule_;
constexpr static int PipelineAsyncMmaStages = 0;
using ArchTag = arch::Sm120;
};
template <
int Stages_,
int SchedulerPipelineStageCount_,
class ClusterShape_,
class KernelSchedule_
>
struct MainloopSm120ArrayTmaWarpSpecializedBlockwiseScaling {
constexpr static int Stages = Stages_;
constexpr static int SchedulerPipelineStageCount = SchedulerPipelineStageCount_;
using ClusterShape = ClusterShape_;
using Schedule = KernelSchedule_;
constexpr static int PipelineAsyncMmaStages = 0;
using ArchTag = arch::Sm120;
static_assert(cute::is_base_of_v<KernelPtrArrayTmaWarpSpecializedCooperative, Schedule> ||
cute::is_base_of_v<KernelPtrArrayTmaWarpSpecializedPingpong, Schedule>,
"KernelSchedule must be one of the Ptr-Array or Grouped Gemm TMA Warp Specialized Cooperative or Pingpong policies.");
};
//////////////////////////////////////////////////////////////////////////////
include/cutlass/gemm/group_array_problem_shape.hpp
+2 -2
View File
@@ -74,7 +74,7 @@ struct GroupProblemShape {
CUTLASS_HOST_DEVICE
bool
is_host_problem_shape_available() {
is_host_problem_shape_available() const {
return host_problem_shapes != nullptr;
}
};
@@ -113,7 +113,7 @@ public:
CUTLASS_HOST_DEVICE
bool
is_host_problem_shape_available() {
is_host_problem_shape_available() const {
return true;
}
private:
include/cutlass/gemm/kernel/sm100_gemm_array_tma_warpspecialized.hpp
+52 -89
View File
@@ -29,8 +29,6 @@
*
**************************************************************************************************/
#pragma once
#include "cutlass/cutlass.h"
@@ -50,6 +48,7 @@
#include "cutlass/gemm/kernel/sm100_tile_scheduler.hpp"
#include "cutlass/gemm/kernel/sm100_tile_scheduler_group.hpp"
#include "cutlass/pipeline/pipeline.hpp"
#include "cutlass/detail/sm100_tmem_helper.hpp"
#include "cute/tensor.hpp"
#include "cute/arch/tmem_allocator_sm100.hpp"
@@ -190,7 +189,6 @@ public:
// Kernel level shared memory storage
struct SharedStorage {
// Barriers should be allocated in lower 8KB of SMEM for SM100
struct PipelineStorage : cute::aligned_struct<16, _1> {
using MainloopPipelineStorage = typename CollectiveMainloop::PipelineStorage;
using EpiLoadPipelineStorage = typename CollectiveEpilogue::PipelineStorage;
@@ -206,7 +204,6 @@ public:
alignas(16) AccumulatorPipelineStorage accumulator;
alignas(16) CLCThrottlePipelineStorage clc_throttle;
alignas(16) arch::ClusterBarrier tmem_dealloc;
alignas(16) arch::ClusterBarrier epilogue_throttle;
} pipelines;
alignas(16) typename TileScheduler::CLCResponse clc_response[SchedulerPipelineStageCount];
@@ -280,7 +277,12 @@ public:
ProblemShape problem_shapes = args.problem_shape;
// Get SM count if needed, otherwise use user supplied SM count
int sm_count = args.hw_info.sm_count;
if (!IsGroupedGemmKernel && sm_count != 0) {
if (IsGroupedGemmKernel && sm_count <= 0) {
CUTLASS_TRACE_HOST(" WARNING: Arguments do not include a valid SM count.\n"
" For optimal performance, populate the arguments KernelHardwareInfo struct with the SM count.");
sm_count = KernelHardwareInfo::query_device_multiprocessor_count(args.hw_info.device_id);
}
else if (!IsGroupedGemmKernel && 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).");
}
@@ -487,7 +489,7 @@ public:
: WarpCategory::Epilogue;
uint32_t lane_predicate = cute::elect_one_sync();
auto cluster_shape = cutlass::detail::select_cluster_shape(ClusterShape{}, cute::cluster_shape());
auto cluster_shape = cutlass::detail::select_cluster_shape(ClusterShape{});
int cluster_size = size(cluster_shape);
uint32_t cta_rank_in_cluster = cute::block_rank_in_cluster();
bool is_first_cta_in_cluster = IsSchedDynamicPersistent ? (cta_rank_in_cluster == 0) : true;
@@ -542,7 +544,7 @@ public:
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 = 4;
epi_load_pipeline_params.initializing_warp = 1;
EpiLoadPipeline epi_load_pipeline(shared_storage.pipelines.epi_load, epi_load_pipeline_params);
// Epilogue Store pipeline
@@ -554,12 +556,11 @@ public:
typename LoadOrderBarrier::Params load_order_barrier_params;
load_order_barrier_params.group_id = (warp_category == WarpCategory::MainloopLoad) ? 0 : 1;
load_order_barrier_params.group_size = NumMainloopLoadThreads;
load_order_barrier_params.initializing_warp = 5;
load_order_barrier_params.initializing_warp = 3;
LoadOrderBarrier load_order_barrier(shared_storage.pipelines.load_order, load_order_barrier_params);
// CLC pipeline
typename CLCPipeline::Params clc_pipeline_params;
if (WarpCategory::Sched == warp_category) {
clc_pipeline_params.role = IsSchedDynamicPersistent ?
CLCPipeline::ThreadCategory::ProducerConsumer :
@@ -568,8 +569,7 @@ public:
else {
clc_pipeline_params.role = CLCPipeline::ThreadCategory::Consumer;
}
clc_pipeline_params.initializing_warp = 1;
clc_pipeline_params.initializing_warp = 4;
clc_pipeline_params.producer_arv_count = 1;
if constexpr (IsSchedDynamicPersistent) {
@@ -608,7 +608,7 @@ public:
// 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;
accumulator_pipeline_params.initializing_warp = 5;
AccumulatorPipeline accumulator_pipeline(shared_storage.pipelines.accumulator,
accumulator_pipeline_params,
cluster_shape,
@@ -641,28 +641,20 @@ public:
// 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;
if constexpr(!IsOverlappingAccum) {
if (WarpCategory::MMA == warp_category && has_mma_peer_cta && lane_predicate) {
tmem_deallocation_result_barrier.init(NumMMAThreads);
if (WarpCategory::MMA == warp_category) {
if constexpr(!IsOverlappingAccum) {
if (has_mma_peer_cta && lane_predicate) {
tmem_deallocation_result_barrier.init(NumMMAThreads);
}
}
}
else {
if (WarpCategory::MMA == warp_category && has_mma_peer_cta && lane_predicate) {
tmem_deallocation_result_barrier.init(NumEpilogueThreads*2);
else {
if (has_mma_peer_cta && lane_predicate) {
tmem_deallocation_result_barrier.init(NumEpilogueThreads*2);
}
else if (lane_predicate) {
tmem_deallocation_result_barrier.init(NumEpilogueThreads);
}
}
else if (WarpCategory::MMA == warp_category && lane_predicate) {
tmem_deallocation_result_barrier.init(NumEpilogueThreads);
}
}
// Initialize smem barrier for prologue throttling. Epilogue warps are stalled until the prologue finishes.
arch::ClusterBarrier& epilogue_throttle_barrier = shared_storage.pipelines.epilogue_throttle;
if (WarpCategory::MMA == warp_category && lane_predicate) {
epilogue_throttle_barrier.init( NumMMAThreads +
(is_first_cta_in_cluster ? NumSchedThreads : 0) +
NumMainloopLoadThreads +
(is_epi_load_needed ? NumEpilogueLoadThreads : 0));
}
// We need this to guarantee that the Pipeline init is visible
@@ -689,22 +681,17 @@ public:
// Calculate mask after cluster barrier arrival
mainloop_pipeline.init_masks(cluster_shape, block_id_in_cluster);
accumulator_pipeline.init_masks(cluster_shape);
accumulator_pipeline.init_masks(cluster_shape, block_id_in_cluster);
// TileID scheduler
TileScheduler scheduler(&shared_storage.clc_response[0], params.scheduler, block_id_in_cluster);
auto work_tile_info = scheduler.initial_work_tile_info(cluster_shape);
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"
//
// ((MMA_TILE_M,MMA_TILE_N),MMA_M,MMA_N,ACC_PIPE) where ACC_PIPE=2 so we can double buffer our accumulators for mainloop and epilogue.
TiledMma tiled_mma;
auto acc_shape = collective_mainloop.partition_accumulator_shape();
Tensor accumulators = cutlass::detail::make_sm100_accumulator<AccumulatorPipelineStageCount, IsOverlappingAccum>(
tiled_mma, acc_shape, EpilogueTile{});
auto tmem_storage = collective_mainloop.template init_tmem_tensors<EpilogueTile, IsOverlappingAccum>(EpilogueTile{});
pipeline_init_wait(cluster_size);
if constexpr (IsGroupedGemmKernel) {
@@ -719,16 +706,17 @@ public:
auto problem_shape_MNKL = append<4>(problem_shape.get_problem_shape(work_tile_info.L_idx), 1);
if (is_participant.main_load) {
auto load_inputs = collective_mainloop.load_init(
problem_shape_MNKL, params.mainloop,
shared_storage.tensors.mainloop,
shared_storage.tensormaps.mainloop,
params.hw_info.sm_count, sm_id, work_tile_info.L_idx);
// Ensure that the prefetched kernel does not touch
// unflushed global memory prior to this instruction
cutlass::arch::wait_on_dependent_grids();
bool do_load_order_arrive = is_epi_load_needed;
auto load_inputs = collective_mainloop.load_init(
problem_shape_MNKL, params.mainloop,
shared_storage.tensors.mainloop,
shared_storage.tensormaps.mainloop,
params.hw_info.sm_count, sm_id, work_tile_info.L_idx);
Tensor gA_mkl = get<0>(load_inputs);
// Fetch a copy of tensormaps for the CTA from Params
auto input_tensormaps = get<rank(load_inputs) - 1>(load_inputs);
@@ -737,9 +725,6 @@ public:
// Even the first tile for a CTA can be from any of the batches.
// And during initialization of the first TMA descriptor on host, we don't initialize to the first batch due to that args value being device-only.
bool did_batch_change = true;
// Signal the epilogue warps to proceed once the prologue is complete
epilogue_throttle_barrier.arrive();
bool requires_clc_query = true;
do {
@@ -824,9 +809,6 @@ public:
}
else if (is_participant.sched) {
// Signal the epilogue warps to proceed once the prologue is complete
epilogue_throttle_barrier.arrive();
// Grouped GEMM uses static tile scheduler
if constexpr (IsSchedDynamicPersistent) {
// Whether a new CLC query must be performed.
@@ -891,18 +873,8 @@ public:
__syncwarp();
tmem_allocation_result_barrier.arrive();
uint32_t tmem_base_ptr = shared_storage.tmem_base_ptr;
accumulators.data() = tmem_base_ptr;
int tmem_non_accumulator_base = tmem_base_ptr + cutlass::detail::find_tmem_tensor_col_offset(accumulators);
auto mma_inputs = collective_mainloop.mma_init(
params.mainloop,
collective_mainloop.slice_accumulator(accumulators, 0),
shared_storage.tensors.mainloop,
tmem_non_accumulator_base /*Start SF TMEM allocation after the accumulator*/);
// Signal the epilogue warps to proceed once the prologue is complete
epilogue_throttle_barrier.arrive();
collective_mainloop.set_tmem_offsets(tmem_storage, tmem_base_ptr);
auto mma_inputs = collective_mainloop.mma_init(tmem_storage, shared_storage.tensors.mainloop);
do {
@@ -917,30 +889,29 @@ public:
++clc_pipe_consumer_state;
}
// Wait for tmem accumulator buffer to become empty with a flipped phase
if constexpr (!IsOverlappingAccum) {
if (is_mma_leader_cta) {
accumulator_pipeline.producer_acquire(accumulator_pipe_producer_state);
}
}
if constexpr (IsGroupedGemmKernel) {
problem_shape_MNKL = append<4>(problem_shape.get_problem_shape(work_tile_info.L_idx), 1);
}
auto k_tile_count = TileScheduler::get_work_k_tile_count(work_tile_info, problem_shape_MNKL, CtaShape_MNK{});
int acc_stage = (IsOverlappingAccum) ? (accumulator_pipe_producer_state.phase() ^ 1) : (accumulator_pipe_producer_state.index());
auto accumulator = collective_mainloop.slice_accumulator(accumulators, acc_stage);
// Accumulator stage slice
int acc_stage = [&] () {
if constexpr (IsOverlappingAccum) {
return accumulator_pipe_producer_state.phase() ^ 1;
}
else {
return accumulator_pipe_producer_state.index();
}
}();
auto accumulator = collective_mainloop.slice_accumulator(tmem_storage, acc_stage);
if (is_mma_leader_cta) {
mainloop_pipe_consumer_state = collective_mainloop.mma(
cute::make_tuple(
mainloop_pipeline, accumulator_pipeline),
cute::make_tuple(
mainloop_pipe_consumer_state, accumulator_pipe_producer_state),
cute::make_tuple(mainloop_pipeline, accumulator_pipeline),
cute::make_tuple(mainloop_pipe_consumer_state, accumulator_pipe_producer_state),
accumulator,
mma_inputs,
cta_coord_mnkl,
k_tile_count
);
);
accumulator_pipeline.producer_commit(accumulator_pipe_producer_state);
}
++accumulator_pipe_producer_state;
@@ -969,7 +940,6 @@ public:
tmem_deallocation_result_barrier.wait(dealloc_barrier_phase);
tmem_deallocation_result_barrier.arrive(mma_peer_cta_rank, is_mma_leader_cta);
}
}
else {
tmem_deallocation_result_barrier.wait(dealloc_barrier_phase);
@@ -997,9 +967,6 @@ public:
bool did_batch_change = true;
constexpr bool IsEpiLoad = true;
// Signal the epilogue warps to proceed once the prologue is complete
epilogue_throttle_barrier.arrive();
do {
int32_t curr_batch = work_tile_info.L_idx;
if (did_batch_change) {
@@ -1069,14 +1036,10 @@ public:
}
else if (is_participant.epilogue) {
// Throttle the epilogue warps to improve prologue performance
static constexpr int epilogue_throttle_phase_bit = 0;
epilogue_throttle_barrier.wait(epilogue_throttle_phase_bit);
// Wait for tmem allocate here
tmem_allocation_result_barrier.arrive_and_wait();
uint32_t tmem_base_ptr = shared_storage.tmem_base_ptr;
accumulators.data() = tmem_base_ptr;
collective_mainloop.set_tmem_offsets(tmem_storage, tmem_base_ptr);
auto warp_idx_in_epi = canonical_warp_idx_sync() - static_cast<int>(WarpCategory::Epilogue);
bool do_tail_store = false;
@@ -1110,7 +1073,7 @@ public:
++clc_pipe_consumer_state;
}
// Accumulator stage slice after making sure allocation has been performed
// Accumulator stage slice
int acc_stage = [&] () {
if constexpr (IsOverlappingAccum) {
return accumulator_pipe_consumer_state.phase();
@@ -1119,6 +1082,7 @@ public:
return accumulator_pipe_consumer_state.index();
}
}();
auto accumulator = collective_mainloop.slice_accumulator(tmem_storage, acc_stage);
// Fusions may need problem shape for the current group
if constexpr (IsGroupedGemmKernel) {
@@ -1139,7 +1103,7 @@ public:
cta_coord_mnkl,
TileShape{},
TiledMma{},
collective_mainloop.slice_accumulator(accumulators, acc_stage),
accumulator,
shared_storage.tensors.epilogue,
cute::make_tuple(epi_store_tensormap, did_batch_change)
);
@@ -1175,7 +1139,6 @@ public:
}
else {
}
}
};
include/cutlass/gemm/kernel/sm100_gemm_array_tma_warpspecialized_mma_transform.hpp
+6 -1
View File
@@ -289,7 +289,12 @@ public:
ProblemShape problem_shapes = args.problem_shape;
// Get SM count if needed, otherwise use user supplied SM count
int sm_count = args.hw_info.sm_count;
if (!IsGroupedGemmKernel && sm_count != 0) {
if (IsGroupedGemmKernel && sm_count <= 0) {
CUTLASS_TRACE_HOST(" WARNING: Arguments do not include a valid SM count.\n"
" For optimal performance, populate the arguments KernelHardwareInfo struct with the SM count.");
sm_count = KernelHardwareInfo::query_device_multiprocessor_count(args.hw_info.device_id);
}
else if (!IsGroupedGemmKernel && 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).");
}
include/cutlass/gemm/kernel/sm100_gemm_tma_warpspecialized.hpp
+36
View File
@@ -172,6 +172,9 @@ public:
using CLCPipeline = cutlass::PipelineCLCFetchAsync<SchedulerPipelineStageCount, ClusterShape>;
using CLCPipelineState = typename CLCPipeline::PipelineState;
using CLCThrottlePipeline = cutlass::PipelineAsync<SchedulerPipelineStageCount>;
using CLCThrottlePipelineState = typename CLCThrottlePipeline::PipelineState;
using TmemAllocator = cute::conditional_t<cute::size(cute::shape<0>(typename TiledMma::ThrLayoutVMNK{})) == 1,
cute::TMEM::Allocator1Sm, cute::TMEM::Allocator2Sm>;
@@ -183,12 +186,14 @@ public:
using LoadOrderBarrierStorage = typename LoadOrderBarrier::SharedStorage;
using CLCPipelineStorage = typename CLCPipeline::SharedStorage;
using AccumulatorPipelineStorage = typename AccumulatorPipeline::SharedStorage;
using CLCThrottlePipelineStorage = typename CLCThrottlePipeline::SharedStorage;
alignas(16) MainloopPipelineStorage mainloop;
alignas(16) EpiLoadPipelineStorage epi_load;
alignas(16) LoadOrderBarrierStorage load_order;
alignas(16) CLCPipelineStorage clc;
alignas(16) AccumulatorPipelineStorage accumulator;
alignas(16) CLCThrottlePipelineStorage clc_throttle;
alignas(16) arch::ClusterBarrier tmem_dealloc;
} pipelines;
@@ -530,6 +535,22 @@ public:
cute::true_type{}, // Perform barrier init
cute::false_type{}); // Delay mask calculation
// CLC throttle pipeline
typename CLCThrottlePipeline::Params clc_throttle_pipeline_params;
if (WarpCategory::MainloopLoad == warp_category) {
clc_throttle_pipeline_params.role = CLCThrottlePipeline::ThreadCategory::Producer;
}
if (WarpCategory::Sched == warp_category) {
clc_throttle_pipeline_params.role = CLCThrottlePipeline::ThreadCategory::Consumer;
}
clc_throttle_pipeline_params.producer_arv_count = NumMainloopLoadThreads;
clc_throttle_pipeline_params.consumer_arv_count = NumSchedThreads;
clc_throttle_pipeline_params.dst_blockid = 0;
clc_throttle_pipeline_params.initializing_warp = 3;
CLCThrottlePipeline clc_throttle_pipeline(shared_storage.pipelines.clc_throttle, clc_throttle_pipeline_params);
CLCThrottlePipelineState clc_pipe_throttle_consumer_state;
CLCThrottlePipelineState clc_pipe_throttle_producer_state = cutlass::make_producer_start_state<CLCThrottlePipeline>();
// Tmem allocator
TmemAllocator tmem_allocator{};
@@ -599,6 +620,7 @@ public:
cutlass::arch::wait_on_dependent_grids();
bool do_load_order_arrive = is_epi_load_needed;
bool requires_clc_query = true;
do {
// Get the number of K tiles to compute for this work as well as the starting K tile offset of the work.
@@ -606,6 +628,14 @@ public:
auto k_tile_count = TileScheduler::get_work_k_tile_count(work_tile_info, problem_shape_MNKL, CtaShape_MNK{});
auto k_tile_prologue = min(MainloopPipeline::Stages, k_tile_count);
if constexpr (IsSchedDynamicPersistent) {
if (is_first_cta_in_cluster && requires_clc_query) {
clc_throttle_pipeline.producer_acquire(clc_pipe_throttle_producer_state);
clc_throttle_pipeline.producer_commit(clc_pipe_throttle_producer_state);
++clc_pipe_throttle_producer_state;
}
}
// Start mainloop prologue loads, arrive on the epilogue residual load barrier, resume mainloop loads
auto [mainloop_producer_state_next, k_tile_iter_next] = collective_mainloop.load(
mainloop_pipeline,
@@ -639,6 +669,7 @@ public:
);
work_tile_info = next_work_tile_info;
cta_coord_mnkl = scheduler.work_tile_to_cta_coord(work_tile_info);
requires_clc_query = increment_pipe;
if (increment_pipe) {
++clc_pipe_consumer_state;
}
@@ -658,6 +689,11 @@ public:
do {
if (requires_clc_query) {
// Throttle CLC query to mitigate workload imbalance caused by skews among persistent workers.
clc_throttle_pipeline.consumer_wait(clc_pipe_throttle_consumer_state);
clc_throttle_pipeline.consumer_release(clc_pipe_throttle_consumer_state);
++clc_pipe_throttle_consumer_state;
// Query next clcID and update producer state
clc_pipe_producer_state = scheduler.advance_to_next_work(clc_pipeline, clc_pipe_producer_state);
}
include/cutlass/gemm/kernel/sm100_sparse_gemm_tma_warpspecialized.hpp
+36
View File
@@ -178,6 +178,9 @@ public:
using CLCPipeline = cutlass::PipelineCLCFetchAsync<SchedulerPipelineStageCount, ClusterShape>;
using CLCPipelineState = typename CLCPipeline::PipelineState;
using CLCThrottlePipeline = cutlass::PipelineAsync<SchedulerPipelineStageCount>;
using CLCThrottlePipelineState = typename CLCThrottlePipeline::PipelineState;
using TmemAllocator = cute::conditional_t<cute::size(cute::shape<0>(typename TiledMma::ThrLayoutVMNK{})) == 1,
cute::TMEM::Allocator1Sm, cute::TMEM::Allocator2Sm>;
@@ -190,12 +193,14 @@ public:
using LoadOrderBarrierStorage = typename LoadOrderBarrier::SharedStorage;
using CLCPipelineStorage = typename CLCPipeline::SharedStorage;
using AccumulatorPipelineStorage = typename AccumulatorPipeline::SharedStorage;
using CLCThrottlePipelineStorage = typename CLCThrottlePipeline::SharedStorage;
alignas(16) MainloopPipelineStorage mainloop;
alignas(16) EpiLoadPipelineStorage epi_load;
alignas(16) LoadOrderBarrierStorage load_order;
alignas(16) CLCPipelineStorage clc;
alignas(16) AccumulatorPipelineStorage accumulator;
alignas(16) CLCThrottlePipelineStorage clc_throttle;
alignas(16) arch::ClusterBarrier tmem_dealloc;
} pipelines;
@@ -580,6 +585,22 @@ public:
cute::true_type{}, // Perform barrier init
cute::false_type{}); // Delay mask calculation
// CLC throttle pipeline
typename CLCThrottlePipeline::Params clc_throttle_pipeline_params;
if (WarpCategory::MainloopLoad == warp_category) {
clc_throttle_pipeline_params.role = CLCThrottlePipeline::ThreadCategory::Producer;
}
if (WarpCategory::Sched == warp_category) {
clc_throttle_pipeline_params.role = CLCThrottlePipeline::ThreadCategory::Consumer;
}
clc_throttle_pipeline_params.producer_arv_count = NumMainloopLoadThreads;
clc_throttle_pipeline_params.consumer_arv_count = NumSchedThreads;
clc_throttle_pipeline_params.dst_blockid = 0;
clc_throttle_pipeline_params.initializing_warp = 3;
CLCThrottlePipeline clc_throttle_pipeline(shared_storage.pipelines.clc_throttle, clc_throttle_pipeline_params);
CLCThrottlePipelineState clc_pipe_throttle_consumer_state;
CLCThrottlePipelineState clc_pipe_throttle_producer_state = cutlass::make_producer_start_state<CLCThrottlePipeline>();
// Tmem allocator
TmemAllocator tmem_allocator{};
@@ -649,12 +670,21 @@ public:
cutlass::arch::wait_on_dependent_grids();
bool do_load_order_arrive = is_epi_load_needed;
bool requires_clc_query = true;
do {
// 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{}, load_inputs.k_tiles);
auto k_tile_count = TileScheduler::get_work_k_tile_count(work_tile_info, problem_shape_MNKL, CtaShape_MNK{});
if constexpr (IsSchedDynamicPersistent) {
if (is_first_cta_in_cluster && requires_clc_query) {
clc_throttle_pipeline.producer_acquire(clc_pipe_throttle_producer_state);
clc_throttle_pipeline.producer_commit(clc_pipe_throttle_producer_state);
++clc_pipe_throttle_producer_state;
}
}
// Start mainloop prologue loads, arrive on the epilogue residual load barrier, resume mainloop loads
auto [mainloop_producer_state_next, unused_] = collective_mainloop.load(
mainloop_pipeline,
@@ -678,6 +708,7 @@ public:
);
work_tile_info = next_work_tile_info;
cta_coord_mnkl = scheduler.work_tile_to_cta_coord(work_tile_info);
requires_clc_query = increment_pipe;
if (increment_pipe) {
++clc_pipe_consumer_state;
}
@@ -697,6 +728,11 @@ public:
do {
if (requires_clc_query) {
// Throttle CLC query to mitigate workload imbalance caused by skews among persistent workers.
clc_throttle_pipeline.consumer_wait(clc_pipe_throttle_consumer_state);
clc_throttle_pipeline.consumer_release(clc_pipe_throttle_consumer_state);
++clc_pipe_throttle_consumer_state;
// Query next clcID and update producer state
clc_pipe_producer_state = scheduler.advance_to_next_work(clc_pipeline, clc_pipe_producer_state);
}
include/cutlass/gemm/kernel/sm100_tile_scheduler_group.hpp
+5 -10
View File
@@ -59,8 +59,7 @@ class PersistentTileSchedulerSm100Group {
public:
using UnderlyingScheduler = PersistentTileSchedulerSm90Group<GroupProblemShape, SchedulerPipelineStageCount>;
using UnderlyingProblemShape = typename GroupProblemShape::UnderlyingProblemShape;
using Params = PersistentTileSchedulerSm100GroupParams<UnderlyingProblemShape>;
using Params = PersistentTileSchedulerSm100GroupParams<GroupProblemShape>;
using WorkTileInfo = typename UnderlyingScheduler::WorkTileInfo;
using Arguments = typename UnderlyingScheduler::Arguments;
using RasterOrder = typename Params::RasterOrder;
@@ -94,7 +93,6 @@ public:
shape_div(tile_shape_mnk, selected_cluster_shape)); // Static Cluster: Blackwell builders expects TileShape to be Cluster's Tile Shape, Hopper doesn't
dim3 problem_blocks = get_tiled_cta_shape_mnl(
problem_shapes.groups(),
problem_shapes,
hw_info,
cta_shape, selected_cluster_shape);
@@ -102,9 +100,7 @@ public:
Params params;
params.initialize(
problem_blocks,
problem_shapes.groups(),
problem_shapes.problem_shapes,
problem_shapes.host_problem_shapes,
problem_shapes,
to_gemm_coord(cta_shape),
to_gemm_coord(selected_cluster_shape),
hw_info,
@@ -144,8 +140,8 @@ public:
template<class BlockShape, class ClusterShape>
CUTLASS_HOST_DEVICE static
dim3
get_tiled_cta_shape_mnl(int groups, GroupProblemShape problem_shapes, KernelHardwareInfo hw_info, BlockShape cta_shape, ClusterShape cluster_shape) {
return UnderlyingScheduler::get_tiled_cta_shape_mnl(groups, problem_shapes, hw_info, cta_shape, cluster_shape);
get_tiled_cta_shape_mnl(GroupProblemShape const &problem_shapes, KernelHardwareInfo hw_info, BlockShape cta_shape, ClusterShape cluster_shape) {
return UnderlyingScheduler::get_tiled_cta_shape_mnl(problem_shapes, hw_info, cta_shape, cluster_shape);
}
// Given the inputs, computes the physical grid we should launch.
@@ -154,13 +150,12 @@ public:
static dim3
get_grid_shape(
Params const& params,
GroupProblemShape problem_shapes,
GroupProblemShape const& problem_shapes,
BlockShape cta_shape,
[[maybe_unused]] AtomThrShape atom_thr_shape,
ClusterShape cluster_shape,
KernelHardwareInfo hw_info) {
dim3 problem_blocks = get_tiled_cta_shape_mnl(
problem_shapes.groups(),
problem_shapes,
hw_info,
cta_shape,
include/cutlass/gemm/kernel/sm90_gemm_tma_warpspecialized_cooperative.hpp
+2 -1
View File
@@ -442,7 +442,8 @@ public:
// Mainloop Load pipeline
using MainloopPipeline = typename CollectiveMainloop::MainloopPipeline;
typename MainloopPipeline::Params mainloop_pipeline_params;
if (warp_group_role == WarpGroupRole::Producer && producer_warp_role == ProducerWarpRole::Mainloop) {
if (warp_group_role == WarpGroupRole::Producer && (producer_warp_role == ProducerWarpRole::Mainloop ||
producer_warp_role == ProducerWarpRole::MainloopAux)) {
mainloop_pipeline_params.role = MainloopPipeline::ThreadCategory::Producer;
}
if (warp_group_role == WarpGroupRole::Consumer0 || warp_group_role == WarpGroupRole::Consumer1) {
include/cutlass/gemm/kernel/sm90_gemm_tma_warpspecialized_pingpong.hpp
+2 -1
View File
@@ -447,7 +447,8 @@ public:
// Mainloop Load pipeline
using MainloopPipeline = typename CollectiveMainloop::MainloopPipeline;
typename MainloopPipeline::Params mainloop_pipeline_params;
if (warp_group_role == WarpGroupRole::Producer && producer_warp_role == ProducerWarpRole::Mainloop) {
if (warp_group_role == WarpGroupRole::Producer && (producer_warp_role == ProducerWarpRole::Mainloop
|| producer_warp_role == ProducerWarpRole::MainloopAux)) {
mainloop_pipeline_params.role = MainloopPipeline::ThreadCategory::Producer;
}
if (warp_group_role == WarpGroupRole::Consumer0 || warp_group_role == WarpGroupRole::Consumer1) {
include/cutlass/gemm/kernel/sm90_tile_scheduler_group.hpp
+18 -21
View File
@@ -94,7 +94,7 @@ public:
};
using ProblemShape = typename GroupProblemShape::UnderlyingProblemShape;
using Params = PersistentTileSchedulerSm90GroupParams<ProblemShape>;
using Params = PersistentTileSchedulerSm90GroupParams<GroupProblemShape>;
using RasterOrder = typename Params::RasterOrder;
using RasterOrderOptions = typename Params::RasterOrderOptions;
static constexpr bool IsDynamicPersistent = false;
@@ -160,7 +160,6 @@ public:
static_assert(cute::is_static<ClusterShape>::value);
dim3 problem_blocks = get_tiled_cta_shape_mnl(
problem_shapes.groups(),
problem_shapes,
hw_info,
tile_shape, cluster_shape);
@@ -168,9 +167,7 @@ public:
Params params;
params.initialize(
problem_blocks,
problem_shapes.groups(),
problem_shapes.problem_shapes,
problem_shapes.host_problem_shapes,
problem_shapes,
to_gemm_coord(tile_shape),
to_gemm_coord(cluster_shape),
hw_info,
@@ -187,7 +184,7 @@ public:
dim3
get_grid_shape(
[[maybe_unused]] Params const& params,
GroupProblemShape problem_shapes,
GroupProblemShape const& problem_shapes,
TileShape tile_shape,
ClusterShape cluster_shape,
KernelHardwareInfo hw_info,
@@ -195,7 +192,6 @@ public:
bool truncate_by_problem_size=true) {
dim3 problem_blocks = get_tiled_cta_shape_mnl(
problem_shapes.groups(),
problem_shapes,
hw_info,
tile_shape, cluster_shape);
@@ -215,7 +211,8 @@ public:
template<class BlockShape, class ClusterShape>
CUTLASS_HOST_DEVICE static
dim3
get_tiled_cta_shape_mnl(int groups, GroupProblemShape problem_shapes, KernelHardwareInfo hw_info, BlockShape cta_shape, ClusterShape cluster_shape) {
get_tiled_cta_shape_mnl(GroupProblemShape const& problem_shapes, KernelHardwareInfo hw_info, BlockShape cta_shape, ClusterShape cluster_shape) {
int groups = problem_shapes.groups();
uint32_t total_ctas = 0;
uint32_t cta_in_N_dim = 1; // We linearize the blocks across all the problems here
@@ -259,20 +256,21 @@ public:
}
int lane_idx = canonical_lane_idx();
if (lane_idx < params_.groups_) {
cached_problem_shapes_[1] = params_.problem_shapes_[lane_idx];
if (lane_idx < params_.problem_shapes_.groups()) {
cached_problem_shapes_[1] = params_.problem_shapes_.get_problem_shape(lane_idx);
}
total_grid_size_ = uint64_t(gridDim.x) * uint64_t(gridDim.y) * uint64_t(gridDim.z);
uint64_t ctas_along_m, ctas_along_n;
if (is_tuple<decltype(cute::shape<0>(params_.problem_shapes_[0]))>::value ||
is_tuple<decltype(cute::shape<1>(params_.problem_shapes_[0]))>::value) {
ctas_along_m = cute::size(cute::ceil_div(cute::shape<0>(params_.problem_shapes_[0]), scheduler_params.cta_shape_.m()));
ctas_along_n = cute::size(cute::ceil_div(cute::shape<1>(params_.problem_shapes_[0]), scheduler_params.cta_shape_.n()));
ProblemShape problem_shape = params_.problem_shapes_.get_problem_shape(0);
if (is_tuple<decltype(cute::shape<0>(problem_shape))>::value ||
is_tuple<decltype(cute::shape<1>(problem_shape))>::value) {
ctas_along_m = cute::size(cute::ceil_div(cute::shape<0>(problem_shape), scheduler_params.cta_shape_.m()));
ctas_along_n = cute::size(cute::ceil_div(cute::shape<1>(problem_shape), scheduler_params.cta_shape_.n()));
}
else {
ctas_along_m = scheduler_params.divmod_cta_shape_m_.divide(cute::shape<0>(params_.problem_shapes_[0]) + scheduler_params.divmod_cta_shape_m_.divisor - 1);
ctas_along_n = scheduler_params.divmod_cta_shape_n_.divide(cute::shape<1>(params_.problem_shapes_[0]) + scheduler_params.divmod_cta_shape_n_.divisor - 1);
ctas_along_m = scheduler_params.divmod_cta_shape_m_.divide(cute::shape<0>(problem_shape) + scheduler_params.divmod_cta_shape_m_.divisor - 1);
ctas_along_n = scheduler_params.divmod_cta_shape_n_.divide(cute::shape<1>(problem_shape) + scheduler_params.divmod_cta_shape_n_.divisor - 1);
}
auto problem_blocks_m = round_up(ctas_along_m, (1 << params_.log_swizzle_size_) * params_.cluster_shape_.m());
auto problem_blocks_n = round_up(ctas_along_n, (1 << params_.log_swizzle_size_) * params_.cluster_shape_.n());
@@ -292,8 +290,7 @@ public:
get_work_idx_m_and_n(
uint64_t linear_idx,
GroupInfo& group_info,
int32_t total_problem_groups,
ProblemShape* problem_shapes,
GroupProblemShape &problem_shapes,
ProblemShape (&cached_problem_shapes)[2],
GemmCoord cta_shape,
GemmCoord cluster_shape,
@@ -308,13 +305,14 @@ public:
// Use a warp to "speculatively" check if the work tile maps to the next 32 groups
int lane_idx = canonical_lane_idx();
int total_problem_groups = problem_shapes.groups();
if (linear_idx >= group_info.total_tiles + group_info.start_linear_idx) {
group_info.group_idx += lane_idx;
for ( ; ; group_info.group_idx += NumThreadsPerWarp) {
cached_problem_shapes[0] = cached_problem_shapes[1];
if (group_info.group_idx + NumThreadsPerWarp < total_problem_groups) {
cached_problem_shapes[1] = problem_shapes[group_info.group_idx + NumThreadsPerWarp];
cached_problem_shapes[1] = problem_shapes.get_problem_shape(group_info.group_idx + NumThreadsPerWarp);
}
if (group_info.group_idx < total_problem_groups) {
uint64_t ctas_along_m, ctas_along_n;
@@ -354,7 +352,7 @@ public:
group_info.total_tiles = __shfl_sync(0xffffffff, group_info.total_tiles, first_succeeding_thread);
group_info.problem_blocks_along_raster_order = __shfl_sync(0xffffffff, group_info.problem_blocks_along_raster_order, first_succeeding_thread);
if (group_info.group_idx + lane_idx < total_problem_groups) {
cached_problem_shapes[1] = problem_shapes[group_info.group_idx + lane_idx];
cached_problem_shapes[1] = problem_shapes.get_problem_shape(group_info.group_idx + lane_idx);
}
break;
}
@@ -419,7 +417,6 @@ public:
return get_work_idx_m_and_n<WorkTileInfo>(
linear_idx,
current_group_info_,
scheduler_params.groups_,
scheduler_params.problem_shapes_,
cached_problem_shapes_,
scheduler_params.cta_shape_,
include/cutlass/gemm/kernel/sm90_tile_scheduler_stream_k.hpp
-1
View File
@@ -1083,7 +1083,6 @@ private:
// output tile. This work will thus be subsumed by the previous stream-K unit.
--unit_idx;
}
return unit_idx;
};
include/cutlass/gemm/kernel/tile_scheduler_params.h
+19 -29
View File
@@ -59,6 +59,14 @@ get_max_cta_occupancy(int max_sm_per_gpc, GemmCoord cluster_shape, int sm_count)
int const min_num_gpc = sm_count < max_sm_per_gpc ? 1 : sm_count / max_sm_per_gpc;
int const max_cta_occupancy_per_gpc = max_sm_per_gpc - (max_sm_per_gpc % cluster_size);
int cta_per_device = min_num_gpc * max_cta_occupancy_per_gpc;
// Suppose max_sm_per_gpc = 20, cluster_size = 8, sm_count = 148
// min_num_gpc = 148 / 20 = 7
// max_cta_occupancy_per_gpc = 20 - (20 % 8) = 16
// cta_per_device = 7 * 16 = 112
// num_gpc_residual = 148 % 20 = 8
// max_cta_occupancy_per_residual_gpc = 8 - (8 % 8) = 8
// cta_per_device += 8 = 120
// cta_per_device = 120 < 148 ? 148 : 120 = 148
// The calculation below allows for larger grid size launch for different GPUs.
int const num_gpc_residual = sm_count < max_sm_per_gpc ? 0 : sm_count % max_sm_per_gpc;
@@ -658,7 +666,6 @@ struct PersistentTileSchedulerSm90StreamKParams {
// number of K tiles per stream-K unit remains above min_iters_per_sk_unit_
uint32_t groups = platform::min(max_groups_problem, uint32_t(max_sk_groups_));
// Grouping is disabled when separate reduction is used because grouping is primarily an attempt
// to improve L2 locality, and L2-locality optimizations are unnecessary when the the kernel
// is a single wave (which is the case for separate reduction).
@@ -1616,19 +1623,10 @@ struct PersistentTileSchedulerSm90StreamKParams {
////////////////////////////////////////////////////////////////////////////////
// Parameters for SM90 persistent group scheduler (only used for Grouped Gemms)
template<class ProblemShape>
template<class GroupProblemShape>
struct PersistentTileSchedulerSm90GroupParams {
enum class RasterOrder {
AlongM,
AlongN
};
enum class RasterOrderOptions {
Heuristic,
AlongM,
AlongN
};
using RasterOrder = cutlass::gemm::kernel::detail::RasterOrder;
using RasterOrderOptions = cutlass::gemm::kernel::detail::RasterOrderOptions;
FastDivmodU64Pow2 divmod_cluster_shape_major_{};
FastDivmodU64Pow2 divmod_cluster_shape_minor_{};
@@ -1640,8 +1638,7 @@ struct PersistentTileSchedulerSm90GroupParams {
int32_t log_swizzle_size_ = 0;
RasterOrder raster_order_ = RasterOrder::AlongN;
int32_t groups_ = 0;
ProblemShape* problem_shapes_ = nullptr;
GroupProblemShape problem_shapes_;
GemmCoord cta_shape_;
GemmCoord cluster_shape_;
@@ -1651,9 +1648,7 @@ struct PersistentTileSchedulerSm90GroupParams {
void
initialize(
dim3 problem_blocks,
int32_t groups,
ProblemShape* problem_shapes,
ProblemShape const* host_problem_shapes,
GroupProblemShape problem_shapes,
GemmCoord cta_shape,
GemmCoord cluster_shape,
KernelHardwareInfo const& hw_info,
@@ -1677,13 +1672,12 @@ struct PersistentTileSchedulerSm90GroupParams {
//
// Set members
//
groups_ = groups;
problem_shapes_ = problem_shapes;
cta_shape_ = cta_shape;
cluster_shape_ = cluster_shape;
blocks_across_problem_ = problem_blocks.x * problem_blocks.y * problem_blocks.z;
pre_processed_problem_shapes = (host_problem_shapes == nullptr) ? false : true;
pre_processed_problem_shapes = problem_shapes.is_host_problem_shape_available();
log_swizzle_size_ = log_swizzle_size;
raster_order_ = raster_order;
@@ -2442,12 +2436,12 @@ struct PersistentTileSchedulerSm100StreamKParams {
////////////////////////////////////////////////////////////////////////////////
// Parameters for SM100 persistent group scheduler (only used for Grouped Gemms)
template<class ProblemShape>
template<class GroupProblemShape>
struct PersistentTileSchedulerSm100GroupParams {
using UnderlyingSm90Params = PersistentTileSchedulerSm90GroupParams<ProblemShape>;
using RasterOrder = typename UnderlyingSm90Params::RasterOrder;
using RasterOrderOptions = typename UnderlyingSm90Params::RasterOrderOptions;
using UnderlyingSm90Params = PersistentTileSchedulerSm90GroupParams<GroupProblemShape>;
using RasterOrder = cutlass::gemm::kernel::detail::RasterOrder;
using RasterOrderOptions = cutlass::gemm::kernel::detail::RasterOrderOptions;
UnderlyingSm90Params params_sm90_{};
@@ -2457,9 +2451,7 @@ struct PersistentTileSchedulerSm100GroupParams {
void
initialize(
dim3 problem_blocks,
int32_t groups,
ProblemShape* problem_shapes,
ProblemShape const* host_problem_shapes,
GroupProblemShape problem_shapes,
GemmCoord cta_shape,
GemmCoord cluster_shape,
KernelHardwareInfo const& hw_info,
@@ -2469,9 +2461,7 @@ struct PersistentTileSchedulerSm100GroupParams {
params_sm90_.initialize(
problem_blocks,
groups,
problem_shapes,
host_problem_shapes,
cta_shape,
cluster_shape,
hw_info,
include/cutlass/numeric_conversion.h
+117 -1
View File
@@ -3968,7 +3968,6 @@ struct NumericArrayConverter<float_e2m1_t, float, N, Round> {
}
};
/////////////////////////////////////////////////////////////////////////////////////////////////
/// Partial specialization for Array<int8_t> <= Array<float>
@@ -4377,6 +4376,123 @@ namespace detail {
}
/////////////////////////////////////////////////////////////////////////////////////////////////
/// Partial specialization for Array<half_t, N> <= Array<float_e2m1_t, N>
template <
FloatRoundStyle Round,
int N
>
struct NumericArrayConverter<cutlass::half_t, cutlass::float_e2m1_t, N, Round> {
using result_element = cutlass::half_t;
using source_element = cutlass::float_e2m1_t;
using result_type = Array<result_element, N>;
using source_type = Array<source_element, N>;
static FloatRoundStyle const round_style = Round;
private:
using result_type_packed_8 = Array<cutlass::half_t, 8>;
using result_type_packed_4 = Array<cutlass::half_t, 4>;
using result_type_packed_2 = Array<cutlass::half_t, 2>;
using source_type_packed_8 = Array<cutlass::float_e2m1_t, 8>;
using source_type_packed_4 = Array<cutlass::float_e2m1_t, 4>;
using source_type_packed_2 = Array<cutlass::float_e2m1_t, 2>;
using ScalarConverter = NumericConverter<cutlass::half_t, cutlass::float_e2m1_t, Round>;
#if defined(CUDA_PTX_FP8_CVT_ENABLED)
CUTLASS_DEVICE
static result_type_packed_8 ptx_convert(source_type_packed_8 const &source) {
result_type_packed_8 out;
uint32_t* out_fp16 = reinterpret_cast<uint32_t*>(&out);
uint32_t const& src_packed = reinterpret_cast<uint32_t const&>(source);
asm volatile( \
"{\n" \
".reg .b8 byte0, byte1, byte2, byte3;\n" \
"mov.b32 {byte0, byte1, byte2, byte3}, %4;\n" \
"cvt.rn.f16x2.e2m1x2 %0, byte0;\n" \
"cvt.rn.f16x2.e2m1x2 %1, byte1;\n" \
"cvt.rn.f16x2.e2m1x2 %2, byte2;\n" \
"cvt.rn.f16x2.e2m1x2 %3, byte3;\n" \
"}\n" : "=r"(out_fp16[0]), "=r"(out_fp16[1]) , "=r"(out_fp16[2]), "=r"(out_fp16[3]): "r"(src_packed));
return out;
}
CUTLASS_DEVICE
static result_type_packed_4 ptx_convert(source_type_packed_4 const &source) {
result_type_packed_4 out;
uint32_t* out_fp16 = reinterpret_cast<uint32_t*>(&out);
uint16_t const& src_packed = reinterpret_cast<uint16_t const&>(source);
asm volatile( \
"{\n" \
".reg .b8 byte0, byte1;\n" \
"mov.b16 {byte0, byte1}, %2;\n" \
"cvt.rn.f16x2.e2m1x2 %0, byte0;\n" \
"cvt.rn.f16x2.e2m1x2 %1, byte1;\n" \
"}\n" : "=r"(out_fp16[0]), "=r"(out_fp16[1]) : "h"(src_packed));
return out;
}
CUTLASS_DEVICE
static result_type_packed_2 ptx_convert(source_type_packed_2 const &source) {
result_type_packed_2 out;
uint32_t* out_fp16 = reinterpret_cast<uint32_t*>(&out);
uint16_t const& src_packed = static_cast<uint16_t const&>(reinterpret_cast<uint8_t const&>(source));
asm volatile( \
"{\n" \
".reg .b8 byte0, byte1;\n" \
"mov.b16 {byte0, byte1}, %1;\n" \
"cvt.rn.f16x2.e2m1x2 %0, byte0;\n" \
"}\n" : "=r"(out_fp16[0]) : "h"(src_packed));
return out;
}
#endif
template <typename PackedResultType, typename PackedSrcType>
CUTLASS_DEVICE
static PackedResultType packed_convert(PackedSrcType const &source) {
static_assert((platform::is_same<PackedSrcType, source_type_packed_2>::value &&
platform::is_same<PackedResultType, result_type_packed_2>::value) ||
(platform::is_same<PackedSrcType, source_type_packed_4>::value &&
platform::is_same<PackedResultType, result_type_packed_4>::value) ||
(platform::is_same<PackedSrcType, source_type_packed_8>::value &&
platform::is_same<PackedResultType, result_type_packed_8>::value),
"Invalid PackedSrcType/PackedResultType must be 2, 4 or 8 to use private convert dispatch.");
#if defined(CUDA_PTX_FP4FP6_CVT_ENABLED)
return ptx_convert(source);
#else
PackedResultType result;
NumericConverter<result_element, source_element, Round> converter;
const int k_packed = PackedResultType::kElements;
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < k_packed; ++i) {
result[i] = converter(source[i]);
}
return result;
#endif
}
friend class detail::VectorizedConverter;
public:
CUTLASS_DEVICE
static result_type convert(source_type const &source) {
result_type result;
using ConverterType = NumericArrayConverter<typename result_type::Element, typename source_type::Element, N, Round>;
detail::VectorizedConverter::convert<ConverterType,
result_type_packed_8, source_type_packed_8,
result_type_packed_4, source_type_packed_4,
result_type_packed_2, source_type_packed_2>(result, source);
return result;
}
CUTLASS_HOST_DEVICE
result_type operator()(source_type const &s) const {
return convert(s);
}
};
/// Partial specialization for Array<cutlass::float_e4m3_t, N> <= Array<cutlass::int2b_t, N>
template <FloatRoundStyle Round, int N>