wehub-resource-sync
ec436095dd
Book-CI / test (macos-latest) (push) Has been cancelled
Book-CI / test (ubuntu-latest) (push) Has been cancelled
Book-CI / test (windows-latest) (push) Has been cancelled
Release Fake Tag / publish (push) Has been cancelled
Deploy / deploy (macos-latest) (push) Has been cancelled
Deploy / deploy (ubuntu-latest) (push) Has been cancelled
Deploy / deploy (windows-latest) (push) Has been cancelled
Release to PyPI / Build & publish sglang-kt (push) Has been cancelled
Release to PyPI / Build kt-kernel (Python 3.11) (push) Has been cancelled
Release to PyPI / Build kt-kernel (Python 3.12) (push) Has been cancelled
Release to PyPI / Publish kt-kernel to PyPI (push) Has been cancelled
172 lines
8.3 KiB
C++
172 lines
8.3 KiB
C++
#ifndef AMX_HPP
|
|
#define AMX_HPP
|
|
#include <emmintrin.h>
|
|
#include <immintrin.h>
|
|
#include <stdlib.h>
|
|
#include <sys/syscall.h>
|
|
#include <tmmintrin.h>
|
|
#include <unistd.h>
|
|
|
|
#include <cassert>
|
|
#include <cstdio>
|
|
#include <stdexcept>
|
|
|
|
#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<intptr_t>(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<qa, qb, float>::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<qa, qb, float>::type::recommended_nth(m); });
|
|
}
|
|
|
|
} // namespace amx
|
|
|
|
#endif // AMX_HPP
|