| 1 | // clang-format off |
| 2 | // adapted from https://github.com/state-spaces/mamba/blob/main/csrc/selective_scan/selective_scan_fwd_kernel.cuh |
| 3 | #include "../torch_utils.h" |
| 4 | #include <torch/csrc/stable/macros.h> |
| 5 | #include "selective_scan.h" |
| 6 | |
| 7 | #ifndef USE_ROCM |
| 8 | #include <cub/block/block_load.cuh> |
| 9 | #include <cub/block/block_store.cuh> |
| 10 | #include <cub/block/block_scan.cuh> |
| 11 | #else |
| 12 | #include <hipcub/hipcub.hpp> |
| 13 | namespace cub = hipcub; |
| 14 | #endif |
| 15 | |
| 16 | #include "selective_scan.h" |
| 17 | #include "static_switch.h" |
| 18 | |
| 19 | template<int kNThreads_, int kNItems_, int kNRows_, bool kIsEvenLen_, |
| 20 | bool kIsVariableB_, bool kIsVariableC_, |
| 21 | bool kHasZ_, bool kVarlen_, typename input_t_, typename weight_t_, typename state_t_> |
| 22 | struct Selective_Scan_fwd_kernel_traits { |
| 23 | static_assert(kNItems_ % 4 == 0); |
| 24 | using input_t = input_t_; |
| 25 | using weight_t = weight_t_; |
| 26 | using state_t = state_t_; |
| 27 | static constexpr int kNThreads = kNThreads_; |
| 28 | // Setting MinBlocksPerMP to be 3 (instead of 2) for 128 threads improves occupancy. |
| 29 | static constexpr int kMinBlocks = kNThreads < 128 ? 5 : 3; |
| 30 | static constexpr int kNItems = kNItems_; |
| 31 | static constexpr int kNRows = kNRows_; |
| 32 | static constexpr int kNBytes = sizeof(input_t); |
| 33 | static_assert(kNBytes == 2 || kNBytes == 4); |
| 34 | static constexpr int kNElts = kNBytes == 4 ? 4 : constexpr_min(8, kNItems); |
| 35 | static_assert(kNItems % kNElts == 0); |
| 36 | static constexpr int kNLoads = kNItems / kNElts; |
| 37 | static constexpr bool kIsEvenLen = kVarlen_ ? false : kIsEvenLen_; |
| 38 | static constexpr bool kIsVariableB = kIsVariableB_; |
| 39 | static constexpr bool kIsVariableC = kIsVariableC_; |
| 40 | static constexpr bool kHasZ = kHasZ_; |
| 41 | static constexpr bool kVarlen = kVarlen_; |
| 42 | |
| 43 | static constexpr bool kDirectIO = kVarlen_ ? false : kIsEvenLen && kNLoads == 1; |
| 44 | static constexpr int kNLoadsIndex = kNItems / 4; |
| 45 | using vec_t = typename BytesToType<kNBytes * kNElts>::Type; |
| 46 | using scan_t = float2; |
| 47 | using BlockLoadT = cub::BlockLoad<input_t, kNThreads, kNItems, cub::BLOCK_LOAD_WARP_TRANSPOSE>; |
| 48 | using BlockLoadVecT = cub::BlockLoad<vec_t, kNThreads, kNLoads, |
| 49 | !kDirectIO ? cub::BLOCK_LOAD_WARP_TRANSPOSE : cub::BLOCK_LOAD_DIRECT>; |
| 50 | using BlockLoadWeightT = cub::BlockLoad<input_t, kNThreads, kNItems , cub::BLOCK_LOAD_WARP_TRANSPOSE>; |
| 51 | using BlockLoadWeightVecT = cub::BlockLoad<vec_t, kNThreads, kNLoads , |
| 52 | !kDirectIO ? cub::BLOCK_LOAD_WARP_TRANSPOSE : cub::BLOCK_LOAD_DIRECT>; |
| 53 | using BlockStoreT = cub::BlockStore<input_t, kNThreads, kNItems, cub::BLOCK_STORE_WARP_TRANSPOSE>; |
| 54 | using BlockStoreVecT = cub::BlockStore<vec_t, kNThreads, kNLoads, |
| 55 | !kDirectIO ? cub::BLOCK_STORE_WARP_TRANSPOSE : cub::BLOCK_STORE_DIRECT>; |
| 56 | // using BlockScanT = cub::BlockScan<scan_t, kNThreads, cub::BLOCK_SCAN_RAKING_MEMOIZE>; |
| 57 | // using BlockScanT = cub::BlockScan<scan_t, kNThreads, cub::BLOCK_SCAN_RAKING>; |
| 58 | using BlockScanT = cub::BlockScan<scan_t, kNThreads, cub::BLOCK_SCAN_WARP_SCANS>; |
| 59 | static constexpr int kSmemIOSize = custom_max({sizeof(typename BlockLoadT::TempStorage), |
| 60 | sizeof(typename BlockLoadVecT::TempStorage), |
| 61 | (int(kIsVariableB) + int(kIsVariableC)) * sizeof(typename BlockLoadWeightT::TempStorage), |
| 62 | (int(kIsVariableB) + int(kIsVariableC)) * sizeof(typename BlockLoadWeightVecT::TempStorage), |
| 63 | sizeof(typename BlockStoreT::TempStorage), |
| 64 | sizeof(typename BlockStoreVecT::TempStorage)}); |
| 65 | static constexpr int kSmemSize = kSmemIOSize + sizeof(typename BlockScanT::TempStorage); |
| 66 | }; |
| 67 | |
| 68 | template<typename Ktraits> |
| 69 | __global__ __launch_bounds__(Ktraits::kNThreads, Ktraits::kMinBlocks) |
| 70 | void selective_scan_fwd_kernel(SSMParamsBase params) { |
| 71 | constexpr bool kIsVariableB = Ktraits::kIsVariableB; |
| 72 | constexpr bool kIsVariableC = Ktraits::kIsVariableC; |
| 73 | constexpr bool kHasZ = Ktraits::kHasZ; |
| 74 | constexpr bool kVarlen = Ktraits::kVarlen; |
| 75 | constexpr int kNItems = Ktraits::kNItems; |
| 76 | constexpr int kNRows = Ktraits::kNRows; |
| 77 | constexpr bool kDirectIO = Ktraits::kDirectIO; |
| 78 | using input_t = typename Ktraits::input_t; |
| 79 | using weight_t = typename Ktraits::weight_t; |
| 80 | using scan_t = typename Ktraits::scan_t; |
| 81 | |
| 82 | // Shared memory. |
| 83 | extern __shared__ char smem_[]; |
| 84 | // cast to lvalue reference of expected type |
| 85 | // char *smem_loadstorescan = smem_ + 2 * MAX_DSTATE * sizeof(weight_t); |
| 86 | // auto& smem_load = reinterpret_cast<typename BlockLoadT::TempStorage&>(smem_ + 2 * MAX_DSTATE * sizeof(weight_t)); |
| 87 | // auto& smem_load = reinterpret_cast<typename BlockLoadT::TempStorage&>(smem_loadstorescan); |
| 88 | auto& smem_load = reinterpret_cast<typename Ktraits::BlockLoadT::TempStorage&>(smem_); |
| 89 | auto& smem_load_weight = reinterpret_cast<typename Ktraits::BlockLoadWeightT::TempStorage&>(smem_); |
| 90 | auto& smem_load_weight1 = *reinterpret_cast<typename Ktraits::BlockLoadWeightT::TempStorage*>(smem_ + sizeof(typename Ktraits::BlockLoadWeightT::TempStorage)); |
| 91 | auto& smem_store = reinterpret_cast<typename Ktraits::BlockStoreT::TempStorage&>(smem_); |
| 92 | auto& smem_scan = *reinterpret_cast<typename Ktraits::BlockScanT::TempStorage*>(smem_ + Ktraits::kSmemIOSize); |
| 93 | // weight_t *smem_a = reinterpret_cast<weight_t *>(smem_ + smem_loadstorescan_size); |
| 94 | // weight_t *smem_bc = reinterpret_cast<weight_t *>(smem_a + MAX_DSTATE); |
| 95 | scan_t *smem_running_prefix = reinterpret_cast<scan_t *>(smem_ + Ktraits::kSmemSize); |
| 96 | |
| 97 | const int batch_id = blockIdx.x; |
| 98 | const int dim_id = blockIdx.y; |
| 99 | const int group_id = dim_id / (params.dim_ngroups_ratio); |
| 100 | int seqlen = params.seqlen; |
| 101 | int sequence_start_index = batch_id; |
| 102 | if constexpr (kVarlen){ |
| 103 | int *query_start_loc = reinterpret_cast<int *>(params.query_start_loc_ptr); |
| 104 | sequence_start_index = query_start_loc[batch_id]; |
| 105 | seqlen = query_start_loc[batch_id + 1] - sequence_start_index; |
| 106 | } |
| 107 | const bool has_initial_state = params.has_initial_state_ptr == nullptr ? false |
| 108 | : reinterpret_cast<bool *>(params.has_initial_state_ptr)[batch_id]; |
| 109 | |
| 110 | const int* cache_indices = params.cache_indices_ptr == nullptr ? nullptr |
| 111 | : reinterpret_cast<int *>(params.cache_indices_ptr); |
| 112 | int cache_index; |
| 113 | if (cache_indices == nullptr) { |
| 114 | cache_index = batch_id; |
| 115 | } else if (params.cache_enabled) { |
| 116 | const int* initial_state_idx = reinterpret_cast<const int*>(params.initial_state_idx_ptr); |
| 117 | cache_index = cache_indices[batch_id * params.cache_indices_stride + initial_state_idx[batch_id]]; |
| 118 | } else { |
| 119 | cache_index = cache_indices[batch_id]; |
| 120 | } |
| 121 | // Skip batch entries whose cache index maps to the null block (padding). |
| 122 | if (cache_indices != nullptr && cache_index == params.null_block_id){ |
| 123 | return; |
| 124 | } |
| 125 | input_t *u = reinterpret_cast<input_t *>(params.u_ptr) + sequence_start_index * params.u_batch_stride |
| 126 | + dim_id * kNRows * params.u_d_stride; |
| 127 | input_t *delta = reinterpret_cast<input_t *>(params.delta_ptr) + sequence_start_index * params.delta_batch_stride |
| 128 | + dim_id * kNRows * params.delta_d_stride; |
| 129 | weight_t *A = reinterpret_cast<weight_t *>(params.A_ptr) + dim_id * kNRows * params.A_d_stride; |
| 130 | weight_t *B = reinterpret_cast<weight_t *>(params.B_ptr) + dim_id * kNRows * params.B_d_stride; |
| 131 | input_t *Bvar = reinterpret_cast<input_t *>(params.B_ptr) + sequence_start_index * params.B_batch_stride + group_id * params.B_group_stride; |
| 132 | weight_t *C = reinterpret_cast<weight_t *>(params.C_ptr) + dim_id * kNRows * params.C_d_stride; |
| 133 | input_t *Cvar = reinterpret_cast<input_t *>(params.C_ptr) + sequence_start_index * params.C_batch_stride + group_id * params.C_group_stride; |
| 134 | |
| 135 | typename Ktraits::state_t *ssm_states; |
| 136 | if (params.cache_enabled) { |
| 137 | // APC mode: ssm_states points to the base, we'll use absolute cache slots later |
| 138 | ssm_states = reinterpret_cast<typename Ktraits::state_t *>(params.ssm_states_ptr) + |
| 139 | dim_id * kNRows * params.ssm_states_dim_stride; |
| 140 | } else { |
| 141 | // Non-APC mode: offset by cache_index as before |
| 142 | ssm_states = reinterpret_cast<typename Ktraits::state_t *>(params.ssm_states_ptr) + |
| 143 | cache_index * params.ssm_states_batch_stride + |
| 144 | dim_id * kNRows * params.ssm_states_dim_stride; |
| 145 | } |
| 146 | |
| 147 | float D_val[kNRows] = {0}; |
| 148 | if (params.D_ptr != nullptr) { |
| 149 | #pragma unroll |
| 150 | for (int r = 0; r < kNRows; ++r) { |
| 151 | D_val[r] = reinterpret_cast<float *>(params.D_ptr)[dim_id * kNRows + r]; |
| 152 | } |
| 153 | } |
| 154 | float delta_bias[kNRows] = {0}; |
| 155 | if (params.delta_bias_ptr != nullptr) { |
| 156 | #pragma unroll |
| 157 | for (int r = 0; r < kNRows; ++r) { |
| 158 | delta_bias[r] = reinterpret_cast<float *>(params.delta_bias_ptr)[dim_id * kNRows + r]; |
| 159 | } |
| 160 | } |
| 161 | |
| 162 | // Use block_size for chunking when APC is enabled, otherwise use 2048 for backwards compatibility |
| 163 | const int block_size = params.cache_enabled ? params.block_size : 2048; |
| 164 | |
| 165 | const int* batch_cache_indices = cache_indices != nullptr ? |
| 166 | cache_indices + batch_id * params.cache_indices_stride : nullptr; |
| 167 | const int* block_idx_first_scheduled = params.block_idx_first_scheduled_token_ptr != nullptr ? |
| 168 | reinterpret_cast<const int*>(params.block_idx_first_scheduled_token_ptr) : nullptr; |
| 169 | const int* block_idx_last_scheduled = params.block_idx_last_scheduled_token_ptr != nullptr ? |
| 170 | reinterpret_cast<const int*>(params.block_idx_last_scheduled_token_ptr) : nullptr; |
| 171 | const int* initial_state_idx = params.initial_state_idx_ptr != nullptr ? |
| 172 | reinterpret_cast<const int*>(params.initial_state_idx_ptr) : nullptr; |
| 173 | const int* cu_chunk_seqlen = params.cu_chunk_seqlen_ptr != nullptr ? |
| 174 | reinterpret_cast<const int*>(params.cu_chunk_seqlen_ptr) : nullptr; |
| 175 | const int* last_chunk_indices = params.last_chunk_indices_ptr != nullptr ? |
| 176 | reinterpret_cast<const int*>(params.last_chunk_indices_ptr) : nullptr; |
| 177 | |
| 178 | const size_t load_cache_slot = params.cache_enabled && batch_cache_indices != nullptr ? batch_cache_indices[initial_state_idx[batch_id]] : cache_index; |
| 179 | |
| 180 | const int block_idx_first = (params.cache_enabled && block_idx_first_scheduled != nullptr) ? |
| 181 | block_idx_first_scheduled[batch_id] : 0; |
| 182 | |
| 183 | // Determine chunk boundaries from pre-computed metadata (APC mode) |
| 184 | // or fall back to simple block_size chunking. |
| 185 | int first_chunk_idx, n_chunks; |
| 186 | int current_position; |
| 187 | |
| 188 | if (cu_chunk_seqlen != nullptr && last_chunk_indices != nullptr) { |
| 189 | const int last_chunk_idx = last_chunk_indices[batch_id]; |
| 190 | first_chunk_idx = (batch_id == 0) ? 0 : last_chunk_indices[batch_id - 1] + 1; |
| 191 | n_chunks = last_chunk_idx - first_chunk_idx + 1; |
| 192 | // Derive current_position: if the first chunk is partial (fills remainder |
| 193 | // of a started block), offset into the block accordingly. |
| 194 | const int first_chunk_tokens = cu_chunk_seqlen[first_chunk_idx + 1] - cu_chunk_seqlen[first_chunk_idx]; |
| 195 | const int chunk_start_offset = (n_chunks > 1 && first_chunk_tokens < block_size) |
| 196 | ? (block_size - first_chunk_tokens) : 0; |
| 197 | current_position = block_idx_first * block_size + chunk_start_offset; |
| 198 | } else { |
| 199 | first_chunk_idx = 0; |
| 200 | n_chunks = (seqlen + block_size - 1) / block_size; |
| 201 | current_position = 0; |
| 202 | } |
| 203 | |
| 204 | int tokens_processed = 0; |
| 205 | |
| 206 | for (int chunk = 0; chunk < n_chunks; ++chunk) { |
| 207 | const int chunk_tokens = (cu_chunk_seqlen != nullptr) |
| 208 | ? cu_chunk_seqlen[first_chunk_idx + chunk + 1] - cu_chunk_seqlen[first_chunk_idx + chunk] |
| 209 | : min(block_size, seqlen - tokens_processed); |
| 210 | if (chunk_tokens <= 0) break; |
| 211 | input_t u_vals[kNRows][kNItems], delta_vals_load[kNRows][kNItems]; |
| 212 | |
| 213 | __syncthreads(); |
| 214 | #pragma unroll |
| 215 | for (int r = 0; r < kNRows; ++r) { |
| 216 | if constexpr (!kDirectIO) { |
| 217 | if (r > 0) { __syncthreads(); } |
| 218 | } |
| 219 | load_input<Ktraits>(u + r * params.u_d_stride, u_vals[r], smem_load, chunk_tokens); |
| 220 | if constexpr (!kDirectIO) { __syncthreads(); } |
| 221 | load_input<Ktraits>(delta + r * params.delta_d_stride, delta_vals_load[r], smem_load, chunk_tokens); |
| 222 | } |
| 223 | u += chunk_tokens; |
| 224 | delta += chunk_tokens; |
| 225 | |
| 226 | float delta_vals[kNRows][kNItems], delta_u_vals[kNRows][kNItems], out_vals[kNRows][kNItems]; |
| 227 | #pragma unroll |
| 228 | for (int r = 0; r < kNRows; ++r) { |
| 229 | #pragma unroll |
| 230 | for (int i = 0; i < kNItems; ++i) { |
| 231 | float u_val = float(u_vals[r][i]); |
| 232 | delta_vals[r][i] = float(delta_vals_load[r][i]) + delta_bias[r]; |
| 233 | if (params.delta_softplus) { |
| 234 | delta_vals[r][i] = delta_vals[r][i] <= 20.f ? log1pf(expf(delta_vals[r][i])) : delta_vals[r][i]; |
| 235 | } |
| 236 | delta_u_vals[r][i] = delta_vals[r][i] * u_val; |
| 237 | out_vals[r][i] = D_val[r] * u_val; |
| 238 | } |
| 239 | } |
| 240 | |
| 241 | __syncthreads(); |
| 242 | for (int state_idx = 0; state_idx < params.dstate; ++state_idx) { |
| 243 | weight_t A_val[kNRows]; |
| 244 | #pragma unroll |
| 245 | for (int r = 0; r < kNRows; ++r) { |
| 246 | A_val[r] = A[state_idx * params.A_dstate_stride + r * params.A_d_stride]; |
| 247 | // Multiply the real part of A with LOG2E so we can use exp2f instead of expf. |
| 248 | constexpr float kLog2e = M_LOG2E; |
| 249 | A_val[r] *= kLog2e; |
| 250 | } |
| 251 | // This variable holds B * C if both B and C are constant across seqlen. If only B varies |
| 252 | // across seqlen, this holds C. If only C varies across seqlen, this holds B. |
| 253 | // If both B and C vary, this is unused. |
| 254 | weight_t BC_val[kNRows]; |
| 255 | weight_t B_vals[kNItems], C_vals[kNItems]; |
| 256 | if constexpr (kIsVariableB) { |
| 257 | load_weight<Ktraits>(Bvar + state_idx * params.B_dstate_stride, B_vals, |
| 258 | smem_load_weight, chunk_tokens); |
| 259 | if constexpr (!kIsVariableC) { |
| 260 | #pragma unroll |
| 261 | for (int r = 0; r < kNRows; ++r) { |
| 262 | BC_val[r] = C[state_idx * params.C_dstate_stride + r * params.C_d_stride]; |
| 263 | } |
| 264 | } |
| 265 | } |
| 266 | if constexpr (kIsVariableC) { |
| 267 | auto &smem_load_weight_C = !kIsVariableB ? smem_load_weight : smem_load_weight1; |
| 268 | load_weight<Ktraits>(Cvar + state_idx * params.C_dstate_stride, C_vals, |
| 269 | smem_load_weight_C, chunk_tokens); |
| 270 | if constexpr (!kIsVariableB) { |
| 271 | #pragma unroll |
| 272 | for (int r = 0; r < kNRows; ++r) { |
| 273 | BC_val[r] = B[state_idx * params.B_dstate_stride + r * params.B_d_stride]; |
| 274 | } |
| 275 | } |
| 276 | } |
| 277 | if constexpr (!kIsVariableB && !kIsVariableC) { |
| 278 | #pragma unroll |
| 279 | for (int r = 0; r < kNRows; ++r) { |
| 280 | BC_val[r] = B[state_idx * params.B_dstate_stride + r * params.B_d_stride] * C[state_idx * params.C_dstate_stride + r * params.C_d_stride]; |
| 281 | } |
| 282 | } |
| 283 | |
| 284 | #pragma unroll |
| 285 | for (int r = 0; r < kNRows; ++r) { |
| 286 | if (r > 0) { __syncthreads(); } // Scan could be using the same smem |
| 287 | scan_t thread_data[kNItems]; |
| 288 | #pragma unroll |
| 289 | for (int i = 0; i < kNItems; ++i) { |
| 290 | thread_data[i] = make_float2(exp2f(delta_vals[r][i] * A_val[r]), |
| 291 | !kIsVariableB ? delta_u_vals[r][i] : B_vals[i] * delta_u_vals[r][i]); |
| 292 | if (threadIdx.x * kNItems + i >= chunk_tokens) { |
| 293 | thread_data[i] = make_float2(1.f, 0.f); |
| 294 | } |
| 295 | } |
| 296 | // Initialize running total |
| 297 | scan_t running_prefix; |
| 298 | if (chunk > 0) { |
| 299 | running_prefix = smem_running_prefix[state_idx + r * MAX_DSTATE]; |
| 300 | } else { |
| 301 | // Load initial state |
| 302 | if (params.cache_enabled && has_initial_state && batch_cache_indices != nullptr) { |
| 303 | size_t state_offset = load_cache_slot * params.ssm_states_batch_stride + |
| 304 | r * params.ssm_states_dim_stride + |
| 305 | state_idx * params.ssm_states_dstate_stride; |
| 306 | running_prefix = make_float2(1.0, float(ssm_states[state_offset])); |
| 307 | } else if (has_initial_state) { |
| 308 | // Non-APC mode: load from current batch position |
| 309 | running_prefix = make_float2(1.0, float(ssm_states[state_idx * params.ssm_states_dstate_stride])); |
| 310 | } else { |
| 311 | // No initial state |
| 312 | running_prefix = make_float2(1.0, 0.0); |
| 313 | } |
| 314 | } |
| 315 | |
| 316 | SSMScanPrefixCallbackOp<weight_t> prefix_op(running_prefix); |
| 317 | typename Ktraits::BlockScanT(smem_scan).InclusiveScan( |
| 318 | thread_data, thread_data, SSMScanOp<weight_t>(), prefix_op |
| 319 | ); |
| 320 | // There's a syncthreads in the scan op, so we don't need to sync here. |
| 321 | // Unless there's only 1 warp, but then it's the same thread (0) reading and writing. |
| 322 | if (threadIdx.x == 0) { |
| 323 | smem_running_prefix[state_idx + r * MAX_DSTATE] = prefix_op.running_prefix; |
| 324 | |
| 325 | // Store state at the end of each aligned chunk when cache is enabled |
| 326 | if (params.cache_enabled && batch_cache_indices != nullptr) { |
| 327 | size_t cache_slot; |
| 328 | if (chunk == n_chunks - 1) { |
| 329 | cache_slot = batch_cache_indices[block_idx_last_scheduled[batch_id]]; |
| 330 | } else { |
| 331 | const int block_idx_completed = (current_position + chunk_tokens - 1) / block_size; |
| 332 | cache_slot = batch_cache_indices[block_idx_completed]; |
| 333 | } |
| 334 | |
| 335 | size_t state_offset = cache_slot * params.ssm_states_batch_stride + |
| 336 | r * params.ssm_states_dim_stride + |
| 337 | state_idx * params.ssm_states_dstate_stride; |
| 338 | |
| 339 | ssm_states[state_offset] = typename Ktraits::state_t(prefix_op.running_prefix.y); |
| 340 | } else if (!params.cache_enabled && chunk == n_chunks - 1) { |
| 341 | // Non-APC mode: store only final state at current batch position |
| 342 | ssm_states[state_idx * params.ssm_states_dstate_stride] = typename Ktraits::state_t(prefix_op.running_prefix.y); |
| 343 | } |
| 344 | } |
| 345 | #pragma unroll |
| 346 | for (int i = 0; i < kNItems; ++i) { |
| 347 | const weight_t C_val = !kIsVariableC |
| 348 | ? BC_val[r] |
| 349 | : (!kIsVariableB ? BC_val[r] * C_vals[i] : C_vals[i]); |
| 350 | out_vals[r][i] += thread_data[i].y * C_val; |
| 351 | } |
| 352 | } |
| 353 | } |
| 354 | input_t *out = reinterpret_cast<input_t *>(params.out_ptr) + sequence_start_index * params.out_batch_stride |
| 355 | + dim_id * kNRows * params.out_d_stride + tokens_processed; |
| 356 | __syncthreads(); |
| 357 | #pragma unroll |
| 358 | for (int r = 0; r < kNRows; ++r) { |
| 359 | if constexpr (!kDirectIO) { |
| 360 | if (r > 0) { __syncthreads(); } |
| 361 | } |
| 362 | store_output<Ktraits>(out + r * params.out_d_stride, out_vals[r], smem_store, chunk_tokens); |
| 363 | } |
| 364 | |
| 365 | if constexpr (kHasZ) { |
| 366 | input_t *z = reinterpret_cast<input_t *>(params.z_ptr) + sequence_start_index * params.z_batch_stride |
| 367 | + dim_id * kNRows * params.z_d_stride + tokens_processed; |
| 368 | input_t *out_z = reinterpret_cast<input_t *>(params.out_z_ptr) + sequence_start_index * params.out_z_batch_stride |
| 369 | + dim_id * kNRows * params.out_z_d_stride + tokens_processed; |
| 370 | #pragma unroll |
| 371 | for (int r = 0; r < kNRows; ++r) { |
| 372 | input_t z_vals[kNItems]; |
| 373 | __syncthreads(); |
| 374 | load_input<Ktraits>(z + r * params.z_d_stride, z_vals, smem_load, chunk_tokens); |
| 375 | #pragma unroll |
| 376 | for (int i = 0; i < kNItems; ++i) { |
| 377 | float z_val = z_vals[i]; |
| 378 | out_vals[r][i] *= z_val / (1 + expf(-z_val)); |
| 379 | } |
| 380 | __syncthreads(); |
| 381 | store_output<Ktraits>(out_z + r * params.out_z_d_stride, out_vals[r], smem_store, chunk_tokens); |
| 382 | } |
| 383 | } |
| 384 | |
| 385 | Bvar += chunk_tokens; |
| 386 | Cvar += chunk_tokens; |
| 387 | |
| 388 | tokens_processed += chunk_tokens; |
| 389 | current_position += chunk_tokens; |
| 390 | } |
| 391 | } |
| 392 | |
| 393 | template<int kNThreads, int kNItems, typename input_t, typename weight_t, typename state_t> |
| 394 | void selective_scan_fwd_launch(SSMParamsBase ¶ms, cudaStream_t stream) { |
| 395 | // Only kNRows == 1 is tested for now, which ofc doesn't differ from previously when we had each block |
| 396 | // processing 1 row. |
| 397 | constexpr int kNRows = 1; |
| 398 | // kIsVariableB, kIsVariableC and kHasZ are all set to True to reduce binary size |
| 399 | constexpr bool kIsVariableB = true; |
| 400 | constexpr bool kIsVariableC = true; |
| 401 | BOOL_SWITCH(params.seqlen % (kNThreads * kNItems) == 0, kIsEvenLen, [&] { |
| 402 | BOOL_SWITCH(params.z_ptr != nullptr , kHasZ, [&] { |
| 403 | BOOL_SWITCH(params.query_start_loc_ptr != nullptr , kVarlen, [&] { |
| 404 | using Ktraits = Selective_Scan_fwd_kernel_traits<kNThreads, kNItems, kNRows, kIsEvenLen, kIsVariableB, kIsVariableC, kHasZ, kVarlen, input_t, weight_t, state_t>; |
| 405 | constexpr int kSmemSize = Ktraits::kSmemSize + kNRows * MAX_DSTATE * sizeof(typename Ktraits::scan_t); |
| 406 | dim3 grid(params.batch, params.dim / kNRows); |
| 407 | auto kernel = &selective_scan_fwd_kernel<Ktraits>; |
| 408 | if (kSmemSize >= 48 * 1024) { |
| 409 | #ifdef USE_ROCM |
| 410 | STD_CUDA_CHECK(hipFuncSetAttribute( |
| 411 | reinterpret_cast<const void*>(kernel), hipFuncAttributeMaxDynamicSharedMemorySize, kSmemSize)); |
| 412 | #else |
| 413 | STD_CUDA_CHECK(cudaFuncSetAttribute( |
| 414 | kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, kSmemSize)); |
| 415 | #endif |
| 416 | } |
| 417 | kernel<<<grid, Ktraits::kNThreads, kSmemSize, stream>>>(params); |
| 418 | STD_CUDA_KERNEL_LAUNCH_CHECK(); |
| 419 | }); |
| 420 | }); |
| 421 | }); |
| 422 | } |
| 423 | |
| 424 | template<typename input_t, typename weight_t, typename state_t> |
| 425 | void selective_scan_fwd_cuda(SSMParamsBase ¶ms, cudaStream_t stream) { |
| 426 | |
| 427 | #ifndef USE_ROCM |
| 428 | if (params.cache_enabled && params.block_size == 1024) { |
| 429 | selective_scan_fwd_launch<64, 16, input_t, weight_t, state_t>(params, stream); |
| 430 | } else if (params.seqlen <= 128) { |
| 431 | selective_scan_fwd_launch<32, 4, input_t, weight_t, state_t>(params, stream); |
| 432 | } else if (params.seqlen <= 256) { |
| 433 | selective_scan_fwd_launch<32, 8, input_t, weight_t, state_t>(params, stream); |
| 434 | } else if (params.seqlen <= 512) { |
| 435 | selective_scan_fwd_launch<32, 16, input_t, weight_t, state_t>(params, stream); |
| 436 | } else if (params.seqlen <= 1024) { |
| 437 | selective_scan_fwd_launch<64, 16, input_t, weight_t, state_t>(params, stream); |
| 438 | } else { |
| 439 | selective_scan_fwd_launch<128, 16, input_t, weight_t, state_t>(params, stream); |
| 440 | } |
| 441 | #else |
| 442 | if (params.cache_enabled && params.block_size == 1024) { |
| 443 | selective_scan_fwd_launch<64, 16, input_t, weight_t, state_t>(params, stream); |
| 444 | } else if (params.seqlen <= 256) { |
| 445 | selective_scan_fwd_launch<64, 4, input_t, weight_t, state_t>(params, stream); |
| 446 | } else if (params.seqlen <= 512) { |
| 447 | selective_scan_fwd_launch<64, 8, input_t, weight_t, state_t>(params, stream); |
| 448 | } else if (params.seqlen <= 1024) { |
| 449 | selective_scan_fwd_launch<64, 16, input_t, weight_t, state_t>(params, stream); |
| 450 | } else { |
| 451 | selective_scan_fwd_launch<128, 16, input_t, weight_t, state_t>(params, stream); |
| 452 | } |
| 453 | #endif |
| 454 | } |
| 455 | |
| 456 | template void selective_scan_fwd_cuda<torch::headeronly::BFloat16, float, torch::headeronly::BFloat16>(SSMParamsBase ¶ms, cudaStream_t stream); |
| 457 | template void selective_scan_fwd_cuda<torch::headeronly::BFloat16, float, float>(SSMParamsBase ¶ms, cudaStream_t stream); |
| 458 | template void selective_scan_fwd_cuda<torch::headeronly::Half, float, torch::headeronly::Half>(SSMParamsBase ¶ms, cudaStream_t stream); |
| 459 | template void selective_scan_fwd_cuda<torch::headeronly::Half, float, float>(SSMParamsBase ¶ms, cudaStream_t stream); |
| 460 | template void selective_scan_fwd_cuda<float, float, float>(SSMParamsBase ¶ms, cudaStream_t stream); |
| 461 | |
| 462 | #define CHECK_SHAPE(x, ...) STD_TORCH_CHECK(x.sizes().equals(torch::headeronly::IntHeaderOnlyArrayRef({__VA_ARGS__})), #x " must have shape (" #__VA_ARGS__ ")") |
| 463 | |
| 464 | #define DISPATCH_WTYPE_ITYPE_FLOAT_AND_HALF_AND_BF16(ITYPE, STYPE, NAME, ...) \ |
| 465 | if (ITYPE == torch::headeronly::ScalarType::Half) { \ |
| 466 | using input_t = torch::headeronly::Half; \ |
| 467 | using weight_t = float; \ |
| 468 | if (STYPE == torch::headeronly::ScalarType::Half) { \ |
| 469 | using state_t = torch::headeronly::Half; \ |
| 470 | __VA_ARGS__(); \ |
| 471 | } else if (STYPE == torch::headeronly::ScalarType::Float) { \ |
| 472 | using state_t = float; \ |
| 473 | __VA_ARGS__(); \ |
| 474 | } else { \ |
| 475 | STD_TORCH_CHECK(false, #NAME " not implemented for state type '", STYPE, "'"); \ |
| 476 | } \ |
| 477 | } else if (ITYPE == torch::headeronly::ScalarType::BFloat16) { \ |
| 478 | using input_t = torch::headeronly::BFloat16; \ |
| 479 | using weight_t = float; \ |
| 480 | if (STYPE == torch::headeronly::ScalarType::BFloat16) { \ |
| 481 | using state_t = torch::headeronly::BFloat16; \ |
| 482 | __VA_ARGS__(); \ |
| 483 | } else if (STYPE == torch::headeronly::ScalarType::Float) { \ |
| 484 | using state_t = float; \ |
| 485 | __VA_ARGS__(); \ |
| 486 | } else { \ |
| 487 | STD_TORCH_CHECK(false, #NAME " not implemented for state type '", STYPE, "'"); \ |
| 488 | } \ |
| 489 | } else if (ITYPE == torch::headeronly::ScalarType::Float) { \ |
| 490 | using input_t = float; \ |
| 491 | using weight_t = float; \ |
| 492 | using state_t = float; \ |
| 493 | __VA_ARGS__(); \ |
| 494 | } else { \ |
| 495 | STD_TORCH_CHECK(false, #NAME " not implemented for input type '", ITYPE, "'"); \ |
| 496 | } |
| 497 | |
| 498 | |
| 499 | template<typename input_t, typename weight_t, typename state_t> |
| 500 | void selective_scan_fwd_cuda(SSMParamsBase ¶ms, cudaStream_t stream); |
| 501 | |
| 502 | void set_ssm_params_fwd(SSMParamsBase ¶ms, |
| 503 | // sizes |
| 504 | const size_t batch, |
| 505 | const size_t dim, |
| 506 | const size_t seqlen, |
| 507 | const size_t dstate, |
| 508 | const size_t n_groups, |
| 509 | const bool is_variable_B, |
| 510 | const bool is_variable_C, |
| 511 | // device pointers |
| 512 | const torch::stable::Tensor u, |
| 513 | const torch::stable::Tensor delta, |
| 514 | const torch::stable::Tensor A, |
| 515 | const torch::stable::Tensor B, |
| 516 | const torch::stable::Tensor C, |
| 517 | const torch::stable::Tensor out, |
| 518 | const torch::stable::Tensor z, |
| 519 | const torch::stable::Tensor out_z, |
| 520 | const std::optional<torch::stable::Tensor>& D, |
| 521 | const std::optional<torch::stable::Tensor>& delta_bias, |
| 522 | const torch::stable::Tensor ssm_states, |
| 523 | bool has_z, |
| 524 | bool delta_softplus, |
| 525 | const std::optional<torch::stable::Tensor>& query_start_loc, |
| 526 | const std::optional<torch::stable::Tensor>& cache_indices, |
| 527 | const std::optional<torch::stable::Tensor>& has_initial_state, |
| 528 | bool varlen, |
| 529 | int64_t null_block_id, |
| 530 | int64_t block_size, |
| 531 | const std::optional<torch::stable::Tensor> &block_idx_first_scheduled_token, |
| 532 | const std::optional<torch::stable::Tensor> &block_idx_last_scheduled_token, |
| 533 | const std::optional<torch::stable::Tensor> &initial_state_idx, |
| 534 | const std::optional<torch::stable::Tensor> &cu_chunk_seqlen, |
| 535 | const std::optional<torch::stable::Tensor> &last_chunk_indices) { |
| 536 | |
| 537 | // Reset the parameters |
| 538 | memset(¶ms, 0, sizeof(params)); |
| 539 | |
| 540 | params.batch = batch; |
| 541 | params.dim = dim; |
| 542 | params.seqlen = seqlen; |
| 543 | params.dstate = dstate; |
| 544 | params.n_groups = n_groups; |
| 545 | params.dim_ngroups_ratio = dim / n_groups; |
| 546 | params.null_block_id = null_block_id; |
| 547 | |
| 548 | params.delta_softplus = delta_softplus; |
| 549 | |
| 550 | params.is_variable_B = is_variable_B; |
| 551 | params.is_variable_C = is_variable_C; |
| 552 | |
| 553 | // Set the pointers and strides. |
| 554 | params.u_ptr = u.data_ptr(); |
| 555 | params.delta_ptr = delta.data_ptr(); |
| 556 | params.A_ptr = A.data_ptr(); |
| 557 | params.B_ptr = B.data_ptr(); |
| 558 | params.C_ptr = C.data_ptr(); |
| 559 | params.D_ptr = D.has_value() ? D.value().data_ptr() : nullptr; |
| 560 | params.delta_bias_ptr = delta_bias.has_value() ? delta_bias.value().data_ptr() : nullptr; |
| 561 | params.out_ptr = out.data_ptr(); |
| 562 | params.ssm_states_ptr = ssm_states.data_ptr(); |
| 563 | params.z_ptr = has_z ? z.data_ptr() : nullptr; |
| 564 | params.out_z_ptr = has_z ? out_z.data_ptr() : nullptr; |
| 565 | params.query_start_loc_ptr = query_start_loc.has_value() ? query_start_loc.value().data_ptr() : nullptr; |
| 566 | params.cache_indices_ptr = cache_indices.has_value() ? cache_indices.value().data_ptr() : nullptr; |
| 567 | params.has_initial_state_ptr = has_initial_state.has_value() ? has_initial_state.value().data_ptr() : nullptr; |
| 568 | |
| 569 | // Set cache parameters - cache is enabled if we have direct cache writing params |
| 570 | params.cache_enabled = block_idx_first_scheduled_token.has_value(); |
| 571 | params.block_size = static_cast<int>(block_size); |
| 572 | |
| 573 | // Set direct cache writing pointers |
| 574 | params.block_idx_first_scheduled_token_ptr = block_idx_first_scheduled_token.has_value() ? block_idx_first_scheduled_token.value().data_ptr() : nullptr; |
| 575 | params.block_idx_last_scheduled_token_ptr = block_idx_last_scheduled_token.has_value() ? block_idx_last_scheduled_token.value().data_ptr() : nullptr; |
| 576 | params.initial_state_idx_ptr = initial_state_idx.has_value() ? initial_state_idx.value().data_ptr() : nullptr; |
| 577 | params.cu_chunk_seqlen_ptr = cu_chunk_seqlen.has_value() ? cu_chunk_seqlen.value().data_ptr() : nullptr; |
| 578 | params.last_chunk_indices_ptr = last_chunk_indices.has_value() ? last_chunk_indices.value().data_ptr() : nullptr; |
| 579 | |
| 580 | // All stride are in elements, not bytes. |
| 581 | params.A_d_stride = A.stride(0); |
| 582 | params.A_dstate_stride = A.stride(1); |
| 583 | |
| 584 | if (varlen){ |
| 585 | params.B_batch_stride = B.stride(2); |
| 586 | params.B_group_stride = B.stride(0); |
| 587 | params.B_dstate_stride = B.stride(1); |
| 588 | params.C_batch_stride = C.stride(2); |
| 589 | params.C_group_stride = C.stride(0); |
| 590 | params.C_dstate_stride = C.stride(1); |
| 591 | |
| 592 | params.u_batch_stride = u.stride(1); |
| 593 | params.u_d_stride = u.stride(0); |
| 594 | params.delta_batch_stride = delta.stride(1); |
| 595 | params.delta_d_stride = delta.stride(0); |
| 596 | if (has_z) { |
| 597 | params.z_batch_stride = z.stride(1); |
| 598 | params.z_d_stride = z.stride(0); |
| 599 | params.out_z_batch_stride = out_z.stride(1); |
| 600 | params.out_z_d_stride = out_z.stride(0); |
| 601 | } |
| 602 | params.out_batch_stride = out.stride(1); |
| 603 | params.out_d_stride = out.stride(0); |
| 604 | |
| 605 | params.ssm_states_batch_stride = ssm_states.stride(0); |
| 606 | params.ssm_states_dim_stride = ssm_states.stride(1); |
| 607 | params.ssm_states_dstate_stride = ssm_states.stride(2); |
| 608 | |
| 609 | params.cache_indices_stride = cache_indices.has_value() ? cache_indices.value().stride(0) : 0; |
| 610 | |
| 611 | } |
| 612 | else{ |
| 613 | if (!is_variable_B) { |
| 614 | params.B_d_stride = B.stride(0); |
| 615 | } else { |
| 616 | params.B_batch_stride = B.stride(0); |
| 617 | params.B_group_stride = B.stride(1); |
| 618 | } |
| 619 | params.B_dstate_stride = !is_variable_B ? B.stride(1) : B.stride(2); |
| 620 | if (!is_variable_C) { |
| 621 | params.C_d_stride = C.stride(0); |
| 622 | } else { |
| 623 | params.C_batch_stride = C.stride(0); |
| 624 | params.C_group_stride = C.stride(1); |
| 625 | } |
| 626 | params.C_dstate_stride = !is_variable_C ? C.stride(1) : C.stride(2); |
| 627 | params.u_batch_stride = u.stride(0); |
| 628 | params.u_d_stride = u.stride(1); |
| 629 | params.delta_batch_stride = delta.stride(0); |
| 630 | params.delta_d_stride = delta.stride(1); |
| 631 | if (has_z) { |
| 632 | params.z_batch_stride = z.stride(0); |
| 633 | params.z_d_stride = z.stride(1); |
| 634 | params.out_z_batch_stride = out_z.stride(0); |
| 635 | params.out_z_d_stride = out_z.stride(1); |
| 636 | } |
| 637 | params.out_batch_stride = out.stride(0); |
| 638 | params.out_d_stride = out.stride(1); |
| 639 | |
| 640 | params.ssm_states_batch_stride = ssm_states.stride(0); |
| 641 | params.ssm_states_dim_stride = ssm_states.stride(1); |
| 642 | params.ssm_states_dstate_stride = ssm_states.stride(2); |
| 643 | |
| 644 | params.cache_indices_stride = cache_indices.has_value() ? cache_indices.value().stride(0) : 0; |
| 645 | } |
| 646 | } |
| 647 | |
| 648 | void selective_scan_fwd(const torch::stable::Tensor &u, const torch::stable::Tensor &delta, |
| 649 | const torch::stable::Tensor &A, const torch::stable::Tensor &B, const torch::stable::Tensor &C, |
| 650 | const std::optional<torch::stable::Tensor> &D_, |
| 651 | const std::optional<torch::stable::Tensor> &z_, |
| 652 | const std::optional<torch::stable::Tensor> &delta_bias_, |
| 653 | bool delta_softplus, |
| 654 | const std::optional<torch::stable::Tensor> &query_start_loc, |
| 655 | const std::optional<torch::stable::Tensor> &cache_indices, |
| 656 | const std::optional<torch::stable::Tensor> &has_initial_state, |
| 657 | const torch::stable::Tensor &ssm_states, |
| 658 | // used to identify padding entries if cache_indices provided |
| 659 | // in case of padding, the kernel will return early |
| 660 | int64_t null_block_id, |
| 661 | int64_t block_size, |
| 662 | const std::optional<torch::stable::Tensor> &block_idx_first_scheduled_token, |
| 663 | const std::optional<torch::stable::Tensor> &block_idx_last_scheduled_token, |
| 664 | const std::optional<torch::stable::Tensor> &initial_state_idx, |
| 665 | const std::optional<torch::stable::Tensor> &cu_chunk_seqlen, |
| 666 | const std::optional<torch::stable::Tensor> &last_chunk_indices) { |
| 667 | auto input_type = u.scalar_type(); |
| 668 | auto weight_type = A.scalar_type(); |
| 669 | STD_TORCH_CHECK(input_type == torch::headeronly::ScalarType::Float || input_type == torch::headeronly::ScalarType::Half || input_type == torch::headeronly::ScalarType::BFloat16); |
| 670 | STD_TORCH_CHECK(weight_type == torch::headeronly::ScalarType::Float); |
| 671 | |
| 672 | const bool is_variable_B = B.dim() >= 3; |
| 673 | const bool is_variable_C = C.dim() >= 3; |
| 674 | |
| 675 | STD_TORCH_CHECK(delta.scalar_type() == input_type); |
| 676 | STD_TORCH_CHECK(B.scalar_type() == (!is_variable_B ? weight_type : input_type)); |
| 677 | STD_TORCH_CHECK(C.scalar_type() == (!is_variable_C ? weight_type : input_type)); |
| 678 | |
| 679 | STD_TORCH_CHECK(u.is_cuda()); |
| 680 | STD_TORCH_CHECK(delta.is_cuda()); |
| 681 | STD_TORCH_CHECK(A.is_cuda()); |
| 682 | STD_TORCH_CHECK(B.is_cuda()); |
| 683 | STD_TORCH_CHECK(C.is_cuda()); |
| 684 | |
| 685 | STD_TORCH_CHECK(u.stride(-1) == 1 || u.size(-1) == 1); |
| 686 | STD_TORCH_CHECK(delta.stride(-1) == 1 || delta.size(-1) == 1); |
| 687 | |
| 688 | const auto sizes = u.sizes(); |
| 689 | const bool varlen = query_start_loc.has_value(); |
| 690 | const int batch_size = varlen ? query_start_loc.value().sizes()[0] - 1 : sizes[0]; |
| 691 | const int dim = varlen ? sizes[0] : sizes[1]; |
| 692 | const int seqlen = varlen ? sizes[1] : sizes[2]; |
| 693 | const int dstate = A.size(1); |
| 694 | const int n_groups = varlen ? B.size(0) : B.size(1); |
| 695 | |
| 696 | STD_TORCH_CHECK(dstate <= 256, "selective_scan only supports state dimension <= 256"); |
| 697 | |
| 698 | if (varlen) { |
| 699 | CHECK_SHAPE(u, dim, seqlen); |
| 700 | CHECK_SHAPE(delta, dim, seqlen); |
| 701 | } else { |
| 702 | CHECK_SHAPE(u, batch_size, dim, seqlen); |
| 703 | CHECK_SHAPE(delta, batch_size, dim, seqlen); |
| 704 | } |
| 705 | CHECK_SHAPE(A, dim, dstate); |
| 706 | STD_TORCH_CHECK(is_variable_B, "is_variable_B = False is disabled in favor of reduced binary size"); |
| 707 | if (varlen) { |
| 708 | CHECK_SHAPE(B, n_groups, dstate, seqlen); |
| 709 | } else { |
| 710 | CHECK_SHAPE(B, batch_size, n_groups, dstate, seqlen); |
| 711 | } |
| 712 | STD_TORCH_CHECK(B.stride(-1) == 1 || B.size(-1) == 1); |
| 713 | |
| 714 | STD_TORCH_CHECK(is_variable_C, "is_variable_C = False is disabled in favor of reduced binary size"); |
| 715 | if (varlen) { |
| 716 | CHECK_SHAPE(C, n_groups, dstate, seqlen); |
| 717 | } else { |
| 718 | CHECK_SHAPE(C, batch_size, n_groups, dstate, seqlen); |
| 719 | } |
| 720 | STD_TORCH_CHECK(C.stride(-1) == 1 || C.size(-1) == 1); |
| 721 | |
| 722 | if (D_.has_value()) { |
| 723 | auto D = D_.value(); |
| 724 | STD_TORCH_CHECK(D.scalar_type() == torch::headeronly::ScalarType::Float); |
| 725 | STD_TORCH_CHECK(D.is_cuda()); |
| 726 | STD_TORCH_CHECK(D.stride(-1) == 1 || D.size(-1) == 1); |
| 727 | CHECK_SHAPE(D, dim); |
| 728 | } |
| 729 | |
| 730 | if (delta_bias_.has_value()) { |
| 731 | auto delta_bias = delta_bias_.value(); |
| 732 | STD_TORCH_CHECK(delta_bias.scalar_type() == torch::headeronly::ScalarType::Float); |
| 733 | STD_TORCH_CHECK(delta_bias.is_cuda()); |
| 734 | STD_TORCH_CHECK(delta_bias.stride(-1) == 1 || delta_bias.size(-1) == 1); |
| 735 | CHECK_SHAPE(delta_bias, dim); |
| 736 | } |
| 737 | |
| 738 | |
| 739 | if (has_initial_state.has_value()) { |
| 740 | auto has_initial_state_ = has_initial_state.value(); |
| 741 | STD_TORCH_CHECK(has_initial_state_.scalar_type() == torch::headeronly::ScalarType::Bool); |
| 742 | STD_TORCH_CHECK(has_initial_state_.is_cuda()); |
| 743 | CHECK_SHAPE(has_initial_state_, batch_size); |
| 744 | } |
| 745 | |
| 746 | |
| 747 | if (query_start_loc.has_value()) { |
| 748 | auto query_start_loc_ = query_start_loc.value(); |
| 749 | STD_TORCH_CHECK(query_start_loc_.scalar_type() == torch::headeronly::ScalarType::Int); |
| 750 | STD_TORCH_CHECK(query_start_loc_.is_cuda()); |
| 751 | } |
| 752 | |
| 753 | |
| 754 | if (cache_indices.has_value()) { |
| 755 | auto cache_indices_ = cache_indices.value(); |
| 756 | STD_TORCH_CHECK(cache_indices_.scalar_type() == torch::headeronly::ScalarType::Int); |
| 757 | STD_TORCH_CHECK(cache_indices_.is_cuda()); |
| 758 | |
| 759 | // cache_indices can be either 1D (batch_size,) for non-APC mode |
| 760 | // or 2D (batch_size, max_positions) for APC mode |
| 761 | const bool is_apc_mode = block_idx_first_scheduled_token.has_value(); |
| 762 | if (is_apc_mode) { |
| 763 | STD_TORCH_CHECK(cache_indices_.dim() == 2, "cache_indices must be 2D for APC mode"); |
| 764 | STD_TORCH_CHECK(cache_indices_.size(0) == batch_size, "cache_indices first dimension must match batch_size"); |
| 765 | } else { |
| 766 | CHECK_SHAPE(cache_indices_, batch_size); |
| 767 | } |
| 768 | } |
| 769 | |
| 770 | |
| 771 | torch::stable::Tensor z, out_z; |
| 772 | const bool has_z = z_.has_value(); |
| 773 | if (has_z) { |
| 774 | z = z_.value(); |
| 775 | STD_TORCH_CHECK(z.scalar_type() == input_type); |
| 776 | STD_TORCH_CHECK(z.is_cuda()); |
| 777 | STD_TORCH_CHECK(z.stride(-1) == 1 || z.size(-1) == 1); |
| 778 | if (varlen){ |
| 779 | CHECK_SHAPE(z, dim, seqlen); |
| 780 | } else { |
| 781 | CHECK_SHAPE(z, batch_size, dim, seqlen); |
| 782 | } |
| 783 | |
| 784 | out_z = z; |
| 785 | } |
| 786 | |
| 787 | // Right now u has BHL layout and delta has HBL layout, and we want out to have HBL layout |
| 788 | torch::stable::Tensor out = delta; |
| 789 | // ssm_states can now be either the same as input_type or float32 |
| 790 | auto state_type = ssm_states.scalar_type(); |
| 791 | STD_TORCH_CHECK(state_type == input_type || state_type == torch::headeronly::ScalarType::Float); |
| 792 | STD_TORCH_CHECK(ssm_states.is_cuda()); |
| 793 | STD_TORCH_CHECK(ssm_states.stride(-1) == 1); |
| 794 | |
| 795 | SSMParamsBase params; |
| 796 | set_ssm_params_fwd(params, batch_size, dim, seqlen, dstate, n_groups, is_variable_B, is_variable_C, |
| 797 | u, delta, A, B, C, out, z, out_z, |
| 798 | D_, |
| 799 | delta_bias_, |
| 800 | ssm_states, |
| 801 | has_z, |
| 802 | delta_softplus, |
| 803 | query_start_loc, |
| 804 | cache_indices, |
| 805 | has_initial_state, |
| 806 | varlen, |
| 807 | null_block_id, |
| 808 | block_size, |
| 809 | block_idx_first_scheduled_token, |
| 810 | block_idx_last_scheduled_token, |
| 811 | initial_state_idx, |
| 812 | cu_chunk_seqlen, |
| 813 | last_chunk_indices |
| 814 | ); |
| 815 | |
| 816 | |
| 817 | const torch::stable::accelerator::DeviceGuard device_guard(u.get_device_index()); |
| 818 | auto stream = get_current_cuda_stream(); |
| 819 | DISPATCH_WTYPE_ITYPE_FLOAT_AND_HALF_AND_BF16(u.scalar_type(), ssm_states.scalar_type(), "selective_scan_fwd", [&] { |
| 820 | selective_scan_fwd_cuda<input_t, weight_t, state_t>(params, stream); |
| 821 | }); |
| 822 | } |