#ifndef AMX_HPP #define AMX_HPP #include #include #include #include #include #include #include #include #include #include "llama.cpp/ggml-quants.h" // Include the split AMX headers #include "amx_config.hpp" #include "amx_kernels.hpp" namespace amx { static inline __m512 exp_avx512(__m512 x) { const __m512 log2e = _mm512_set1_ps(1.44269504089f); const __m512 c1 = _mm512_set1_ps(0.69314718056f); __m512 y = _mm512_mul_ps(x, log2e); __m512i int_part = _mm512_cvtps_epi32(y); __m512 frac_part = _mm512_sub_ps(y, _mm512_cvtepi32_ps(int_part)); const __m512 poly_1 = _mm512_set1_ps(0.9999999995f); const __m512 poly_2 = _mm512_set1_ps(0.6931471805f); const __m512 poly_3 = _mm512_set1_ps(0.2402265069f); const __m512 poly_4 = _mm512_set1_ps(0.0555041087f); const __m512 poly_5 = _mm512_set1_ps(0.0096181291f); const __m512 poly_6 = _mm512_set1_ps(0.0013333558f); __m512 frac_exp = _mm512_fmadd_ps( _mm512_fmadd_ps(_mm512_fmadd_ps(_mm512_fmadd_ps(_mm512_fmadd_ps(poly_6, frac_part, poly_5), frac_part, poly_4), frac_part, poly_3), frac_part, poly_2), frac_part, poly_1); __m512 two_pow_i = _mm512_scalef_ps(_mm512_set1_ps(1.0f), _mm512_cvtepi32_ps(int_part)); return _mm512_mul_ps(two_pow_i, frac_exp); } static inline __m512 act_fn(__m512 gate_val, __m512 up_val, float swiglu_limit = 0.0f) { // DeepSeek V4-Flash 2604B asymmetric SwiGLU clamp. swiglu_limit > 0 // applies the same clamp the trtllm `gemm1_clamp_limit` and the sglang // deep_gemm `_apply_swiglu_limit` use: // gate = clamp(gate, max=limit) // one-sided (pre-silu) // up = clamp(up, min=-limit, max=limit) // symmetric // The branch is on a runtime float; for swiglu_limit==0.0f (every non- // MXFP4 dtype today) the predictor stays on the fall-through path and // adds at most one cmp+jmp per 32-lane tile. // Origin: kt-sglang è€Šćˆ. if (swiglu_limit > 0.0f) { const __m512 pos_lim = _mm512_set1_ps(swiglu_limit); const __m512 neg_lim = _mm512_set1_ps(-swiglu_limit); gate_val = _mm512_min_ps(gate_val, pos_lim); up_val = _mm512_min_ps(up_val, pos_lim); up_val = _mm512_max_ps(up_val, neg_lim); } __m512 neg_gate_val = _mm512_sub_ps(_mm512_setzero_ps(), gate_val); // Clamp neg_gate_val to avoid exp overflow (exp(88) overflows for float32) const __m512 max_exp_input = _mm512_set1_ps(88.0f); neg_gate_val = _mm512_min_ps(neg_gate_val, max_exp_input); __m512 exp_neg_gate = exp_avx512(neg_gate_val); __m512 denom = _mm512_add_ps(_mm512_set1_ps(1.0f), exp_neg_gate); __m512 act_val = _mm512_div_ps(gate_val, denom); return _mm512_mul_ps(act_val, up_val); } // MiniMax M3 "swigluoai" activation + DeepSeek V4 "silu" unified entry point. // alpha > 0 → swigluoai: gate * sigmoid(gate * alpha) * (up + 1), symmetric clamp // alpha == 0 → falls back to standard silu path above static inline __m512 act_fn(__m512 gate_val, __m512 up_val, float swiglu_limit, float swiglu_alpha) { if (swiglu_alpha > 0.0f) { // --- MiniMax M3 swigluoai path --- // Symmetric clamp on both gate and up (differs from V4 which clamps gate one-sided) if (swiglu_limit > 0.0f) { const __m512 pos_lim = _mm512_set1_ps(swiglu_limit); const __m512 neg_lim = _mm512_set1_ps(-swiglu_limit); gate_val = _mm512_min_ps(gate_val, pos_lim); gate_val = _mm512_max_ps(gate_val, neg_lim); up_val = _mm512_min_ps(up_val, pos_lim); up_val = _mm512_max_ps(up_val, neg_lim); } // sigmoid(gate * alpha) = 1 / (1 + exp(-gate * alpha)) __m512 neg_ga = _mm512_mul_ps(gate_val, _mm512_set1_ps(-swiglu_alpha)); neg_ga = _mm512_min_ps(neg_ga, _mm512_set1_ps(88.0f)); __m512 exp_neg = exp_avx512(neg_ga); __m512 sigmoid_val = _mm512_div_ps(_mm512_set1_ps(1.0f), _mm512_add_ps(_mm512_set1_ps(1.0f), exp_neg)); // gate * sigmoid(gate * alpha) * (up + 1) __m512 up_plus_1 = _mm512_add_ps(up_val, _mm512_set1_ps(1.0f)); return _mm512_mul_ps(_mm512_mul_ps(gate_val, sigmoid_val), up_plus_1); } return act_fn(gate_val, up_val, swiglu_limit); } #define AMX_DISPATCH_QTYPES(QA, QB, ...) \ [&] { \ switch (QB) { \ case GGML_TYPE_Q8_0: { \ using qb = block_q8_0; \ switch (QA) { \ case GGML_TYPE_Q4_0: { \ using qa = block_q4_0; \ return __VA_ARGS__(); \ } \ case GGML_TYPE_Q8_0: { \ using qa = block_q8_0; \ return __VA_ARGS__(); \ } \ default: \ throw std::runtime_error("Unsupported quantized data type"); \ } \ } \ case GGML_TYPE_Q8_K: { \ using qb = block_q8_K; \ switch (QA) { \ case GGML_TYPE_Q4_K: { \ using qa = block_q4_K; \ return __VA_ARGS__(); \ } \ default: \ throw std::runtime_error("Unsupported quantized data type"); \ } \ } \ case GGML_TYPE_BF16: { \ using qb = ggml_bf16_t; \ switch (QA) { \ case GGML_TYPE_BF16: { \ using qa = ggml_bf16_t; \ return __VA_ARGS__(); \ } \ default: \ throw std::runtime_error("Unsupported quantized data type"); \ } \ } \ default: \ throw std::runtime_error("Unsupported quantized data type"); \ } \ }() inline void gemm(int m, int n, int k, const void* a, size_t lda, int type_a, const void* b, size_t ldb, int type_b, void* c, size_t ldc, int type_c, int ith, int nth) { assert(reinterpret_cast(c) % 64 == 0); assert(ldc % 64 == 0); assert(type_c == GGML_TYPE_F32); float* cs = (float*)c; AMX_DISPATCH_QTYPES(type_a, type_b, [&]() { mat_mul(m, n, k, (qa*)a, lda, (qb*)b, ldb, cs, ldc, ith, nth); }); } inline void init_tile(int type_a, int type_b, int type_c) { #ifdef HAVE_AMX enable_amx(); assert(type_c == GGML_TYPE_F32); AMX_DISPATCH_QTYPES(type_a, type_b, []() { return GemmKernel::type::config(); }); #endif } inline int recommended_nth(int m, int n, int k, int type_a, int type_b, int type_c) { assert(type_c == GGML_TYPE_F32); return AMX_DISPATCH_QTYPES(type_a, type_b, [&]() { return GemmKernel::type::recommended_nth(m); }); } } // namespace amx #endif // AMX_HPP