| 1 | // SPDX-License-Identifier: Apache-2.0 |
| 2 | // SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | #pragma once |
| 5 | |
| 6 | #if defined(ARM_BF16_SUPPORT) |
| 7 | |
| 8 | template <int K, int BLOCK_N, bool has_bias, bool has_silu> |
| 9 | struct tinygemm_kernel<at::BFloat16, K, BLOCK_N, has_bias, has_silu> { |
| 10 | static inline bfloat16x8_t load_bf16x8(const at::BFloat16* ptr) { |
| 11 | return vld1q_bf16(reinterpret_cast<const bfloat16_t*>(ptr)); |
| 12 | } |
| 13 | |
| 14 | // Input layout: A[m][channel] |
| 15 | // Weight layout: B[N/MICRO_N, K/2, MICRO_N, 2] |
| 16 | // The packed weights are reused across all M positions. |
| 17 | static inline void tinygemm_neon_8( |
| 18 | const at::BFloat16* __restrict__ A, const at::BFloat16* __restrict__ B, |
| 19 | at::BFloat16* __restrict__ C, const at::BFloat16* __restrict__ bias, |
| 20 | const at::BFloat16* __restrict__ conv_states, bool has_initial_state, |
| 21 | int64_t M, int64_t lda, bool is_first_token) { |
| 22 | bfloat16x8_t va[4]; |
| 23 | // Load 4 pairs of kernel weights into 4 registers |
| 24 | const bfloat16x8_t vb0 = load_bf16x8(B); // channels 0-3, taps 0-1 |
| 25 | const bfloat16x8_t vb1 = load_bf16x8(B + 8); // channels 4-7, taps 0-1 |
| 26 | const bfloat16x8_t vb2 = load_bf16x8(B + 16); // channels 0-3, taps 2-3 |
| 27 | const bfloat16x8_t vb3 = load_bf16x8(B + 24); // channels 4-7, taps 2-3 |
| 28 | |
| 29 | // Load the input data. A single register holds the values for a |
| 30 | // single token across all 8 channels for this microtile. |
| 31 | auto load_a = [&](int64_t m) { |
| 32 | if (m < 0 && is_first_token) { |
| 33 | if (!has_initial_state) { |
| 34 | return vreinterpretq_bf16_u16(vdupq_n_u16(0)); |
| 35 | } |
| 36 | return load_bf16x8(conv_states + (m + K - 1) * lda); |
| 37 | } |
| 38 | return load_bf16x8(A + m * lda); |
| 39 | }; |
| 40 | |
| 41 | // Load the previous inputs from the cached state if available. |
| 42 | va[1] = load_a(-3); |
| 43 | va[2] = load_a(-2); |
| 44 | va[3] = load_a(-1); |
| 45 | |
| 46 | // Iterate over the sequence block, updating the previous conv inputs |
| 47 | // and reading the next token's input data (8 channels). |
| 48 | for (int64_t m = 0; m < M; ++m) { |
| 49 | va[0] = va[1]; |
| 50 | va[1] = va[2]; |
| 51 | va[2] = va[3]; |
| 52 | va[3] = load_a(m); |
| 53 | |
| 54 | float32x4_t vc0; |
| 55 | float32x4_t vc1; |
| 56 | if constexpr (has_bias) { |
| 57 | const bfloat16x8_t bias_vec = load_bf16x8(bias); |
| 58 | vc0 = vcvtq_low_f32_bf16(bias_vec); |
| 59 | vc1 = vcvtq_high_f32_bf16(bias_vec); |
| 60 | } else { |
| 61 | vc0 = vdupq_n_f32(0.0f); |
| 62 | vc1 = vdupq_n_f32(0.0f); |
| 63 | } |
| 64 | |
| 65 | // Use the zip operation to match the input layout to the weights layout |
| 66 | bfloat16x8_t va_pair = |
| 67 | vreinterpretq_bf16_u16( // channels 0-3, inputs for taps 0-1 |
| 68 | vzip1q_u16(vreinterpretq_u16_bf16(va[0]), |
| 69 | vreinterpretq_u16_bf16(va[1]))); |
| 70 | vc0 = vbfdotq_f32(vc0, va_pair, |
| 71 | vb0); // 4 independent length-two dot products |
| 72 | va_pair = vreinterpretq_bf16_u16( // channels 4-7, inputs for taps 0-1 |
| 73 | vzip2q_u16(vreinterpretq_u16_bf16(va[0]), |
| 74 | vreinterpretq_u16_bf16(va[1]))); |
| 75 | vc1 = vbfdotq_f32(vc1, va_pair, vb1); |
| 76 | va_pair = vreinterpretq_bf16_u16( // channels 0-3, inputs for taps 2-3 |
| 77 | vzip1q_u16(vreinterpretq_u16_bf16(va[2]), |
| 78 | vreinterpretq_u16_bf16(va[3]))); |
| 79 | // Add to the previous result for the same channels |
| 80 | vc0 = vbfdotq_f32(vc0, va_pair, vb2); |
| 81 | va_pair = vreinterpretq_bf16_u16( // channels 4-7, inputs for taps 2-3 |
| 82 | vzip2q_u16(vreinterpretq_u16_bf16(va[2]), |
| 83 | vreinterpretq_u16_bf16(va[3]))); |
| 84 | vc1 = vbfdotq_f32(vc1, va_pair, vb3); |
| 85 | |
| 86 | using fVec = at::vec::Vectorized<float>; |
| 87 | fVec out0(vc0); |
| 88 | fVec out1(vc1); |
| 89 | if constexpr (has_silu) { |
| 90 | const fVec one(1.0f); |
| 91 | out0 = out0 / (one + out0.neg().fexp_u20()); |
| 92 | out1 = out1 / (one + out1.neg().fexp_u20()); |
| 93 | } |
| 94 | |
| 95 | // Convert the 2 x 4 fp32 results to 1 x 8 bf16 |
| 96 | const bfloat16x8_t out = |
| 97 | vcombine_bf16(vcvt_bf16_f32(static_cast<float32x4_t>(out0)), |
| 98 | vcvt_bf16_f32(static_cast<float32x4_t>(out1))); |
| 99 | vst1q_bf16(reinterpret_cast<bfloat16_t*>(C + m * lda), out); |
| 100 | } |
| 101 | } |
| 102 | |
| 103 | static inline void apply(const at::BFloat16* __restrict__ A, |
| 104 | const at::BFloat16* __restrict__ B, |
| 105 | at::BFloat16* __restrict__ C, |
| 106 | const at::BFloat16* __restrict__ bias, |
| 107 | const at::BFloat16* __restrict__ conv_states, |
| 108 | bool has_initial_state, int64_t M, int64_t lda, |
| 109 | bool is_first_token) { |
| 110 | static_assert(K == 4); |
| 111 | constexpr int64_t MICRO_N = 8; |
| 112 | static_assert(BLOCK_N % MICRO_N == 0); |
| 113 | |
| 114 | // BLOCK_N is always 32 or 64. M is the block sequence length. |
| 115 | // Input layout: A[m][channel] |
| 116 | // Weight layout: B[N/MICRO_N, K/2, MICRO_N, 2] |
| 117 | // Use an 8 channel microtile and iterate over the 8 channel blocks. |
| 118 | for (int64_t n = 0; n < BLOCK_N; n += MICRO_N) { |
| 119 | tinygemm_neon_8(A + n, B + n * K, C + n, has_bias ? bias + n : nullptr, |
| 120 | conv_states ? conv_states + n : nullptr, |
| 121 | has_initial_state, M, lda, is_first_token); |
| 122 | } |
| 123 | } |
| 124 | }; |
| 125 | |
| 126 | #endif |