| 1 | /* |
| 2 | * SPDX-License-Identifier: Apache-2.0 |
| 3 | * SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 4 | */ |
| 5 | |
| 6 | #include <cstdint> |
| 7 | #include <string> |
| 8 | #include <cuda_bf16.h> |
| 9 | #include <cuda_fp16.h> |
| 10 | #include <cuda_runtime.h> |
| 11 | |
| 12 | #include "../torch_utils.h" |
| 13 | #include "../../cuda_compat.h" |
| 14 | |
| 15 | namespace { |
| 16 | |
| 17 | template <typename StateT> |
| 18 | __device__ __forceinline__ void cp_async_16b(StateT* smem_ptr, |
| 19 | const StateT* gmem_ptr) { |
| 20 | const uint32_t smem_addr = |
| 21 | static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr)); |
| 22 | asm volatile("cp.async.cg.shared.global [%0], [%1], 16;\n" |
| 23 | : |
| 24 | : "r"(smem_addr), "l"(gmem_ptr)); |
| 25 | } |
| 26 | |
| 27 | __device__ __forceinline__ void cp_async_commit() { |
| 28 | asm volatile("cp.async.commit_group;\n" ::); |
| 29 | } |
| 30 | |
| 31 | __device__ __forceinline__ void cp_async_wait_all() { |
| 32 | asm volatile("cp.async.wait_all;\n" ::: "memory"); |
| 33 | } |
| 34 | |
| 35 | template <typename StateT, int ChunkV, int DimK, int Stages> |
| 36 | __device__ __forceinline__ void copy_state_chunk(StateT* shared_state, |
| 37 | const StateT* state, int chunk, |
| 38 | int thread, int threads) { |
| 39 | constexpr int kElementsPerCopy = 16 / sizeof(StateT); |
| 40 | constexpr int kCopiesPerChunk = ChunkV * DimK / kElementsPerCopy; |
| 41 | const int stage = chunk % Stages; |
| 42 | for (int copy = thread; copy < kCopiesPerChunk; copy += threads) { |
| 43 | const int element = copy * kElementsPerCopy; |
| 44 | cp_async_16b(shared_state + stage * ChunkV * DimK + element, |
| 45 | state + chunk * ChunkV * DimK + element); |
| 46 | } |
| 47 | cp_async_commit(); |
| 48 | } |
| 49 | |
| 50 | template <typename StateT> |
| 51 | __device__ __forceinline__ float4 load_state4(const StateT* state); |
| 52 | |
| 53 | template <> |
| 54 | __device__ __forceinline__ float4 load_state4<float>(const float* state) { |
| 55 | return *reinterpret_cast<const float4*>(state); |
| 56 | } |
| 57 | |
| 58 | template <> |
| 59 | __device__ __forceinline__ float4 |
| 60 | load_state4<__nv_bfloat16>(const __nv_bfloat16* state) { |
| 61 | const __nv_bfloat162 lo = *reinterpret_cast<const __nv_bfloat162*>(state); |
| 62 | const __nv_bfloat162 hi = *reinterpret_cast<const __nv_bfloat162*>(state + 2); |
| 63 | return make_float4(__bfloat162float(lo.x), __bfloat162float(lo.y), |
| 64 | __bfloat162float(hi.x), __bfloat162float(hi.y)); |
| 65 | } |
| 66 | |
| 67 | template <typename StateT> |
| 68 | __device__ __forceinline__ void store_state4(StateT* state, float4 value); |
| 69 | |
| 70 | template <> |
| 71 | __device__ __forceinline__ void store_state4<float>(float* state, |
| 72 | float4 value) { |
| 73 | *reinterpret_cast<float4*>(state) = value; |
| 74 | } |
| 75 | |
| 76 | template <> |
| 77 | __device__ __forceinline__ void store_state4<__nv_bfloat16>( |
| 78 | __nv_bfloat16* state, float4 value) { |
| 79 | *reinterpret_cast<__nv_bfloat162*>(state) = |
| 80 | __floats2bfloat162_rn(value.x, value.y); |
| 81 | *reinterpret_cast<__nv_bfloat162*>(state + 2) = |
| 82 | __floats2bfloat162_rn(value.z, value.w); |
| 83 | } |
| 84 | |
| 85 | constexpr int kDimK = 128; |
| 86 | constexpr int kDimV = 128; |
| 87 | constexpr int kThreads = 256; |
| 88 | constexpr int kWarps = kThreads / 32; |
| 89 | constexpr int kChunkV = 32; |
| 90 | constexpr int kNumChunks = kDimV / kChunkV; |
| 91 | constexpr int kRowsPerWarp = kChunkV / kWarps; |
| 92 | constexpr int kMaxMtpTokens = 8; |
| 93 | constexpr int kDtBiasFloat32 = 0; |
| 94 | constexpr int kDtBiasBFloat16 = 1; |
| 95 | constexpr int kDtBiasFloat16 = 2; |
| 96 | |
| 97 | struct GdnDecodeStrides { |
| 98 | int64_t mixed_row; |
| 99 | int64_t a_row; |
| 100 | int64_t b_row; |
| 101 | int64_t gate_row; |
| 102 | int64_t state_slot; |
| 103 | }; |
| 104 | |
| 105 | __device__ __forceinline__ float sigmoid_fast(float x) { |
| 106 | return 1.0f / (1.0f + __expf(-x)); |
| 107 | } |
| 108 | |
| 109 | __device__ __forceinline__ float silu_fast(float x) { |
| 110 | return x * sigmoid_fast(x); |
| 111 | } |
| 112 | |
| 113 | __device__ __forceinline__ float softplus_fast(float x) { |
| 114 | return x > 20.0f ? x : log1pf(__expf(x)); |
| 115 | } |
| 116 | |
| 117 | __device__ __forceinline__ float load_dt_bias(const void* dt_bias, int head, |
| 118 | int dt_bias_type) { |
| 119 | if (dt_bias_type == kDtBiasBFloat16) { |
| 120 | return __bfloat162float(static_cast<const __nv_bfloat16*>(dt_bias)[head]); |
| 121 | } |
| 122 | if (dt_bias_type == kDtBiasFloat16) { |
| 123 | return __half2float(static_cast<const __half*>(dt_bias)[head]); |
| 124 | } |
| 125 | return static_cast<const float*>(dt_bias)[head]; |
| 126 | } |
| 127 | |
| 128 | __device__ __forceinline__ float warp_reduce_sum(float value) { |
| 129 | #pragma unroll |
| 130 | for (int offset = 16; offset > 0; offset >>= 1) { |
| 131 | value += __shfl_xor_sync(0xffffffffu, value, offset); |
| 132 | } |
| 133 | return value; |
| 134 | } |
| 135 | |
| 136 | struct Sum2 { |
| 137 | float x; |
| 138 | float y; |
| 139 | }; |
| 140 | |
| 141 | __device__ __forceinline__ Sum2 warp_reduce_sum_pair(float x, float y) { |
| 142 | #pragma unroll |
| 143 | for (int offset = 16; offset > 0; offset >>= 1) { |
| 144 | x += __shfl_xor_sync(0xffffffffu, x, offset); |
| 145 | y += __shfl_xor_sync(0xffffffffu, y, offset); |
| 146 | } |
| 147 | return {x, y}; |
| 148 | } |
| 149 | |
| 150 | template <typename StateT, int ValueHeadsPerKeyHead, bool SigmoidGate> |
| 151 | __global__ __launch_bounds__(kThreads, 2) void gdn_decode_post_conv_mtp_kernel( |
| 152 | const __nv_bfloat16* __restrict__ mixed_qkv, |
| 153 | const __nv_bfloat16* __restrict__ a, const __nv_bfloat16* __restrict__ b, |
| 154 | const float* __restrict__ a_log, const void* __restrict__ dt_bias, |
| 155 | const int* __restrict__ state_indices, const int* __restrict__ cu_seqlens, |
| 156 | const int* __restrict__ num_accepted_tokens, StateT* __restrict__ state, |
| 157 | const __nv_bfloat16* __restrict__ output_gate, |
| 158 | const void* __restrict__ norm_weight, __nv_bfloat16* __restrict__ out, |
| 159 | int H, int HV, int state_indices_width, int dt_bias_type, |
| 160 | bool norm_weight_is_bf16, float scale, float norm_eps, |
| 161 | GdnDecodeStrides strides) { |
| 162 | const int request = blockIdx.x; |
| 163 | const int value_head = blockIdx.y; |
| 164 | const int tid = threadIdx.x; |
| 165 | const int lane = tid & 31; |
| 166 | const int warp = tid >> 5; |
| 167 | const int bos = cu_seqlens[request]; |
| 168 | const int eos = cu_seqlens[request + 1]; |
| 169 | const int num_tokens = eos - bos; |
| 170 | if (num_tokens <= 0) { |
| 171 | return; |
| 172 | } |
| 173 | |
| 174 | const int accepted = num_accepted_tokens[request]; |
| 175 | const int source_slot = |
| 176 | accepted > 0 && accepted <= state_indices_width |
| 177 | ? state_indices[request * state_indices_width + accepted - 1] |
| 178 | : 0; |
| 179 | if (source_slot <= 0 || num_tokens > kMaxMtpTokens) { |
| 180 | for (int linear = tid; linear < num_tokens * kDimV; linear += kThreads) { |
| 181 | const int token = bos + linear / kDimV; |
| 182 | const int value = linear % kDimV; |
| 183 | const int64_t out_offset = |
| 184 | (static_cast<int64_t>(token) * HV + value_head) * kDimV + value; |
| 185 | out[out_offset] = __float2bfloat16(0.0f); |
| 186 | } |
| 187 | return; |
| 188 | } |
| 189 | |
| 190 | const int key_head = value_head / ValueHeadsPerKeyHead; |
| 191 | __shared__ StateT shared_state[2][kChunkV][kDimK]; |
| 192 | __shared__ float shared_q[kMaxMtpTokens][kDimK]; |
| 193 | __shared__ float shared_k[kMaxMtpTokens][kDimK]; |
| 194 | __shared__ __nv_bfloat16 shared_v[kMaxMtpTokens][kDimV]; |
| 195 | __shared__ __nv_bfloat16 shared_out[kMaxMtpTokens][kDimV]; |
| 196 | __shared__ float shared_decay[kMaxMtpTokens]; |
| 197 | __shared__ float shared_beta[kMaxMtpTokens]; |
| 198 | |
| 199 | StateT* source_state = |
| 200 | state + static_cast<int64_t>(source_slot) * strides.state_slot + |
| 201 | value_head * kDimV * kDimK; |
| 202 | copy_state_chunk<StateT, kChunkV, kDimK, 2>(&shared_state[0][0][0], |
| 203 | source_state, 0, tid, kThreads); |
| 204 | |
| 205 | if (warp < num_tokens) { |
| 206 | const int t = warp; |
| 207 | const int token = bos + t; |
| 208 | const int64_t mixed_base = static_cast<int64_t>(token) * strides.mixed_row; |
| 209 | float q_values[4]; |
| 210 | float k_values[4]; |
| 211 | float q_square = 0.0f; |
| 212 | float k_square = 0.0f; |
| 213 | #pragma unroll |
| 214 | for (int i = 0; i < 4; ++i) { |
| 215 | const int dim = lane + i * 32; |
| 216 | q_values[i] = |
| 217 | __bfloat162float(mixed_qkv[mixed_base + key_head * kDimK + dim]); |
| 218 | k_values[i] = __bfloat162float( |
| 219 | mixed_qkv[mixed_base + H * kDimK + key_head * kDimK + dim]); |
| 220 | shared_v[t][dim] = |
| 221 | mixed_qkv[mixed_base + 2 * H * kDimK + value_head * kDimV + dim]; |
| 222 | q_square += q_values[i] * q_values[i]; |
| 223 | k_square += k_values[i] * k_values[i]; |
| 224 | } |
| 225 | const Sum2 qk_sums = warp_reduce_sum_pair(q_square, k_square); |
| 226 | const float q_scale = __shfl_sync( |
| 227 | 0xffffffffu, lane == 0 ? rsqrtf(qk_sums.x + 1.0e-6f) * scale : 0.0f, 0); |
| 228 | const float k_scale = __shfl_sync( |
| 229 | 0xffffffffu, lane == 0 ? rsqrtf(qk_sums.y + 1.0e-6f) : 0.0f, 0); |
| 230 | #pragma unroll |
| 231 | for (int i = 0; i < 4; ++i) { |
| 232 | const int dim = lane + i * 32; |
| 233 | shared_q[t][dim] = q_values[i] * q_scale; |
| 234 | shared_k[t][dim] = k_values[i] * k_scale; |
| 235 | } |
| 236 | if (lane == 0) { |
| 237 | const float a_value = __bfloat162float( |
| 238 | a[static_cast<int64_t>(token) * strides.a_row + value_head]); |
| 239 | const float b_value = __bfloat162float( |
| 240 | b[static_cast<int64_t>(token) * strides.b_row + value_head]); |
| 241 | const float g = -__expf(a_log[value_head]) * |
| 242 | softplus_fast(a_value + load_dt_bias(dt_bias, value_head, |
| 243 | dt_bias_type)); |
| 244 | shared_decay[t] = __expf(g); |
| 245 | shared_beta[t] = sigmoid_fast(b_value); |
| 246 | } |
| 247 | } |
| 248 | __syncthreads(); |
| 249 | |
| 250 | const int k_base = lane * 4; |
| 251 | int rows[kRowsPerWarp]; |
| 252 | #pragma unroll |
| 253 | for (int row = 0; row < kRowsPerWarp; ++row) { |
| 254 | rows[row] = warp + row * kWarps; |
| 255 | } |
| 256 | |
| 257 | #pragma unroll |
| 258 | for (int chunk = 0; chunk < kNumChunks; ++chunk) { |
| 259 | cp_async_wait_all(); |
| 260 | __syncthreads(); |
| 261 | if (chunk + 1 < kNumChunks) { |
| 262 | copy_state_chunk<StateT, kChunkV, kDimK, 2>( |
| 263 | &shared_state[0][0][0], source_state, chunk + 1, tid, kThreads); |
| 264 | } |
| 265 | |
| 266 | float h[kRowsPerWarp][4]; |
| 267 | #pragma unroll |
| 268 | for (int row = 0; row < kRowsPerWarp; ++row) { |
| 269 | const float4 state_value = |
| 270 | load_state4(&shared_state[chunk & 1][rows[row]][k_base]); |
| 271 | h[row][0] = state_value.x; |
| 272 | h[row][1] = state_value.y; |
| 273 | h[row][2] = state_value.z; |
| 274 | h[row][3] = state_value.w; |
| 275 | } |
| 276 | |
| 277 | for (int t = 0; t < num_tokens; ++t) { |
| 278 | const float4 q4 = *reinterpret_cast<const float4*>(&shared_q[t][k_base]); |
| 279 | const float4 k4 = *reinterpret_cast<const float4*>(&shared_k[t][k_base]); |
| 280 | const float q_values[4] = {q4.x, q4.y, q4.z, q4.w}; |
| 281 | const float k_values[4] = {k4.x, k4.y, k4.z, k4.w}; |
| 282 | |
| 283 | float dot_hk[kRowsPerWarp] = {0.0f, 0.0f, 0.0f, 0.0f}; |
| 284 | #pragma unroll |
| 285 | for (int row = 0; row < kRowsPerWarp; ++row) { |
| 286 | #pragma unroll |
| 287 | for (int i = 0; i < 4; ++i) { |
| 288 | h[row][i] *= shared_decay[t]; |
| 289 | dot_hk[row] += h[row][i] * k_values[i]; |
| 290 | } |
| 291 | } |
| 292 | const Sum2 dot_hk_01 = warp_reduce_sum_pair(dot_hk[0], dot_hk[1]); |
| 293 | const Sum2 dot_hk_23 = warp_reduce_sum_pair(dot_hk[2], dot_hk[3]); |
| 294 | const float reduced_hk[kRowsPerWarp] = {dot_hk_01.x, dot_hk_01.y, |
| 295 | dot_hk_23.x, dot_hk_23.y}; |
| 296 | |
| 297 | float dot_hq[kRowsPerWarp] = {0.0f, 0.0f, 0.0f, 0.0f}; |
| 298 | #pragma unroll |
| 299 | for (int row = 0; row < kRowsPerWarp; ++row) { |
| 300 | const int value = chunk * kChunkV + rows[row]; |
| 301 | const float delta = |
| 302 | (__bfloat162float(shared_v[t][value]) - reduced_hk[row]) * |
| 303 | shared_beta[t]; |
| 304 | #pragma unroll |
| 305 | for (int i = 0; i < 4; ++i) { |
| 306 | h[row][i] += k_values[i] * delta; |
| 307 | dot_hq[row] += h[row][i] * q_values[i]; |
| 308 | } |
| 309 | } |
| 310 | const Sum2 dot_hq_01 = warp_reduce_sum_pair(dot_hq[0], dot_hq[1]); |
| 311 | const Sum2 dot_hq_23 = warp_reduce_sum_pair(dot_hq[2], dot_hq[3]); |
| 312 | if (lane == 0) { |
| 313 | shared_out[t][chunk * kChunkV + rows[0]] = |
| 314 | __float2bfloat16(dot_hq_01.x); |
| 315 | shared_out[t][chunk * kChunkV + rows[1]] = |
| 316 | __float2bfloat16(dot_hq_01.y); |
| 317 | shared_out[t][chunk * kChunkV + rows[2]] = |
| 318 | __float2bfloat16(dot_hq_23.x); |
| 319 | shared_out[t][chunk * kChunkV + rows[3]] = |
| 320 | __float2bfloat16(dot_hq_23.y); |
| 321 | } |
| 322 | |
| 323 | const int destination_slot = |
| 324 | state_indices[request * state_indices_width + t]; |
| 325 | if (destination_slot > 0) { |
| 326 | StateT* destination_state = |
| 327 | state + |
| 328 | static_cast<int64_t>(destination_slot) * strides.state_slot + |
| 329 | value_head * kDimV * kDimK; |
| 330 | #pragma unroll |
| 331 | for (int row = 0; row < kRowsPerWarp; ++row) { |
| 332 | const int value = chunk * kChunkV + rows[row]; |
| 333 | store_state4(destination_state + value * kDimK + k_base, |
| 334 | make_float4(h[row][0], h[row][1], h[row][2], h[row][3])); |
| 335 | } |
| 336 | } |
| 337 | } |
| 338 | } |
| 339 | __syncthreads(); |
| 340 | |
| 341 | if (warp < num_tokens) { |
| 342 | const int t = warp; |
| 343 | float output_values[4]; |
| 344 | float sum_square = 0.0f; |
| 345 | #pragma unroll |
| 346 | for (int i = 0; i < 4; ++i) { |
| 347 | const int value = lane + i * 32; |
| 348 | output_values[i] = __bfloat162float(shared_out[t][value]); |
| 349 | sum_square += output_values[i] * output_values[i]; |
| 350 | } |
| 351 | sum_square = warp_reduce_sum(sum_square); |
| 352 | const float rstd = |
| 353 | rsqrtf(sum_square / static_cast<float>(kDimV) + norm_eps); |
| 354 | const int token = bos + t; |
| 355 | #pragma unroll |
| 356 | for (int i = 0; i < 4; ++i) { |
| 357 | const int value = lane + i * 32; |
| 358 | const float gate_input = __bfloat162float( |
| 359 | output_gate[static_cast<int64_t>(token) * strides.gate_row + |
| 360 | value_head * kDimV + value]); |
| 361 | const float gate = |
| 362 | SigmoidGate ? sigmoid_fast(gate_input) : silu_fast(gate_input); |
| 363 | const float weight = |
| 364 | norm_weight_is_bf16 |
| 365 | ? __bfloat162float( |
| 366 | static_cast<const __nv_bfloat16*>(norm_weight)[value]) |
| 367 | : static_cast<const float*>(norm_weight)[value]; |
| 368 | const int64_t out_offset = |
| 369 | (static_cast<int64_t>(token) * HV + value_head) * kDimV + value; |
| 370 | out[out_offset] = |
| 371 | __float2bfloat16(output_values[i] * rstd * weight * gate); |
| 372 | } |
| 373 | } |
| 374 | } |
| 375 | |
| 376 | template <typename StateT, int ValueHeadsPerKeyHead, bool SigmoidGate> |
| 377 | void launch_gdn_decode_post_conv_mtp( |
| 378 | torch::stable::Tensor const& mixed_qkv, torch::stable::Tensor const& a_log, |
| 379 | torch::stable::Tensor const& dt_bias, |
| 380 | torch::stable::Tensor const& state_indices, |
| 381 | torch::stable::Tensor const& cu_seqlens, |
| 382 | torch::stable::Tensor const& num_accepted_tokens, |
| 383 | torch::stable::Tensor& state, torch::stable::Tensor const& norm_weight, |
| 384 | torch::stable::Tensor& out, const __nv_bfloat16* a, const __nv_bfloat16* b, |
| 385 | const __nv_bfloat16* output_gate, int num_key_heads, int num_value_heads, |
| 386 | double scale, double norm_eps, GdnDecodeStrides strides) { |
| 387 | using torch::headeronly::ScalarType; |
| 388 | |
| 389 | const auto dt_bias_scalar_type = dt_bias.scalar_type(); |
| 390 | const int dt_bias_type = |
| 391 | dt_bias_scalar_type == ScalarType::Float |
| 392 | ? kDtBiasFloat32 |
| 393 | : (dt_bias_scalar_type == ScalarType::BFloat16 ? kDtBiasBFloat16 |
| 394 | : kDtBiasFloat16); |
| 395 | torch::stable::accelerator::DeviceGuard const device_guard( |
| 396 | mixed_qkv.get_device_index()); |
| 397 | const cudaStream_t stream = |
| 398 | get_current_cuda_stream(mixed_qkv.get_device_index()); |
| 399 | const int num_requests = static_cast<int>(state_indices.size(0)); |
| 400 | const dim3 grid(num_requests, num_value_heads); |
| 401 | gdn_decode_post_conv_mtp_kernel<StateT, ValueHeadsPerKeyHead, SigmoidGate> |
| 402 | <<<grid, kThreads, 0, stream>>>( |
| 403 | static_cast<const __nv_bfloat16*>(mixed_qkv.data_ptr()), a, b, |
| 404 | static_cast<const float*>(a_log.data_ptr()), dt_bias.data_ptr(), |
| 405 | static_cast<const int*>(state_indices.data_ptr()), |
| 406 | static_cast<const int*>(cu_seqlens.data_ptr()), |
| 407 | static_cast<const int*>(num_accepted_tokens.data_ptr()), |
| 408 | static_cast<StateT*>(state.data_ptr()), output_gate, |
| 409 | norm_weight.data_ptr(), static_cast<__nv_bfloat16*>(out.data_ptr()), |
| 410 | num_key_heads, num_value_heads, |
| 411 | static_cast<int>(state_indices.size(1)), dt_bias_type, |
| 412 | norm_weight.scalar_type() == ScalarType::BFloat16, |
| 413 | static_cast<float>(scale), static_cast<float>(norm_eps), strides); |
| 414 | const cudaError_t error = cudaGetLastError(); |
| 415 | STD_TORCH_CHECK(error == cudaSuccess, |
| 416 | "GDN decode MTP post-conv kernel launch failed: ", |
| 417 | cudaGetErrorString(error)); |
| 418 | } |
| 419 | |
| 420 | } // namespace |
| 421 | |
| 422 | void fused_gdn_decode_post_conv_mtp( |
| 423 | torch::stable::Tensor const& mixed_qkv, torch::stable::Tensor const& a, |
| 424 | torch::stable::Tensor const& b, torch::stable::Tensor const& a_log, |
| 425 | torch::stable::Tensor const& dt_bias, |
| 426 | torch::stable::Tensor const& state_indices, |
| 427 | torch::stable::Tensor const& cu_seqlens, |
| 428 | torch::stable::Tensor const& num_accepted_tokens, |
| 429 | torch::stable::Tensor& state, torch::stable::Tensor const& output_gate, |
| 430 | torch::stable::Tensor const& norm_weight, torch::stable::Tensor& out, |
| 431 | double scale, double norm_eps, const std::string& output_gate_activation) { |
| 432 | using torch::headeronly::ScalarType; |
| 433 | |
| 434 | STD_TORCH_CHECK( |
| 435 | mixed_qkv.is_cuda() && mixed_qkv.scalar_type() == ScalarType::BFloat16, |
| 436 | "mixed_qkv must be a CUDA bfloat16 tensor"); |
| 437 | STD_TORCH_CHECK(a.is_cuda() && a.scalar_type() == ScalarType::BFloat16, |
| 438 | "a must be a CUDA bfloat16 tensor"); |
| 439 | STD_TORCH_CHECK(b.is_cuda() && b.scalar_type() == ScalarType::BFloat16, |
| 440 | "b must be a CUDA bfloat16 tensor"); |
| 441 | STD_TORCH_CHECK(a_log.is_cuda() && a_log.scalar_type() == ScalarType::Float, |
| 442 | "A_log must be a CUDA float32 tensor"); |
| 443 | const auto dt_bias_scalar_type = dt_bias.scalar_type(); |
| 444 | STD_TORCH_CHECK( |
| 445 | dt_bias.is_cuda() && (dt_bias_scalar_type == ScalarType::Float || |
| 446 | dt_bias_scalar_type == ScalarType::BFloat16 || |
| 447 | dt_bias_scalar_type == ScalarType::Half), |
| 448 | "dt_bias must be a CUDA float32, bfloat16, or float16 tensor"); |
| 449 | STD_TORCH_CHECK( |
| 450 | state_indices.is_cuda() && state_indices.scalar_type() == ScalarType::Int, |
| 451 | "state_indices must be a CUDA int32 tensor"); |
| 452 | STD_TORCH_CHECK( |
| 453 | cu_seqlens.is_cuda() && cu_seqlens.scalar_type() == ScalarType::Int, |
| 454 | "cu_seqlens must be a CUDA int32 tensor"); |
| 455 | STD_TORCH_CHECK(num_accepted_tokens.is_cuda() && |
| 456 | num_accepted_tokens.scalar_type() == ScalarType::Int, |
| 457 | "num_accepted_tokens must be a CUDA int32 tensor"); |
| 458 | const auto state_scalar_type = state.scalar_type(); |
| 459 | STD_TORCH_CHECK( |
| 460 | state.is_cuda() && (state_scalar_type == ScalarType::Float || |
| 461 | state_scalar_type == ScalarType::BFloat16), |
| 462 | "state must be a CUDA float32 or bfloat16 tensor"); |
| 463 | STD_TORCH_CHECK(output_gate.is_cuda() && |
| 464 | output_gate.scalar_type() == ScalarType::BFloat16, |
| 465 | "output_gate must be a CUDA bfloat16 tensor"); |
| 466 | STD_TORCH_CHECK(norm_weight.is_cuda() && |
| 467 | (norm_weight.scalar_type() == ScalarType::Float || |
| 468 | norm_weight.scalar_type() == ScalarType::BFloat16), |
| 469 | "norm_weight must be a CUDA float32 or bfloat16 tensor"); |
| 470 | STD_TORCH_CHECK(out.is_cuda() && out.scalar_type() == ScalarType::BFloat16, |
| 471 | "out must be a CUDA bfloat16 tensor"); |
| 472 | STD_TORCH_CHECK( |
| 473 | output_gate_activation == "silu" || output_gate_activation == "sigmoid", |
| 474 | "output_gate_activation must be 'silu' or 'sigmoid'"); |
| 475 | |
| 476 | STD_TORCH_CHECK(mixed_qkv.dim() == 2, |
| 477 | "mixed_qkv must have shape [L, 2 * H * 128 + HV * 128]"); |
| 478 | const int num_tokens = static_cast<int>(mixed_qkv.size(0)); |
| 479 | STD_TORCH_CHECK(num_tokens > 0, |
| 480 | "GDN decode MTP fusion requires at least one token"); |
| 481 | STD_TORCH_CHECK( |
| 482 | state.dim() == 4 && state.size(2) == kDimV && state.size(3) == kDimK, |
| 483 | "state must have shape [slots, HV, 128, 128]"); |
| 484 | const int num_value_heads = static_cast<int>(state.size(1)); |
| 485 | const int64_t key_width = |
| 486 | mixed_qkv.size(1) - static_cast<int64_t>(num_value_heads) * kDimV; |
| 487 | STD_TORCH_CHECK(key_width > 0 && key_width % (2 * kDimK) == 0, |
| 488 | "mixed_qkv width is inconsistent with state"); |
| 489 | const int num_key_heads = static_cast<int>(key_width / (2 * kDimK)); |
| 490 | const int value_heads_per_key_head = num_value_heads / num_key_heads; |
| 491 | STD_TORCH_CHECK( |
| 492 | num_value_heads % num_key_heads == 0 && |
| 493 | ((value_heads_per_key_head >= 1 && value_heads_per_key_head <= 4) || |
| 494 | value_heads_per_key_head == 8), |
| 495 | "GDN decode MTP fusion requires HV/H in {1, 2, 3, 4, 8}"); |
| 496 | |
| 497 | STD_TORCH_CHECK(state_indices.dim() == 2 && state_indices.size(0) > 0 && |
| 498 | state_indices.size(1) > 0 && |
| 499 | state_indices.size(1) <= kMaxMtpTokens, |
| 500 | "state_indices must have shape [N, S] with 1 <= S <= 8"); |
| 501 | const int num_requests = static_cast<int>(state_indices.size(0)); |
| 502 | STD_TORCH_CHECK( |
| 503 | cu_seqlens.dim() == 1 && cu_seqlens.numel() == num_requests + 1, |
| 504 | "cu_seqlens must have N + 1 elements"); |
| 505 | STD_TORCH_CHECK(num_accepted_tokens.dim() == 1 && |
| 506 | num_accepted_tokens.numel() == num_requests, |
| 507 | "num_accepted_tokens must have N elements"); |
| 508 | STD_TORCH_CHECK( |
| 509 | a.dim() == 2 && a.size(0) == num_tokens && a.size(1) == num_value_heads, |
| 510 | "a must have shape [L, HV]"); |
| 511 | STD_TORCH_CHECK( |
| 512 | b.dim() == 2 && b.size(0) == num_tokens && b.size(1) == num_value_heads, |
| 513 | "b must have shape [L, HV]"); |
| 514 | STD_TORCH_CHECK(a_log.is_contiguous() && a_log.numel() == num_value_heads, |
| 515 | "A_log must be contiguous with HV elements"); |
| 516 | STD_TORCH_CHECK(dt_bias.is_contiguous() && dt_bias.numel() == num_value_heads, |
| 517 | "dt_bias must be contiguous with HV elements"); |
| 518 | STD_TORCH_CHECK(state_indices.is_contiguous(), |
| 519 | "state_indices must be contiguous"); |
| 520 | STD_TORCH_CHECK(cu_seqlens.is_contiguous(), "cu_seqlens must be contiguous"); |
| 521 | STD_TORCH_CHECK(num_accepted_tokens.is_contiguous(), |
| 522 | "num_accepted_tokens must be contiguous"); |
| 523 | STD_TORCH_CHECK(output_gate.dim() == 3 && output_gate.size(0) == num_tokens && |
| 524 | output_gate.size(1) == num_value_heads && |
| 525 | output_gate.size(2) == kDimV, |
| 526 | "output_gate must have shape [L, HV, 128]"); |
| 527 | STD_TORCH_CHECK(norm_weight.is_contiguous() && norm_weight.numel() == kDimV, |
| 528 | "norm_weight must be contiguous with 128 elements"); |
| 529 | STD_TORCH_CHECK(out.dim() == 3 && out.size(0) == num_tokens && |
| 530 | out.size(1) == num_value_heads && out.size(2) == kDimV, |
| 531 | "out must have shape [L, HV, 128]"); |
| 532 | STD_TORCH_CHECK(mixed_qkv.stride(1) == 1, |
| 533 | "mixed_qkv channels must be contiguous"); |
| 534 | STD_TORCH_CHECK(a.stride(1) == 1 && b.stride(1) == 1, |
| 535 | "a and b heads must be contiguous"); |
| 536 | STD_TORCH_CHECK(state.stride(0) >= num_value_heads * kDimV * kDimK && |
| 537 | state.stride(1) == kDimV * kDimK && |
| 538 | state.stride(2) == kDimK && state.stride(3) == 1, |
| 539 | "state must have contiguous [HV, 128, 128] slot contents"); |
| 540 | const int state_elements_per_copy = |
| 541 | state_scalar_type == ScalarType::Float ? 4 : 8; |
| 542 | STD_TORCH_CHECK(reinterpret_cast<uintptr_t>(state.data_ptr()) % 16 == 0 && |
| 543 | state.stride(0) % state_elements_per_copy == 0, |
| 544 | "state slots must preserve 16-byte alignment"); |
| 545 | STD_TORCH_CHECK(output_gate.stride(2) == 1 && output_gate.stride(1) == kDimV, |
| 546 | "output_gate head rows must be contiguous"); |
| 547 | STD_TORCH_CHECK(out.is_contiguous(), "out must be contiguous"); |
| 548 | STD_TORCH_CHECK(norm_eps >= 0.0, "norm_eps must be non-negative"); |
| 549 | |
| 550 | const GdnDecodeStrides strides{mixed_qkv.stride(0), a.stride(0), b.stride(0), |
| 551 | output_gate.stride(0), state.stride(0)}; |
| 552 | const auto* a_ptr = static_cast<const __nv_bfloat16*>(a.data_ptr()); |
| 553 | const auto* b_ptr = static_cast<const __nv_bfloat16*>(b.data_ptr()); |
| 554 | const auto* output_gate_ptr = |
| 555 | static_cast<const __nv_bfloat16*>(output_gate.data_ptr()); |
| 556 | const auto launch = [&]<typename StateT, int ValueHeadsPerKeyHead, |
| 557 | bool SigmoidGate>() { |
| 558 | launch_gdn_decode_post_conv_mtp<StateT, ValueHeadsPerKeyHead, SigmoidGate>( |
| 559 | mixed_qkv, a_log, dt_bias, state_indices, cu_seqlens, |
| 560 | num_accepted_tokens, state, norm_weight, out, a_ptr, b_ptr, |
| 561 | output_gate_ptr, num_key_heads, num_value_heads, scale, norm_eps, |
| 562 | strides); |
| 563 | }; |
| 564 | const auto dispatch_state_type = [&]<int ValueHeadsPerKeyHead>() { |
| 565 | if (state_scalar_type == ScalarType::Float) { |
| 566 | if (output_gate_activation == "sigmoid") { |
| 567 | launch.template operator()<float, ValueHeadsPerKeyHead, true>(); |
| 568 | } else { |
| 569 | launch.template operator()<float, ValueHeadsPerKeyHead, false>(); |
| 570 | } |
| 571 | } else { |
| 572 | if (output_gate_activation == "sigmoid") { |
| 573 | launch.template operator()<__nv_bfloat16, ValueHeadsPerKeyHead, true>(); |
| 574 | } else { |
| 575 | launch |
| 576 | .template operator()<__nv_bfloat16, ValueHeadsPerKeyHead, false>(); |
| 577 | } |
| 578 | } |
| 579 | }; |
| 580 | switch (value_heads_per_key_head) { |
| 581 | case 1: |
| 582 | dispatch_state_type.template operator()<1>(); |
| 583 | break; |
| 584 | case 2: |
| 585 | dispatch_state_type.template operator()<2>(); |
| 586 | break; |
| 587 | case 3: |
| 588 | dispatch_state_type.template operator()<3>(); |
| 589 | break; |
| 590 | case 4: |
| 591 | dispatch_state_type.template operator()<4>(); |
| 592 | break; |
| 593 | default: |
| 594 | dispatch_state_type.template operator()<8>(); |
| 595 | break; |
| 596 | } |
| 597 | } |