@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
cuh 148 lines 5.36 KB
Raw
1 #pragma once
2
3 // clang-format will break include orders
4 // clang-format off
5
6 #include "cutlass/cutlass.h"
7
8 #include "cute/tensor.hpp"
9 #include "cute/atom/mma_atom.hpp"
10 #include "cutlass/numeric_types.h"
11
12 #include "cutlass/gemm/device/gemm_universal_adapter.h"
13 #include "cutlass/gemm/kernel/gemm_universal.hpp"
14 #include "cutlass/epilogue/collective/collective_builder.hpp"
15 #include "cutlass/gemm/collective/collective_builder.hpp"
16
17 #include "libtorch_stable/core/math.hpp"
18 #include "libtorch_stable/cutlass_extensions/common.hpp"
19 // clang-format on
20
21 /*
22 Epilogues defined in,
23 csrc/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp,
24 must contain a public type named EVTCompute of type Sm90EVT, as well as a
25 static prepare_args function that constructs an EVTCompute::Arguments struct.
26 */
27
28 using namespace cute;
29
30 namespace vllm {
31
32 template <typename ElementAB_, typename ElementD_,
33 template <typename, typename, typename> typename Epilogue_,
34 typename TileShape, typename ClusterShape, typename KernelSchedule,
35 typename EpilogueSchedule>
36 struct cutlass_3x_gemm {
37 using ElementAB = ElementAB_;
38 using ElementD = ElementD_;
39 using ElementAcc =
40 typename std::conditional<std::is_same_v<ElementAB, int8_t>, int32_t,
41 float>::type;
42
43 using Epilogue = Epilogue_<ElementAcc, ElementD, TileShape>;
44
45 using StrideD = Stride<int64_t, Int<1>, Int<0>>;
46 using ElementC = void;
47 using StrideC = StrideD;
48
49 using EVTCompute = typename Epilogue::EVTCompute;
50
51 // These are the minimum alignments needed for the kernels to compile
52 static constexpr int AlignmentAB =
53 128 / cutlass::sizeof_bits<ElementAB>::value;
54 static constexpr int AlignmentCD =
55 128 / cutlass::sizeof_bits<ElementD>::value;
56
57 using CollectiveEpilogue =
58 typename cutlass::epilogue::collective::CollectiveBuilder<
59 cutlass::arch::Sm90, cutlass::arch::OpClassTensorOp, TileShape,
60 ClusterShape, cutlass::epilogue::collective::EpilogueTileAuto,
61 ElementAcc, float, ElementC, StrideC, AlignmentCD, ElementD, StrideD,
62 AlignmentCD, EpilogueSchedule, EVTCompute>::CollectiveOp;
63
64 static constexpr size_t CEStorageSize =
65 sizeof(typename CollectiveEpilogue::SharedStorage);
66 using Stages = typename cutlass::gemm::collective::StageCountAutoCarveout<
67 static_cast<int>(CEStorageSize)>;
68
69 // clang-format off
70 using CollectiveMainloop =
71 typename cutlass::gemm::collective::CollectiveBuilder<
72 cutlass::arch::Sm90, cutlass::arch::OpClassTensorOp,
73 ElementAB, cutlass::layout::RowMajor, AlignmentAB,
74 ElementAB, cutlass::layout::ColumnMajor, AlignmentAB,
75 ElementAcc, TileShape, ClusterShape,
76 Stages,
77 KernelSchedule>::CollectiveOp;
78 // clang-format on
79
80 using KernelType = enable_sm90_or_later<cutlass::gemm::kernel::GemmUniversal<
81 cute::Shape<int, int, int, int>, CollectiveMainloop, CollectiveEpilogue,
82 cutlass::gemm::PersistentScheduler>>;
83
84 struct GemmKernel : public KernelType {};
85 };
86
87 template <typename ElementAB_, typename ElementD_,
88 template <typename, typename, typename> typename Epilogue_,
89 typename TileShape, typename ClusterShape, typename KernelSchedule,
90 typename EpilogueSchedule>
91 struct cutlass_3x_gemm_sm120 {
92 using ElementAB = ElementAB_;
93 using LayoutA = cutlass::layout::RowMajor;
94 static constexpr int AlignmentA =
95 128 / cutlass::sizeof_bits<ElementAB>::value;
96
97 using LayoutB = cutlass::layout::ColumnMajor;
98 static constexpr int AlignmentB =
99 128 / cutlass::sizeof_bits<ElementAB>::value;
100
101 using ElementC = void;
102 using LayoutC = cutlass::layout::RowMajor;
103 static constexpr int AlignmentC =
104 128 / cutlass::sizeof_bits<ElementD_>::value;
105
106 using ElementD = ElementD_;
107 using LayoutD = cutlass::layout::RowMajor;
108 static constexpr int AlignmentD = AlignmentC;
109
110 using ElementAcc =
111 typename std::conditional<std::is_same_v<ElementAB, int8_t>, int32_t,
112 float>::type;
113 using Epilogue = Epilogue_<ElementAcc, ElementD, TileShape>;
114
115 // MMA type
116 using ElementAccumulator = float;
117
118 // Epilogue types
119 using ElementBias = cutlass::half_t;
120 using ElementCompute = float;
121 using ElementAux = ElementD;
122 using LayoutAux = LayoutD;
123 using ElementAmax = float;
124
125 using EVTCompute = typename Epilogue::EVTCompute;
126
127 using CollectiveEpilogue =
128 typename cutlass::epilogue::collective::CollectiveBuilder<
129 cutlass::arch::Sm120, cutlass::arch::OpClassTensorOp, TileShape,
130 ClusterShape, cutlass::epilogue::collective::EpilogueTileAuto,
131 ElementAccumulator, ElementCompute, ElementC, LayoutC, AlignmentC,
132 ElementD, LayoutD, AlignmentD, EpilogueSchedule,
133 EVTCompute>::CollectiveOp;
134
135 using CollectiveMainloop =
136 typename cutlass::gemm::collective::CollectiveBuilder<
137 cutlass::arch::Sm120, cutlass::arch::OpClassTensorOp, ElementAB,
138 LayoutA, AlignmentA, ElementAB, LayoutB, AlignmentB,
139 ElementAccumulator, TileShape, ClusterShape,
140 cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(
141 sizeof(typename CollectiveEpilogue::SharedStorage))>,
142 KernelSchedule>::CollectiveOp;
143
144 using GemmKernel = enable_sm120_family<cutlass::gemm::kernel::GemmUniversal<
145 Shape<int, int, int, int>, CollectiveMainloop, CollectiveEpilogue, void>>;
146 };
147
148 } // namespace vllm