@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
hpp 173 lines 5.22 KB
Raw
1 #pragma once
2
3 #include "torch_utils.h"
4
5 // This header is shared between _C (unstable ABI, used by machete) and
6 // _C_stable_libtorch (stable ABI, used by W4A8/sparse). TORCH_TARGET_VERSION
7 // is defined only for the stable target, so we switch includes and types
8 // accordingly. TorchTensor (not Tensor) avoids ambiguity with cute::Tensor.
9 #ifdef TORCH_TARGET_VERSION
10 #include <torch/csrc/stable/tensor.h>
11 #include <torch/headeronly/util/BFloat16.h>
12 #include <torch/headeronly/util/Half.h>
13 using TorchTensor = torch::stable::Tensor;
14 #else
15 using TorchTensor = torch::Tensor;
16 #endif
17
18 #include "cute/layout.hpp"
19 #include "cutlass/layout/matrix.h"
20 #include "cutlass/bfloat16.h"
21 #include "cutlass/half.h"
22
23 using ColumnMajor = typename cutlass::layout::ColumnMajor;
24 using RowMajor = typename cutlass::layout::RowMajor;
25
26 namespace cute {
27
28 namespace detail {
29
30 template <class T, class F, class G, int... I>
31 CUTE_HOST_DEVICE constexpr auto tapply_with_idx(T&& t, F&& f, G&& g,
32 seq<I...>) {
33 return g(f(cute::get<I>(static_cast<T&&>(t)), I)...);
34 }
35
36 template <class F, int... I>
37 CUTE_HOST_DEVICE constexpr auto make_shape_from_idx(F&& f, seq<I...>) {
38 return make_shape(f(I)...);
39 }
40
41 }; // namespace detail
42
43 template <class T, class F>
44 CUTE_HOST_DEVICE constexpr auto transform_with_idx(T const& t, F&& f) {
45 if constexpr (cute::is_tuple<T>::value) {
46 return detail::tapply_with_idx(
47 t, f, [](auto const&... a) { return cute::make_tuple(a...); },
48 tuple_seq<T>{});
49 } else {
50 return f(t);
51 }
52
53 CUTE_GCC_UNREACHABLE;
54 }
55
56 // calls: make_shape(f(0), f(1), ..., f(N-1))
57 template <int N, class F>
58 CUTE_HOST_DEVICE constexpr auto make_shape_from_idx(F&& f) {
59 return detail::make_shape_from_idx(f, make_seq<N>{});
60 }
61
62 }; // namespace cute
63
64 // Make a layout from a tensor with `rank(Stride{})`, where the shape is the
65 // shape of the passed in tensor and the strides are of type `Stride` and
66 // contain the strides of the passed in tensor, checking that any static strides
67 // in `Stride{}` match the strides of the passed in tensor.
68 // If `tensor.dim() < rank(Stride{})`, the shape is padded with 1s and the extra
69 // strides are set to be 0 or 1.
70 template <typename Stride>
71 static inline auto make_cute_layout(TorchTensor const& tensor,
72 std::string_view name = "tensor") {
73 TORCH_UTILS_CHECK(tensor.dim() <= rank(Stride{}));
74 auto stride = cute::transform_with_idx(Stride{}, [&](auto const& stride_ele,
75 auto const& idx) {
76 using StrideEle = std::decay_t<decltype(stride_ele)>;
77
78 if (idx < tensor.dim()) {
79 if constexpr (cute::is_static_v<StrideEle>) {
80 TORCH_UTILS_CHECK(StrideEle::value == tensor.stride(idx), "Expected ",
81 name, ".stride(", idx, ") to be ", StrideEle::value);
82 return StrideEle{};
83 } else {
84 if (tensor.size(idx) == 1) {
85 // use 0 stride for dim with size 1, this is easier for
86 // cute/cutlass to optimize (helps the TMA code flatten dims)
87 return StrideEle{0};
88 } else {
89 return tensor.stride(idx);
90 }
91 }
92 } else {
93 // Extra strides are assumed to be 0 or 1
94 if constexpr (cute::is_static_v<StrideEle>) {
95 static_assert(StrideEle::value == 0 || StrideEle::value == 1);
96 }
97 return StrideEle{};
98 }
99 });
100
101 auto shape = cute::make_shape_from_idx<rank(Stride{})>([&](auto const& idx) {
102 if (idx < tensor.dim())
103 return tensor.size(idx);
104 else
105 return int64_t(1);
106 });
107
108 return make_layout(shape, stride);
109 }
110
111 template <typename Stride>
112 static inline auto maybe_make_cute_layout(
113 std::optional<TorchTensor> const& tensor,
114 std::string_view name = "tensor") {
115 using Layout = decltype(make_cute_layout<Stride>(*tensor));
116
117 if (tensor) {
118 return std::optional<Layout>{make_cute_layout<Stride>(*tensor, name)};
119 } else {
120 return std::optional<Layout>{};
121 }
122 }
123
124 //
125 // Torch Type to Cutlass Type (equivalent_cutlass_type)
126 //
127
128 template <typename T>
129 struct equivalent_cutlass_type {
130 using type = T;
131 };
132
133 template <typename T>
134 using equivalent_cutlass_type_t = typename equivalent_cutlass_type<T>::type;
135
136 template <>
137 struct equivalent_cutlass_type<torch::headeronly::Half> {
138 using type = cutlass::half_t;
139 };
140
141 template <>
142 struct equivalent_cutlass_type<torch::headeronly::BFloat16> {
143 using type = cutlass::bfloat16_t;
144 };
145
146 //
147 // equivalent_scalar_t (basically inverse of equivalent_cutlass_type)
148 //
149
150 // Return a `torch::headeronly::CppTypeToScalarType<T>` compatible type, i.e.
151 // get the C++ type equivalent to T, e.g.: `cutlass::half_t -> Half`
152 template <typename T>
153 struct equivalent_scalar_type {
154 using type = T;
155 };
156
157 template <typename T>
158 using equivalent_scalar_type_t = typename equivalent_scalar_type<T>::type;
159
160 template <>
161 struct equivalent_scalar_type<cutlass::half_t> {
162 using type = torch::headeronly::Half;
163 };
164
165 template <>
166 struct equivalent_scalar_type<cutlass::bfloat16_t> {
167 using type = torch::headeronly::BFloat16;
168 };
169
170 // get equivalent torch::headeronly::ScalarType tag from compile time type
171 template <typename T>
172 static inline constexpr torch::headeronly::ScalarType equivalent_scalar_type_v =
173 torch::headeronly::CppTypeToScalarType<equivalent_scalar_type_t<T>>::value;