@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
h 280 lines 9.53 KB
Raw
1 /******************************************************************************
2 * Copyright (c) 2023, Tri Dao.
3 ******************************************************************************/
4 // clang-format off
5 // adapted from https://github.com/state-spaces/mamba/blob/main/csrc/selective_scan/selective_scan.h
6
7 #pragma once
8
9 #ifndef USE_ROCM
10 #include <cuda_bf16.h>
11 #else
12 #include <hip/hip_bf16.h>
13 #endif
14 #include <cuda_fp16.h>
15
16 #include <torch/headeronly/util/Half.h>
17 #include <torch/headeronly/util/BFloat16.h>
18 ////////////////////////////////////////////////////////////////////////////////////////////////////
19
20 struct SSMParamsBase {
21 using index_t = size_t;
22
23 int batch, dim, seqlen, dstate, n_groups;
24 int dim_ngroups_ratio;
25 bool is_variable_B;
26 bool is_variable_C;
27 int64_t null_block_id;
28
29 bool delta_softplus;
30 bool cache_enabled;
31 int block_size;
32
33 index_t A_d_stride;
34 index_t A_dstate_stride;
35 index_t B_batch_stride;
36 index_t B_d_stride;
37 index_t B_dstate_stride;
38 index_t B_group_stride;
39 index_t C_batch_stride;
40 index_t C_d_stride;
41 index_t C_dstate_stride;
42 index_t C_group_stride;
43 index_t u_batch_stride;
44 index_t u_d_stride;
45 index_t delta_batch_stride;
46 index_t delta_d_stride;
47 index_t z_batch_stride;
48 index_t z_d_stride;
49 index_t out_batch_stride;
50 index_t out_d_stride;
51 index_t out_z_batch_stride;
52 index_t out_z_d_stride;
53 index_t ssm_states_batch_stride;
54 index_t ssm_states_dim_stride;
55 index_t ssm_states_dstate_stride;
56 index_t cache_indices_stride;
57
58 // Common data pointers.
59 void *__restrict__ A_ptr;
60 void *__restrict__ B_ptr;
61 void *__restrict__ C_ptr;
62 void *__restrict__ D_ptr;
63 void *__restrict__ u_ptr;
64 void *__restrict__ delta_ptr;
65 void *__restrict__ delta_bias_ptr;
66 void *__restrict__ out_ptr;
67 void *__restrict__ ssm_states_ptr;
68 void *__restrict__ z_ptr;
69 void *__restrict__ out_z_ptr;
70
71 void *__restrict__ query_start_loc_ptr;
72 void *__restrict__ cache_indices_ptr;
73 void *__restrict__ has_initial_state_ptr;
74
75 void *__restrict__ block_idx_first_scheduled_token_ptr; // (batch,) - first block to write
76 void *__restrict__ block_idx_last_scheduled_token_ptr; // (batch,) - last block to write
77 void *__restrict__ initial_state_idx_ptr; // (batch,) - index of the initial state to use
78 void *__restrict__ cu_chunk_seqlen_ptr; // (nchunks+1,) - cumulative chunk token offsets
79 void *__restrict__ last_chunk_indices_ptr; // (batch,) - index of last chunk per sequence
80 };
81
82
83
84
85 #ifndef USE_ROCM
86
87 constexpr size_t custom_max(std::initializer_list<size_t> ilist)
88 {
89 return std::max(ilist);
90 }
91
92 template<typename T>
93 constexpr T constexpr_min(T a, T b) {
94 return std::min(a, b);
95 }
96
97 #else
98 constexpr size_t custom_max(std::initializer_list<size_t> ilist)
99 {
100 return *std::max_element(ilist.begin(), ilist.end());
101 }
102
103 template<typename T>
104 constexpr T constexpr_min(T a, T b) {
105 return a < b ? a : b;
106 }
107 #endif
108
109
110 #define MAX_DSTATE 256
111
112
113 inline __device__ float2 operator+(const float2 & a, const float2 & b){
114 return {a.x + b.x, a.y + b.y};
115 }
116
117 inline __device__ float3 operator+(const float3 &a, const float3 &b) {
118 return {a.x + b.x, a.y + b.y, a.z + b.z};
119 }
120
121 inline __device__ float4 operator+(const float4 & a, const float4 & b){
122 return {a.x + b.x, a.y + b.y, a.z + b.z, a.w + b.w};
123 }
124
125 ////////////////////////////////////////////////////////////////////////////////////////////////////
126
127 template<int BYTES> struct BytesToType {};
128
129 template<> struct BytesToType<16> {
130 using Type = uint4;
131 static_assert(sizeof(Type) == 16);
132 };
133
134 template<> struct BytesToType<8> {
135 using Type = uint64_t;
136 static_assert(sizeof(Type) == 8);
137 };
138
139 template<> struct BytesToType<4> {
140 using Type = uint32_t;
141 static_assert(sizeof(Type) == 4);
142 };
143
144 template<> struct BytesToType<2> {
145 using Type = uint16_t;
146 static_assert(sizeof(Type) == 2);
147 };
148
149 template<> struct BytesToType<1> {
150 using Type = uint8_t;
151 static_assert(sizeof(Type) == 1);
152 };
153
154 ////////////////////////////////////////////////////////////////////////////////////////////////////
155
156 template<typename scalar_t, int N>
157 struct Converter{
158 static inline __device__ void to_float(const scalar_t (&src)[N], float (&dst)[N]) {
159 #pragma unroll
160 for (int i = 0; i < N; ++i) { dst[i] = src[i]; }
161 }
162 };
163
164 template<int N>
165 struct Converter<torch::headeronly::Half, N>{
166 static inline __device__ void to_float(const torch::headeronly::Half (&src)[N], float (&dst)[N]) {
167 static_assert(N % 2 == 0);
168 auto &src2 = reinterpret_cast<const half2 (&)[N / 2]>(src);
169 auto &dst2 = reinterpret_cast<float2 (&)[N / 2]>(dst);
170 #pragma unroll
171 for (int i = 0; i < N / 2; ++i) { dst2[i] = __half22float2(src2[i]); }
172 }
173 };
174
175 #if __CUDA_ARCH__ >= 800
176 template<int N>
177 struct Converter<torch::headeronly::BFloat16, N>{
178 static inline __device__ void to_float(const torch::headeronly::BFloat16 (&src)[N], float (&dst)[N]) {
179 static_assert(N % 2 == 0);
180 auto &src2 = reinterpret_cast<const nv_bfloat162 (&)[N / 2]>(src);
181 auto &dst2 = reinterpret_cast<float2 (&)[N / 2]>(dst);
182 #pragma unroll
183 for (int i = 0; i < N / 2; ++i) { dst2[i] = __bfloat1622float2(src2[i]); }
184 }
185 };
186 #endif
187
188 ////////////////////////////////////////////////////////////////////////////////////////////////////
189
190
191 template<typename scalar_t> struct SSMScanOp;
192
193 template<>
194 struct SSMScanOp<float> {
195 __device__ __forceinline__ float2 operator()(const float2 &ab0, const float2 &ab1) const {
196 return make_float2(ab1.x * ab0.x, ab1.x * ab0.y + ab1.y);
197 }
198 };
199
200 // A stateful callback functor that maintains a running prefix to be applied
201 // during consecutive scan operations.
202 template <typename scalar_t> struct SSMScanPrefixCallbackOp {
203 using scan_t = std::conditional_t<std::is_same_v<scalar_t, float>, float2, float4>;
204 scan_t running_prefix;
205 // Constructor
206 __device__ SSMScanPrefixCallbackOp(scan_t running_prefix_) : running_prefix(running_prefix_) {}
207 // Callback operator to be entered by the first warp of threads in the block.
208 // Thread-0 is responsible for returning a value for seeding the block-wide scan.
209 __device__ scan_t operator()(scan_t block_aggregate) {
210 scan_t old_prefix = running_prefix;
211 running_prefix = SSMScanOp<scalar_t>()(running_prefix, block_aggregate);
212 return old_prefix;
213 }
214 };
215
216 ////////////////////////////////////////////////////////////////////////////////////////////////////
217
218 template<typename Ktraits>
219 inline __device__ void load_input(typename Ktraits::input_t *u,
220 typename Ktraits::input_t (&u_vals)[Ktraits::kNItems],
221 typename Ktraits::BlockLoadT::TempStorage &smem_load,
222 int seqlen) {
223 if constexpr (Ktraits::kIsEvenLen && !Ktraits::kVarlen) {
224 auto& smem_load_vec = reinterpret_cast<typename Ktraits::BlockLoadVecT::TempStorage&>(smem_load);
225 using vec_t = typename Ktraits::vec_t;
226 typename Ktraits::BlockLoadVecT(smem_load_vec).Load(
227 reinterpret_cast<vec_t*>(u),
228 reinterpret_cast<vec_t(&)[Ktraits::kNLoads]>(u_vals)
229 #ifdef USE_ROCM
230 , Ktraits::kNThreads * Ktraits::kNLoads
231 #endif
232
233 );
234 } else {
235 typename Ktraits::BlockLoadT(smem_load).Load(u, u_vals, seqlen, 0.f);
236 }
237 }
238
239
240 template<typename Ktraits>
241 inline __device__ void load_weight(typename Ktraits::input_t *Bvar,
242 typename Ktraits::weight_t (&B_vals)[Ktraits::kNItems],
243 typename Ktraits::BlockLoadWeightT::TempStorage &smem_load_weight,
244 int seqlen) {
245 constexpr int kNItems = Ktraits::kNItems;
246 typename Ktraits::input_t B_vals_load[kNItems];
247 if constexpr (Ktraits::kIsEvenLen && !Ktraits::kVarlen) {
248 auto& smem_load_weight_vec = reinterpret_cast<typename Ktraits::BlockLoadWeightVecT::TempStorage&>(smem_load_weight);
249 using vec_t = typename Ktraits::vec_t;
250 typename Ktraits::BlockLoadWeightVecT(smem_load_weight_vec).Load(
251 reinterpret_cast<vec_t*>(Bvar),
252 reinterpret_cast<vec_t(&)[Ktraits::kNLoads]>(B_vals_load)
253 );
254 } else {
255 typename Ktraits::BlockLoadWeightT(smem_load_weight).Load(Bvar, B_vals_load, seqlen, 0.f);
256 }
257 // #pragma unroll
258 // for (int i = 0; i < kNItems; ++i) { B_vals[i] = B_vals_load[i]; }
259 Converter<typename Ktraits::input_t, kNItems>::to_float(B_vals_load, B_vals);
260 }
261
262 template<typename Ktraits>
263 inline __device__ void store_output(typename Ktraits::input_t *out,
264 const float (&out_vals)[Ktraits::kNItems],
265 typename Ktraits::BlockStoreT::TempStorage &smem_store,
266 int seqlen) {
267 typename Ktraits::input_t write_vals[Ktraits::kNItems];
268 #pragma unroll
269 for (int i = 0; i < Ktraits::kNItems; ++i) { write_vals[i] = out_vals[i]; }
270 if constexpr (Ktraits::kIsEvenLen && !Ktraits::kVarlen) {
271 auto& smem_store_vec = reinterpret_cast<typename Ktraits::BlockStoreVecT::TempStorage&>(smem_store);
272 using vec_t = typename Ktraits::vec_t;
273 typename Ktraits::BlockStoreVecT(smem_store_vec).Store(
274 reinterpret_cast<vec_t*>(out),
275 reinterpret_cast<vec_t(&)[Ktraits::kNLoads]>(write_vals)
276 );
277 } else {
278 typename Ktraits::BlockStoreT(smem_store).Store(out, write_vals, seqlen);
279 }
280 }