| 1 | #pragma once |
| 2 | |
| 3 | #include "cutlass/gemm/collective/collective_builder.hpp" |
| 4 | |
| 5 | namespace cutlass::gemm::collective { |
| 6 | using namespace cute; |
| 7 | |
| 8 | // |
| 9 | // VLLMCollectiveBuilder is a wrapper around CollectiveBuilder that allows for |
| 10 | // for custom kernel tags, allowing you to build custom collectives. Without |
| 11 | // touching the cutlass library headers, using `CutlassKernelTag` will mean it |
| 12 | // will resort to using the standard cutlass collective builder. |
| 13 | // |
| 14 | |
| 15 | // Use the default Cutlass collective builder, i.e. use an unmodified cutless |
| 16 | // collective |
| 17 | struct CutlassKernelTag {}; |
| 18 | |
| 19 | template <class KernelTag, class ArchTag, class OpClass, class ElementA, |
| 20 | class GmemLayoutA, int AlignmentA, class ElementB, class GmemLayoutB, |
| 21 | int AlignmentB, class ElementAccumulator, class TileShape_MNK, |
| 22 | class ClusterShape_MNK, class StageCountType, |
| 23 | class KernelScheduleType, class Enable = void> |
| 24 | struct VLLMCollectiveBuilder { |
| 25 | static_assert(sizeof(ElementA) == 0, |
| 26 | "Could not build a collective for given parameters."); |
| 27 | }; |
| 28 | |
| 29 | template <class ArchTag, class OpClass, class ElementA, class GmemLayoutA, |
| 30 | int AlignmentA, class ElementB, class GmemLayoutB, int AlignmentB, |
| 31 | class ElementAccumulator, class TileShape_MNK, class ClusterShape_MNK, |
| 32 | class StageCountType, class KernelScheduleType> |
| 33 | struct VLLMCollectiveBuilder< |
| 34 | CutlassKernelTag, ArchTag, OpClass, ElementA, GmemLayoutA, AlignmentA, |
| 35 | ElementB, GmemLayoutB, AlignmentB, ElementAccumulator, TileShape_MNK, |
| 36 | ClusterShape_MNK, StageCountType, KernelScheduleType> { |
| 37 | using CollectiveOp = typename CollectiveBuilder< |
| 38 | ArchTag, OpClass, ElementA, GmemLayoutA, AlignmentA, ElementB, |
| 39 | GmemLayoutB, AlignmentB, ElementAccumulator, TileShape_MNK, |
| 40 | ClusterShape_MNK, StageCountType, KernelScheduleType>::CollectiveOp; |
| 41 | }; |
| 42 | |
| 43 | }; // namespace cutlass::gemm::collective |