Safetensors
GGUF
Turkish
llama
Llama-3
instruct
finetune
chatml
gpt4
synthetic data
distillation
function calling
json mode
axolotl
roleplaying
chat
Instructions to use tda45/TdAI with libraries, inference providers, notebooks, and local apps. Follow these links to get started.
- Libraries
- llama-cpp-python
How to use tda45/TdAI with llama-cpp-python:
# !pip install llama-cpp-python from llama_cpp import Llama llm = Llama.from_pretrained( repo_id="tda45/TdAI", filename="llama.cpp/models/ggml-vocab-aquila.gguf", )
output = llm( "Once upon a time,", max_tokens=512, echo=True ) print(output)
- Notebooks
- Google Colab
- Kaggle
- Local Apps Settings
- llama.cpp
How to use tda45/TdAI with llama.cpp:
Install (macOS, Linux)
curl -LsSf https://llama.app/install.sh | sh # Start a local OpenAI-compatible server with a web UI: llama serve -hf tda45/TdAI # Run inference directly in the terminal: llama cli -hf tda45/TdAI
Install from WinGet (Windows)
winget install llama.cpp # Start a local OpenAI-compatible server with a web UI: llama serve -hf tda45/TdAI # Run inference directly in the terminal: llama cli -hf tda45/TdAI
Use pre-built binary
# Download pre-built binary from: # https://github.com/ggerganov/llama.cpp/releases # Start a local OpenAI-compatible server with a web UI: ./llama-server -hf tda45/TdAI # Run inference directly in the terminal: ./llama-cli -hf tda45/TdAI
Build from source code
git clone https://github.com/ggerganov/llama.cpp.git cd llama.cpp cmake -B build cmake --build build -j --target llama-server llama-cli # Start a local OpenAI-compatible server with a web UI: ./build/bin/llama-server -hf tda45/TdAI # Run inference directly in the terminal: ./build/bin/llama-cli -hf tda45/TdAI
Use Docker
docker model run hf.co/tda45/TdAI
- LM Studio
- Jan
- Ollama
How to use tda45/TdAI with Ollama:
ollama run hf.co/tda45/TdAI
- Unsloth Studio
How to use tda45/TdAI with Unsloth Studio:
Install Unsloth Studio (macOS, Linux, WSL)
curl -fsSL https://unsloth.ai/install.sh | sh # Run unsloth studio unsloth studio -H 0.0.0.0 -p 8888 # Then open http://localhost:8888 in your browser # Search for tda45/TdAI to start chatting
Install Unsloth Studio (Windows)
irm https://unsloth.ai/install.ps1 | iex # Run unsloth studio unsloth studio -H 0.0.0.0 -p 8888 # Then open http://localhost:8888 in your browser # Search for tda45/TdAI to start chatting
Using HuggingFace Spaces for Unsloth
# No setup required # Open https://huggingface.co/spaces/unsloth/studio in your browser # Search for tda45/TdAI to start chatting
- Atomic Chat new
- Docker Model Runner
How to use tda45/TdAI with Docker Model Runner:
docker model run hf.co/tda45/TdAI
- Lemonade
How to use tda45/TdAI with Lemonade:
Pull the model
# Download Lemonade from https://lemonade-server.ai/ lemonade pull tda45/TdAI
Run and chat with the model
lemonade run user.TdAI-{{QUANT_TAG}}List all available models
lemonade list
| namespace spacemit_kernels::rvv { | |
| namespace { | |
| auto align_up(size_t value, size_t alignment) { | |
| return (value + alignment - 1) / alignment * alignment; | |
| } | |
| static inline bool flash_attn_ext_supported_d_vlen1024_vf16(int64_t d) { | |
| return d > 0 && d <= 128; | |
| } | |
| static inline bool flash_attn_ext_supported_shape_vlen1024_vf16(int64_t DK, int64_t DV) { | |
| return flash_attn_ext_supported_d_vlen1024_vf16(DK) && flash_attn_ext_supported_d_vlen1024_vf16(DV); | |
| } | |
| static inline float reduce_sum_f32m4_vlen1024(vfloat32m4_t v, size_t vl) { | |
| vfloat32m1_t s_v = __riscv_vfmv_v_f_f32m1(0.0f, 1); | |
| s_v = __riscv_vfredusum_vs_f32m4_f32m1(v, s_v, vl); | |
| return __riscv_vfmv_f_s_f32m1_f32(s_v); | |
| } | |
| static inline float reduce_sum_f32m2_vlen1024(vfloat32m2_t v, size_t vl) { | |
| vfloat32m1_t s_v = __riscv_vfmv_v_f_f32m1(0.0f, 1); | |
| s_v = __riscv_vfredusum_vs_f32m2_f32m1(v, s_v, vl); | |
| return __riscv_vfmv_f_s_f32m1_f32(s_v); | |
| } | |
| // Adapted from ggml_v_expf_m2 in vec.h. This is accurate enough for softmax. | |
| static inline vfloat32m2_t rvv_expf_approx_f32m2(vfloat32m2_t x, size_t vl) { | |
| const vfloat32m2_t r = __riscv_vfmv_v_f_f32m2(0x1.8p23f, vl); | |
| const vfloat32m2_t z = __riscv_vfmacc_vf_f32m2(r, 0x1.715476p+0f, x, vl); | |
| const vfloat32m2_t n = __riscv_vfsub_vv_f32m2(z, r, vl); | |
| const vfloat32m2_t b = | |
| __riscv_vfnmsac_vf_f32m2(__riscv_vfnmsac_vf_f32m2(x, 0x1.62e4p-1f, n, vl), 0x1.7f7d1cp-20f, n, vl); | |
| const vuint32m2_t e = __riscv_vsll_vx_u32m2(__riscv_vreinterpret_v_f32m2_u32m2(z), 23, vl); | |
| const vfloat32m2_t k = __riscv_vreinterpret_v_u32m2_f32m2(__riscv_vadd_vx_u32m2(e, 0x3f800000, vl)); | |
| const vbool16_t c = __riscv_vmfgt_vf_f32m2_b16(__riscv_vfabs_v_f32m2(n, vl), 126.0f, vl); | |
| const vfloat32m2_t u = __riscv_vfmul_vv_f32m2(b, b, vl); | |
| const vfloat32m2_t j = __riscv_vfmacc_vv_f32m2( | |
| __riscv_vfmul_vf_f32m2(b, 0x1.ffffecp-1f, vl), | |
| __riscv_vfmacc_vv_f32m2( | |
| __riscv_vfmacc_vf_f32m2(__riscv_vfmv_v_f_f32m2(0x1.fffdb6p-2f, vl), 0x1.555e66p-3f, b, vl), | |
| __riscv_vfmacc_vf_f32m2(__riscv_vfmv_v_f_f32m2(0x1.573e2ep-5f, vl), 0x1.0e4020p-7f, b, vl), u, vl), | |
| u, vl); | |
| if (!__riscv_vcpop_m_b16(c, vl)) { | |
| return __riscv_vfmacc_vv_f32m2(k, j, k, vl); | |
| } | |
| const vbool16_t dm = __riscv_vmfle_vf_f32m2_b16(n, 0.0f, vl); | |
| const vuint32m2_t d = __riscv_vmerge_vxm_u32m2(__riscv_vmv_v_x_u32m2(0, vl), 0x82000000, dm, vl); | |
| const vfloat32m2_t s1 = __riscv_vreinterpret_v_u32m2_f32m2(__riscv_vadd_vx_u32m2(d, 0x7f000000, vl)); | |
| const vfloat32m2_t s2 = __riscv_vreinterpret_v_u32m2_f32m2(__riscv_vsub_vv_u32m2(e, d, vl)); | |
| const vfloat32m2_t r1 = | |
| __riscv_vmerge_vvm_f32m2(__riscv_vfmacc_vv_f32m2(k, k, j, vl), | |
| __riscv_vfmul_vv_f32m2(__riscv_vfmacc_vv_f32m2(s2, s2, j, vl), s1, vl), c, vl); | |
| return __riscv_vmerge_vvm_f32m2(r1, __riscv_vfmul_vv_f32m2(s1, s1, vl), | |
| __riscv_vmfgt_vf_f32m2_b16(__riscv_vfabs_v_f32m2(n, vl), 192.0f, vl), vl); | |
| } | |
| static inline vfloat32m2_t rvv_tanh_approx_f32m2(vfloat32m2_t x, size_t vl) { | |
| const vfloat32m2_t abs_x = __riscv_vfabs_v_f32m2(x, vl); | |
| const vfloat32m2_t neg_2_abs = __riscv_vfmul_vf_f32m2(abs_x, -2.0f, vl); | |
| const vfloat32m2_t exp_term = rvv_expf_approx_f32m2(neg_2_abs, vl); | |
| const vfloat32m2_t numerator = __riscv_vfsub_vf_f32m2(exp_term, 1.0f, vl); | |
| const vfloat32m2_t denominator = __riscv_vfadd_vf_f32m2(exp_term, 1.0f, vl); | |
| const vfloat32m2_t tanh_abs = __riscv_vfneg_v_f32m2(__riscv_vfdiv_vv_f32m2(numerator, denominator, vl), vl); | |
| const vbool16_t neg_mask = __riscv_vmflt_vf_f32m2_b16(x, 0.0f, vl); | |
| const vfloat32m2_t tanh_neg = __riscv_vfneg_v_f32m2(tanh_abs, vl); | |
| return __riscv_vmerge_vvm_f32m2(tanh_abs, tanh_neg, neg_mask, vl); | |
| } | |
| static void rvv_softcap_tanh_inplace_f32(float * dst, int64_t dst_stride, int64_t tile_rows, int64_t n, float softcap) { | |
| for (int tq = 0; tq < tile_rows; ++tq, dst += dst_stride) { | |
| float * dst_row = dst; | |
| int64_t remaining = n; | |
| while (remaining > 0) { | |
| const size_t vl = __riscv_vsetvl_e32m2(remaining); | |
| vfloat32m2_t v = __riscv_vle32_v_f32m2(dst_row, vl); | |
| v = rvv_tanh_approx_f32m2(v, vl); | |
| v = __riscv_vfmul_vf_f32m2(v, softcap, vl); | |
| __riscv_vse32_v_f32m2(dst_row, v, vl); | |
| dst_row += vl; | |
| remaining -= vl; | |
| } | |
| } | |
| } | |
| static inline float rvv_softmax_exp_inplace_f32(float * dst, int64_t n, float max_value) { | |
| float row_sum = 0.0f; | |
| while (n > 0) { | |
| const size_t vl = __riscv_vsetvl_e32m2(n); | |
| vfloat32m2_t v = __riscv_vle32_v_f32m2(dst, vl); | |
| v = __riscv_vfsub_vf_f32m2(v, max_value, vl); | |
| v = rvv_expf_approx_f32m2(v, vl); | |
| __riscv_vse32_v_f32m2(dst, v, vl); | |
| row_sum += reduce_sum_f32m2_vlen1024(v, vl); | |
| dst += vl; | |
| n -= vl; | |
| } | |
| return row_sum; | |
| } | |
| static inline float rvv_add_max_inplace_f32(float * dst, const float * src, int64_t n) { | |
| float max_val = -INFINITY; | |
| while (n > 0) { | |
| const size_t vl = __riscv_vsetvl_e32m4(n); | |
| vfloat32m4_t vdst = __riscv_vle32_v_f32m4(dst, vl); | |
| vfloat32m4_t vsrc = __riscv_vle32_v_f32m4(src, vl); | |
| vdst = __riscv_vfadd_vv_f32m4(vdst, vsrc, vl); | |
| __riscv_vse32_v_f32m4(dst, vdst, vl); | |
| vfloat32m1_t seed = __riscv_vfmv_v_f_f32m1(max_val, 1); | |
| seed = __riscv_vfredmax_vs_f32m4_f32m1(vdst, seed, vl); | |
| max_val = __riscv_vfmv_f_s_f32m1_f32(seed); | |
| dst += vl; | |
| src += vl; | |
| n -= vl; | |
| } | |
| return max_val; | |
| } | |
| static inline float rvv_softcap_add_max_inplace_f32(float * dst, const float * src, int64_t n, float softcap) { | |
| if (softcap == 0.0f) { | |
| return rvv_add_max_inplace_f32(dst, src, n); | |
| } | |
| float max_val = -INFINITY; | |
| while (n > 0) { | |
| const size_t vl = __riscv_vsetvl_e32m2(n); | |
| vfloat32m2_t vdst = __riscv_vle32_v_f32m2(dst, vl); | |
| vfloat32m2_t vsrc = __riscv_vle32_v_f32m2(src, vl); | |
| vdst = rvv_tanh_approx_f32m2(vdst, vl); | |
| vdst = __riscv_vfmul_vf_f32m2(vdst, softcap, vl); | |
| vdst = __riscv_vfadd_vv_f32m2(vdst, vsrc, vl); | |
| __riscv_vse32_v_f32m2(dst, vdst, vl); | |
| vfloat32m1_t seed = __riscv_vfmv_v_f_f32m1(max_val, 1); | |
| seed = __riscv_vfredmax_vs_f32m2_f32m1(vdst, seed, vl); | |
| max_val = __riscv_vfmv_f_s_f32m1_f32(seed); | |
| dst += vl; | |
| src += vl; | |
| n -= vl; | |
| } | |
| return max_val; | |
| } | |
| static inline void rvv_zero_f32(float * dst, int64_t n) { | |
| while (n > 0) { | |
| const size_t vl = __riscv_vsetvl_e32m4(n); | |
| const vfloat32m4_t z = __riscv_vfmv_v_f_f32m4(0.0f, vl); | |
| __riscv_vse32_v_f32m4(dst, z, vl); | |
| dst += vl; | |
| n -= vl; | |
| } | |
| } | |
| static inline void rvv_scale_f32(float * dst, float scale, int64_t n) { | |
| while (n > 0) { | |
| const size_t vl = __riscv_vsetvl_e32m4(n); | |
| vfloat32m4_t v = __riscv_vle32_v_f32m4(dst, vl); | |
| v = __riscv_vfmul_vf_f32m4(v, scale, vl); | |
| __riscv_vse32_v_f32m4(dst, v, vl); | |
| dst += vl; | |
| n -= vl; | |
| } | |
| } | |
| static inline void rvv_add_inplace_f32(float * dst, | |
| int64_t dst_stride, | |
| const float * src, | |
| int64_t src_stride, | |
| int64_t tile_rows, | |
| int64_t n) { | |
| for (int tq = 0; tq < tile_rows; ++tq, dst += dst_stride, src += src_stride) { | |
| int64_t remaining = n; | |
| float * dst_row = dst; | |
| const float * src_row = src; | |
| while (remaining > 0) { | |
| const size_t vl = __riscv_vsetvl_e32m4(remaining); | |
| vfloat32m4_t vdst = __riscv_vle32_v_f32m4(dst_row, vl); | |
| vfloat32m4_t vsrc = __riscv_vle32_v_f32m4(src_row, vl); | |
| vdst = __riscv_vfadd_vv_f32m4(vdst, vsrc, vl); | |
| __riscv_vse32_v_f32m4(dst_row, vdst, vl); | |
| dst_row += vl; | |
| src_row += vl; | |
| remaining -= vl; | |
| } | |
| } | |
| } | |
| static inline float rvv_max_f32(const float * src, int64_t n) { | |
| float max_val = -INFINITY; | |
| while (n > 0) { | |
| const size_t vl = __riscv_vsetvl_e32m4(n); | |
| const vfloat32m4_t v = __riscv_vle32_v_f32m4(src, vl); | |
| vfloat32m1_t seed = __riscv_vfmv_v_f_f32m1(max_val, 1); | |
| seed = __riscv_vfredmax_vs_f32m4_f32m1(v, seed, vl); | |
| max_val = __riscv_vfmv_f_s_f32m1_f32(seed); | |
| src += vl; | |
| n -= vl; | |
| } | |
| return max_val; | |
| } | |
| static void rvv_pack_f32_as_scaled_f16(void * dst, | |
| int64_t dst_row_stride, | |
| const void * src, | |
| int64_t src_row_stride, | |
| int64_t tile_rows, | |
| int64_t n, | |
| float scale) { | |
| for (int tq = 0; tq < tile_rows; ++tq) { | |
| const float * row_ptr = (const float *) ((const char *) src + tq * src_row_stride); | |
| _Float16 * dst_row_ptr = (_Float16 *) ((char *) dst + tq * dst_row_stride); | |
| int64_t remaining = n; | |
| while (remaining > 0) { | |
| const size_t vl = __riscv_vsetvl_e32m4(remaining); | |
| vfloat32m4_t v32 = __riscv_vle32_v_f32m4(row_ptr, vl); | |
| v32 = __riscv_vfmul_vf_f32m4(v32, scale, vl); | |
| const vfloat16m2_t v16 = __riscv_vfncvt_f_f_w_f16m2(v32, vl); | |
| __riscv_vse16_v_f16m2(dst_row_ptr, v16, vl); | |
| dst_row_ptr += vl; | |
| row_ptr += vl; | |
| remaining -= vl; | |
| } | |
| } | |
| } | |
| static void rvv_pack_scaled_f16_as_f32(void * dst, | |
| int64_t dst_row_stride, | |
| const void * src, | |
| int64_t src_row_stride, | |
| int64_t tile_rows, | |
| int64_t n, | |
| float scale) { | |
| for (int tq = 0; tq < tile_rows; ++tq) { | |
| const _Float16 * row_ptr = (const _Float16 *) ((const char *) src + tq * src_row_stride); | |
| float * dst_row_ptr = (float *) ((char *) dst + tq * dst_row_stride); | |
| int64_t remaining = n; | |
| while (remaining > 0) { | |
| const size_t vl = __riscv_vsetvl_e16m2(remaining); | |
| const vfloat16m2_t v16 = __riscv_vle16_v_f16m2(row_ptr, vl); | |
| vfloat32m4_t v32 = __riscv_vfwcvt_f_f_v_f32m4(v16, vl); | |
| v32 = __riscv_vfmul_vf_f32m4(v32, scale, vl); | |
| __riscv_vse32_v_f32m4(dst_row_ptr, v32, vl); | |
| dst_row_ptr += vl; | |
| row_ptr += vl; | |
| remaining -= vl; | |
| } | |
| } | |
| } | |
| static void rvv_pack_scaled_f32_as_f32(void * dst, | |
| int64_t dst_row_stride, | |
| const void * src, | |
| int64_t src_row_stride, | |
| int64_t tile_rows, | |
| int64_t n, | |
| float * scale) { | |
| for (int tq = 0; tq < tile_rows; ++tq) { | |
| const float * row_ptr = (const float *) ((const char *) src + tq * src_row_stride); | |
| float * dst_row_ptr = (float *) ((char *) dst + tq * dst_row_stride); | |
| int64_t remaining = n; | |
| while (remaining > 0) { | |
| const size_t vl = __riscv_vsetvl_e32m4(remaining); | |
| vfloat32m4_t v32 = __riscv_vle32_v_f32m4(row_ptr, vl); | |
| v32 = __riscv_vfmul_vf_f32m4(v32, scale[tq], vl); | |
| __riscv_vse32_v_f32m4(dst_row_ptr, v32, vl); | |
| dst_row_ptr += vl; | |
| row_ptr += vl; | |
| remaining -= vl; | |
| } | |
| } | |
| } | |
| static inline void rvv_transposed_s32_mn_to_nm(int8_t * dst, | |
| int64_t n_dst_stride, | |
| int8_t * src, | |
| int64_t m_src_stride, | |
| int64_t m, | |
| int64_t n) { | |
| int8_t * in = src; | |
| int8_t * out = dst; | |
| __asm__ volatile( | |
| "vsetvli t0, zero, e32, m1, tu, mu \n\t" | |
| "mul t3, t0, %[os0] \n\t" | |
| "srli t2, %[isz0], 3 \n\t" | |
| "blez t2, M1%= \n\t" | |
| "LOOP_M8%=: \n\t" | |
| "addi a1, %[dst], 0 \n\t" | |
| "addi s1, %[src], 0 \n\t" | |
| "add s2, %[src], %[is0] \n\t" | |
| "add s3, s2, %[is0] \n\t" | |
| "add s4, s3, %[is0] \n\t" | |
| "add s5, s4, %[is0] \n\t" | |
| "add s6, s5, %[is0] \n\t" | |
| "add s7, s6, %[is0] \n\t" | |
| "add s8, s7, %[is0] \n\t" | |
| "addi t1, %[isz1], 0 \n\t" | |
| "LOOP_M8N%=: \n\t" | |
| "vsetvli t0, t1, e32, m1, tu, mu \n\t" | |
| "sub t1, t1, t0 \n\t" | |
| "vle32.v v0, (s1) \n\t" | |
| "sh2add s1, t0, s1 \n\t" | |
| "vle32.v v1, (s2) \n\t" | |
| "sh2add s2, t0, s2 \n\t" | |
| "vle32.v v2, (s3) \n\t" | |
| "sh2add s3, t0, s3 \n\t" | |
| "vle32.v v3, (s4) \n\t" | |
| "sh2add s4, t0, s4 \n\t" | |
| "vle32.v v4, (s5) \n\t" | |
| "sh2add s5, t0, s5 \n\t" | |
| "vle32.v v5, (s6) \n\t" | |
| "sh2add s6, t0, s6 \n\t" | |
| "vle32.v v6, (s7) \n\t" | |
| "sh2add s7, t0, s7 \n\t" | |
| "vle32.v v7, (s8) \n\t" | |
| "sh2add s8, t0, s8 \n\t" | |
| "vssseg8e32.v v0, (a1), %[os0] \n\t" | |
| "add a1, a1, t3 \n\t" | |
| "bnez t1, LOOP_M8N%= \n\t" | |
| "sh3add %[src], %[is0], %[src] \n\t" | |
| "addi %[dst], %[dst], 32 \n\t" | |
| "addi t2, t2, -1 \n\t" | |
| "bnez t2, LOOP_M8%= \n\t" | |
| "M1%=: \n\t" | |
| "andi t2, %[isz0], 7 \n\t" | |
| "blez t2, END%= \n\t" | |
| "LOOP_M1%=: \n\t" | |
| "addi a1, %[dst], 0 \n\t" | |
| "addi s1, %[src], 0 \n\t" | |
| "addi t1, %[isz1], 0 \n\t" | |
| "LOOP_M1N%=: \n\t" | |
| "vsetvli t0, t1, e32, m1, tu, mu \n\t" | |
| "sub t1, t1, t0 \n\t" | |
| "vle32.v v0, (s1) \n\t" | |
| "sh2add s1, t0, s1 \n\t" | |
| "vsse32.v v0, (a1), %[os0] \n\t" | |
| "add a1, a1, t3 \n\t" | |
| "bnez t1, LOOP_M1N%= \n\t" | |
| "add %[src], %[is0], %[src] \n\t" | |
| "addi %[dst], %[dst], 4 \n\t" | |
| "addi t2, t2, -1 \n\t" | |
| "bnez t2, LOOP_M1%= \n\t" | |
| "END%=: \n\t" | |
| : [src] "+r"(in), [dst] "+r"(out), [isz0] "+r"(m) | |
| : [isz1] "r"(n), [is0] "r"(m_src_stride), [os0] "r"(n_dst_stride) | |
| : "cc", "t0", "t1", "t2", "t3", "s1", "s2", "s3", "s4", "s5", "s6", "s7", "s8", "a1"); | |
| } | |
| static inline void rvv_transposed_s16_mn_to_nm(int8_t * dst, | |
| int64_t n_dst_stride, | |
| int8_t * src, | |
| int64_t m_src_stride, | |
| int64_t m, | |
| int64_t n) { | |
| int8_t * in = src; | |
| int8_t * out = dst; | |
| __asm__ volatile( | |
| "vsetvli t0, zero, e16, m1, tu, mu \n\t" | |
| "mul t3, t0, %[os0] \n\t" | |
| "srli t2, %[isz0], 3 \n\t" | |
| "blez t2, M1%= \n\t" | |
| "LOOP_M8%=: \n\t" | |
| "addi a1, %[dst], 0 \n\t" | |
| "addi s1, %[src], 0 \n\t" | |
| "add s2, %[src], %[is0] \n\t" | |
| "add s3, s2, %[is0] \n\t" | |
| "add s4, s3, %[is0] \n\t" | |
| "add s5, s4, %[is0] \n\t" | |
| "add s6, s5, %[is0] \n\t" | |
| "add s7, s6, %[is0] \n\t" | |
| "add s8, s7, %[is0] \n\t" | |
| "addi t1, %[isz1], 0 \n\t" | |
| "LOOP_M8N%=: \n\t" | |
| "vsetvli t0, t1, e16, m1, tu, mu \n\t" | |
| "sub t1, t1, t0 \n\t" | |
| "vle16.v v0, (s1) \n\t" | |
| "sh1add s1, t0, s1 \n\t" | |
| "vle16.v v1, (s2) \n\t" | |
| "sh1add s2, t0, s2 \n\t" | |
| "vle16.v v2, (s3) \n\t" | |
| "sh1add s3, t0, s3 \n\t" | |
| "vle16.v v3, (s4) \n\t" | |
| "sh1add s4, t0, s4 \n\t" | |
| "vle16.v v4, (s5) \n\t" | |
| "sh1add s5, t0, s5 \n\t" | |
| "vle16.v v5, (s6) \n\t" | |
| "sh1add s6, t0, s6 \n\t" | |
| "vle16.v v6, (s7) \n\t" | |
| "sh1add s7, t0, s7 \n\t" | |
| "vle16.v v7, (s8) \n\t" | |
| "sh1add s8, t0, s8 \n\t" | |
| "vssseg8e16.v v0, (a1), %[os0] \n\t" | |
| "add a1, a1, t3 \n\t" | |
| "bnez t1, LOOP_M8N%= \n\t" | |
| "sh3add %[src], %[is0], %[src] \n\t" | |
| "addi %[dst], %[dst], 16 \n\t" | |
| "addi t2, t2, -1 \n\t" | |
| "bnez t2, LOOP_M8%= \n\t" | |
| "M1%=: \n\t" | |
| "andi t2, %[isz0], 7 \n\t" | |
| "blez t2, END%= \n\t" | |
| "LOOP_M1%=: \n\t" | |
| "addi a1, %[dst], 0 \n\t" | |
| "addi s1, %[src], 0 \n\t" | |
| "addi t1, %[isz1], 0 \n\t" | |
| "LOOP_M1N%=: \n\t" | |
| "vsetvli t0, t1, e16, m1, tu, mu \n\t" | |
| "sub t1, t1, t0 \n\t" | |
| "vle16.v v0, (s1) \n\t" | |
| "sh1add s1, t0, s1 \n\t" | |
| "vsse16.v v0, (a1), %[os0] \n\t" | |
| "add a1, a1, t3 \n\t" | |
| "bnez t1, LOOP_M1N%= \n\t" | |
| "add %[src], %[is0], %[src] \n\t" | |
| "addi %[dst], %[dst], 2 \n\t" | |
| "addi t2, t2, -1 \n\t" | |
| "bnez t2, LOOP_M1%= \n\t" | |
| "END%=: \n\t" | |
| : [src] "+r"(in), [dst] "+r"(out), [isz0] "+r"(m) | |
| : [isz1] "r"(n), [is0] "r"(m_src_stride), [os0] "r"(n_dst_stride) | |
| : "cc", "t0", "t1", "t2", "t3", "s1", "s2", "s3", "s4", "s5", "s6", "s7", "s8", "a1"); | |
| } | |
| static inline void rvv_qk_dot_tile_f16_x1(float * dst, | |
| const _Float16 * q_row, | |
| const _Float16 * k_pack, | |
| int64_t dk, | |
| int64_t kv_tile) { | |
| const size_t vl = __riscv_vsetvl_e16m1(kv_tile); | |
| vfloat32m2_t acc = __riscv_vfmv_v_f_f32m2(0.0f, vl); | |
| for (int64_t d = 0; d < dk; ++d) { | |
| const vfloat16m1_t k_vec = __riscv_vle16_v_f16m1(k_pack + d * ggml_fa_tile_config::KV, vl); | |
| acc = __riscv_vfwmacc_vf_f32m2(acc, q_row[d], k_vec, vl); | |
| } | |
| __riscv_vse32_v_f32m2(dst, acc, vl); | |
| } | |
| static inline void rvv_qk_dot_tile_f16_x4(float * dst0, | |
| float * dst1, | |
| float * dst2, | |
| float * dst3, | |
| const _Float16 * q0, | |
| const _Float16 * q1, | |
| const _Float16 * q2, | |
| const _Float16 * q3, | |
| const _Float16 * k_pack, | |
| int64_t dk, | |
| int64_t kv_tile) { | |
| const size_t vl = __riscv_vsetvl_e16m1(kv_tile); | |
| vfloat32m2_t acc0 = __riscv_vfmv_v_f_f32m2(0.0f, vl); | |
| vfloat32m2_t acc1 = __riscv_vfmv_v_f_f32m2(0.0f, vl); | |
| vfloat32m2_t acc2 = __riscv_vfmv_v_f_f32m2(0.0f, vl); | |
| vfloat32m2_t acc3 = __riscv_vfmv_v_f_f32m2(0.0f, vl); | |
| for (int64_t d = 0; d < dk; ++d) { | |
| const vfloat16m1_t k_vec = __riscv_vle16_v_f16m1(k_pack + d * ggml_fa_tile_config::KV, vl); | |
| acc0 = __riscv_vfwmacc_vf_f32m2(acc0, q0[d], k_vec, vl); | |
| acc1 = __riscv_vfwmacc_vf_f32m2(acc1, q1[d], k_vec, vl); | |
| acc2 = __riscv_vfwmacc_vf_f32m2(acc2, q2[d], k_vec, vl); | |
| acc3 = __riscv_vfwmacc_vf_f32m2(acc3, q3[d], k_vec, vl); | |
| } | |
| __riscv_vse32_v_f32m2(dst0, acc0, vl); | |
| __riscv_vse32_v_f32m2(dst1, acc1, vl); | |
| __riscv_vse32_v_f32m2(dst2, acc2, vl); | |
| __riscv_vse32_v_f32m2(dst3, acc3, vl); | |
| } | |
| static inline void rvv_pv_accumulate_f16_x1(float * dst, | |
| const float * prob, | |
| const _Float16 * v_pack, | |
| int64_t kv_tile, | |
| int64_t dv) { | |
| int64_t d_left = dv; | |
| int64_t d_off = 0; | |
| while (d_left > 0) { | |
| const size_t vl = __riscv_vsetvl_e16m2(d_left); | |
| vfloat32m4_t acc = __riscv_vle32_v_f32m4(dst + d_off, vl); | |
| for (int64_t tk = 0; tk < kv_tile; ++tk) { | |
| const vfloat16m2_t v16 = __riscv_vle16_v_f16m2(v_pack + tk * dv + d_off, vl); | |
| const vfloat32m4_t v32 = __riscv_vfwcvt_f_f_v_f32m4(v16, vl); | |
| acc = __riscv_vfmacc_vf_f32m4(acc, prob[tk], v32, vl); | |
| } | |
| __riscv_vse32_v_f32m4(dst + d_off, acc, vl); | |
| d_left -= vl; | |
| d_off += vl; | |
| } | |
| } | |
| static inline void rvv_pv_accumulate_f16_x4(float * dst0, | |
| float * dst1, | |
| float * dst2, | |
| float * dst3, | |
| const float * prob0, | |
| const float * prob1, | |
| const float * prob2, | |
| const float * prob3, | |
| const _Float16 * v_pack, | |
| int64_t kv_tile, | |
| int64_t dv) { | |
| int64_t d_left = dv; | |
| int64_t d_off = 0; | |
| while (d_left > 0) { | |
| const size_t vl = __riscv_vsetvl_e16m2(d_left); | |
| vfloat32m4_t acc0 = __riscv_vle32_v_f32m4(dst0 + d_off, vl); | |
| vfloat32m4_t acc1 = __riscv_vle32_v_f32m4(dst1 + d_off, vl); | |
| vfloat32m4_t acc2 = __riscv_vle32_v_f32m4(dst2 + d_off, vl); | |
| vfloat32m4_t acc3 = __riscv_vle32_v_f32m4(dst3 + d_off, vl); | |
| for (int64_t tk = 0; tk < kv_tile; ++tk) { | |
| const vfloat16m2_t v16 = __riscv_vle16_v_f16m2(v_pack + tk * dv + d_off, vl); | |
| const vfloat32m4_t v32 = __riscv_vfwcvt_f_f_v_f32m4(v16, vl); | |
| acc0 = __riscv_vfmacc_vf_f32m4(acc0, prob0[tk], v32, vl); | |
| acc1 = __riscv_vfmacc_vf_f32m4(acc1, prob1[tk], v32, vl); | |
| acc2 = __riscv_vfmacc_vf_f32m4(acc2, prob2[tk], v32, vl); | |
| acc3 = __riscv_vfmacc_vf_f32m4(acc3, prob3[tk], v32, vl); | |
| } | |
| __riscv_vse32_v_f32m4(dst0 + d_off, acc0, vl); | |
| __riscv_vse32_v_f32m4(dst1 + d_off, acc1, vl); | |
| __riscv_vse32_v_f32m4(dst2 + d_off, acc2, vl); | |
| __riscv_vse32_v_f32m4(dst3 + d_off, acc3, vl); | |
| d_left -= vl; | |
| d_off += vl; | |
| } | |
| } | |
| static inline void rvv_qk_dot_tile(float * dst, | |
| const float * q_row, | |
| const float * k_pack, | |
| int64_t dk, | |
| int64_t kv_tile, | |
| float scale) { | |
| const size_t vl = __riscv_vsetvl_e32m4(kv_tile); | |
| vfloat32m4_t acc = __riscv_vfmv_v_f_f32m4(0.0f, vl); | |
| for (int64_t d = 0; d < dk; ++d) { | |
| const vfloat32m4_t k_vec = __riscv_vle32_v_f32m4(k_pack + d * kv_tile, vl); | |
| acc = __riscv_vfmacc_vf_f32m4(acc, q_row[d] * scale, k_vec, vl); | |
| } | |
| __riscv_vse32_v_f32m4(dst, acc, vl); | |
| } | |
| static inline void rvv_pv_accumulate(float * dst, | |
| const float * prob, | |
| const float * v_pack, | |
| int64_t kv_tile, | |
| int64_t dv) { | |
| int64_t d_left = dv; | |
| int64_t d_off = 0; | |
| while (d_left > 0) { | |
| const size_t vl = __riscv_vsetvl_e32m4(d_left); | |
| vfloat32m4_t acc = __riscv_vle32_v_f32m4(dst + d_off, vl); | |
| for (int64_t tk = 0; tk < kv_tile; ++tk) { | |
| const vfloat32m4_t v_vec = __riscv_vle32_v_f32m4(v_pack + tk * dv + d_off, vl); | |
| acc = __riscv_vfmacc_vf_f32m4(acc, prob[tk], v_vec, vl); | |
| } | |
| __riscv_vse32_v_f32m4(dst + d_off, acc, vl); | |
| d_left -= vl; | |
| d_off += vl; | |
| } | |
| } | |
| static void permute_transpose_impl(const ggml_tensor * src0, | |
| ggml_tensor * dst, | |
| int64_t batch, | |
| int64_t m, | |
| int64_t n, | |
| int64_t batch_stride, | |
| int64_t m_src_stride, | |
| int64_t n_src_stride, | |
| int64_t n_dst_stride, | |
| int ith, | |
| int nth) { | |
| GGML_ASSERT(n_src_stride == sizeof(int32_t) || n_src_stride == sizeof(int16_t)); | |
| if (n_src_stride == sizeof(int32_t)) { | |
| for (int64_t bi = ith; bi < batch; bi += nth) { | |
| rvv_transposed_s32_mn_to_nm((int8_t *) ((char *) dst->data + bi * batch_stride), n_dst_stride, | |
| (int8_t *) ((char *) src0->data + bi * batch_stride), m_src_stride, m, n); | |
| } | |
| } else if (n_src_stride == sizeof(int16_t)) { | |
| for (int64_t bi = ith; bi < batch; bi += nth) { | |
| rvv_transposed_s32_mn_to_nm((int8_t *) ((char *) dst->data + bi * batch_stride), n_dst_stride, | |
| (int8_t *) ((char *) src0->data + bi * batch_stride), m_src_stride, m, n); | |
| } | |
| } else { | |
| GGML_ABORT("not implemented"); | |
| } | |
| } | |
| template <size_t QLEN> | |
| static void flash_attn_ext_f16_one_chunk_inner_vlen1024_vf16_mrow(float ** pq, | |
| const char * k_data_row, | |
| const char * v_data_row, | |
| const ggml_fp16_t * mp, | |
| float ** sinks, | |
| float ** dst, | |
| float scale, | |
| float logit_softcap, | |
| float slope, | |
| int64_t nek1, | |
| int64_t nbk1, | |
| int64_t nbv1, | |
| int64_t DV, | |
| int64_t DK, | |
| void * tcm_buffer, | |
| size_t tcm_buffer_size) { | |
| GGML_ASSERT(flash_attn_ext_supported_shape_vlen1024_vf16(DK, DV)); | |
| float S[QLEN] = { 0.0f }; // sum | |
| float M[QLEN] = { -INFINITY }; // maximum KQ value | |
| _Float16 * kq16_buffer = (_Float16 *) tcm_buffer; | |
| _Float16 * qv_buffer = kq16_buffer + QLEN * DV; | |
| const size_t qkv_temp_buffer_size = (QLEN * DV + QLEN * DK) * sizeof(_Float16); | |
| char * kv_tile_buffer = (char *) (qv_buffer + QLEN * DK); | |
| { | |
| vfloat16m2_t VKQ16_v = __riscv_vfmv_v_f_f16m2(0.0f, DV); | |
| for (int64_t i = 0; i < QLEN; ++i) { | |
| __riscv_vse16_v_f16m2(kq16_buffer + i * DV, VKQ16_v, DV); | |
| vfloat16m2_t Q_q_v = __riscv_vfncvt_f_f_w_f16m2(__riscv_vle32_v_f32m4(pq[i], DK), DK); | |
| __riscv_vse16_v_f16m2(qv_buffer + i * DK, Q_q_v, DK); | |
| } | |
| } | |
| const uintptr_t scratch_addr = reinterpret_cast<uintptr_t>(kv_tile_buffer); | |
| const size_t scratch_size = tcm_buffer_size > qkv_temp_buffer_size ? tcm_buffer_size - qkv_temp_buffer_size : 0; | |
| const uintptr_t kq_tile_addr = align_up(scratch_addr, alignof(float)); | |
| const size_t scratch_prefix = kq_tile_addr - scratch_addr; | |
| const size_t packed_tile_size = | |
| QLEN * sizeof(float) + DK * sizeof(_Float16) + DV * sizeof(_Float16) + sizeof(float); | |
| const int64_t max_ic_tile_step = ((int64_t) __riscv_vsetvlmax_e16m1()) & ~((int64_t) 7); | |
| const int64_t max_fit_by_tcm = | |
| scratch_size > scratch_prefix ? (int64_t) ((scratch_size - scratch_prefix) / packed_tile_size) : 0; | |
| const int64_t ic_tile_step = std::min(max_ic_tile_step, max_fit_by_tcm) & ~((int64_t) 7); | |
| const uintptr_t k_tile_addr = kq_tile_addr + QLEN * ic_tile_step * sizeof(float); | |
| const uintptr_t v_tile_addr = k_tile_addr + DK * ic_tile_step * sizeof(_Float16); | |
| const uintptr_t mv_tile_addr = v_tile_addr + ic_tile_step * DV * sizeof(_Float16); | |
| if (ic_tile_step >= 8) { | |
| float * kq_tile_buffer = reinterpret_cast<float *>(kq_tile_addr); | |
| _Float16 * k_tile_pack = reinterpret_cast<_Float16 *>(k_tile_addr); | |
| _Float16 * v_tile_pack = reinterpret_cast<_Float16 *>(v_tile_addr); | |
| float * mv_tile_pack = reinterpret_cast<float *>(mv_tile_addr); | |
| const int64_t k_tile_byte_stride = ic_tile_step * (int64_t) sizeof(_Float16); | |
| int64_t ic_step = 0; | |
| for (int64_t ic = 0; ic < nek1; ++ic) { | |
| const float mv = mp ? slope * ((_Float16 *) mp)[ic] : 0.0f; | |
| if (mv != -INFINITY) { | |
| const _Float16 * k_data = (const _Float16 *) (k_data_row + ic * nbk1); | |
| const _Float16 * v_data = (const _Float16 *) (v_data_row + ic * nbv1); | |
| const vfloat16m2_t k_data_v = __riscv_vle16_v_f16m2(k_data, DK); | |
| const vfloat16m2_t v_data_v = __riscv_vle16_v_f16m2(v_data, DV); | |
| __riscv_vsse16_v_f16m2(k_tile_pack + ic_step, k_tile_byte_stride, k_data_v, DK); | |
| __riscv_vse16_v_f16m2(v_tile_pack + ic_step * DV, v_data_v, DV); | |
| mv_tile_pack[ic_step] = mv; | |
| ic_step++; | |
| } | |
| if (ic_step > 0 && (ic_step == ic_tile_step || ic == (nek1 - 1))) { | |
| if constexpr (QLEN == 4) { | |
| const size_t qk_vl = __riscv_vsetvl_e16m1(ic_step); | |
| vfloat32m2_t qk_acc0 = __riscv_vfmv_v_f_f32m2(0.0f, qk_vl); | |
| vfloat32m2_t qk_acc1 = __riscv_vfmv_v_f_f32m2(0.0f, qk_vl); | |
| vfloat32m2_t qk_acc2 = __riscv_vfmv_v_f_f32m2(0.0f, qk_vl); | |
| vfloat32m2_t qk_acc3 = __riscv_vfmv_v_f_f32m2(0.0f, qk_vl); | |
| for (int64_t d = 0; d < DK; ++d) { | |
| const vfloat16m1_t k_vec = __riscv_vle16_v_f16m1(k_tile_pack + d * ic_tile_step, qk_vl); | |
| qk_acc0 = __riscv_vfwmacc_vf_f32m2(qk_acc0, qv_buffer[0 * DK + d], k_vec, qk_vl); | |
| qk_acc1 = __riscv_vfwmacc_vf_f32m2(qk_acc1, qv_buffer[1 * DK + d], k_vec, qk_vl); | |
| qk_acc2 = __riscv_vfwmacc_vf_f32m2(qk_acc2, qv_buffer[2 * DK + d], k_vec, qk_vl); | |
| qk_acc3 = __riscv_vfwmacc_vf_f32m2(qk_acc3, qv_buffer[3 * DK + d], k_vec, qk_vl); | |
| } | |
| qk_acc0 = __riscv_vfmul_vf_f32m2(qk_acc0, scale, qk_vl); | |
| qk_acc1 = __riscv_vfmul_vf_f32m2(qk_acc1, scale, qk_vl); | |
| qk_acc2 = __riscv_vfmul_vf_f32m2(qk_acc2, scale, qk_vl); | |
| qk_acc3 = __riscv_vfmul_vf_f32m2(qk_acc3, scale, qk_vl); | |
| __riscv_vse32_v_f32m2(kq_tile_buffer + 0 * ic_tile_step, qk_acc0, qk_vl); | |
| __riscv_vse32_v_f32m2(kq_tile_buffer + 1 * ic_tile_step, qk_acc1, qk_vl); | |
| __riscv_vse32_v_f32m2(kq_tile_buffer + 2 * ic_tile_step, qk_acc2, qk_vl); | |
| __riscv_vse32_v_f32m2(kq_tile_buffer + 3 * ic_tile_step, qk_acc3, qk_vl); | |
| } else { | |
| static_assert(QLEN == 2, "unsupported QLEN"); | |
| const size_t qk_vl = __riscv_vsetvl_e16m1(ic_step); | |
| vfloat32m2_t qk_acc0 = __riscv_vfmv_v_f_f32m2(0.0f, qk_vl); | |
| vfloat32m2_t qk_acc1 = __riscv_vfmv_v_f_f32m2(0.0f, qk_vl); | |
| for (int64_t d = 0; d < DK; ++d) { | |
| const vfloat16m1_t k_vec = __riscv_vle16_v_f16m1(k_tile_pack + d * ic_tile_step, qk_vl); | |
| qk_acc0 = __riscv_vfwmacc_vf_f32m2(qk_acc0, qv_buffer[0 * DK + d], k_vec, qk_vl); | |
| qk_acc1 = __riscv_vfwmacc_vf_f32m2(qk_acc1, qv_buffer[1 * DK + d], k_vec, qk_vl); | |
| } | |
| qk_acc0 = __riscv_vfmul_vf_f32m2(qk_acc0, scale, qk_vl); | |
| qk_acc1 = __riscv_vfmul_vf_f32m2(qk_acc1, scale, qk_vl); | |
| __riscv_vse32_v_f32m2(kq_tile_buffer + 0 * ic_tile_step, qk_acc0, qk_vl); | |
| __riscv_vse32_v_f32m2(kq_tile_buffer + 1 * ic_tile_step, qk_acc1, qk_vl); | |
| } | |
| for (int i = 0; i < QLEN; ++i) { | |
| float * row_ptr = kq_tile_buffer + i * ic_tile_step; | |
| const float tile_max = | |
| rvv_softcap_add_max_inplace_f32(row_ptr, mv_tile_pack, ic_step, logit_softcap); | |
| const float Mold = M[i]; | |
| if (tile_max > Mold) { | |
| const float ms = expf(Mold - tile_max); | |
| M[i] = tile_max; | |
| S[i] *= ms; | |
| vfloat16m2_t VKQ16_v = __riscv_vle16_v_f16m2(kq16_buffer + i * DV, DV); | |
| VKQ16_v = __riscv_vfmul_vf_f16m2(VKQ16_v, (_Float16) ms, DV); | |
| __riscv_vse16_v_f16m2(kq16_buffer + i * DV, VKQ16_v, DV); | |
| } | |
| S[i] += rvv_softmax_exp_inplace_f32(row_ptr, ic_step, M[i]); | |
| } | |
| if constexpr (QLEN == 4) { | |
| vfloat16m2_t pv_acc0 = __riscv_vle16_v_f16m2(kq16_buffer + 0 * DV, DV); | |
| vfloat16m2_t pv_acc1 = __riscv_vle16_v_f16m2(kq16_buffer + 1 * DV, DV); | |
| vfloat16m2_t pv_acc2 = __riscv_vle16_v_f16m2(kq16_buffer + 2 * DV, DV); | |
| vfloat16m2_t pv_acc3 = __riscv_vle16_v_f16m2(kq16_buffer + 3 * DV, DV); | |
| for (int64_t tk = 0; tk < ic_step; ++tk) { | |
| const vfloat16m2_t v16 = __riscv_vle16_v_f16m2(v_tile_pack + tk * DV, DV); | |
| pv_acc0 = | |
| __riscv_vfmacc_vf_f16m2(pv_acc0, (_Float16) kq_tile_buffer[0 * ic_tile_step + tk], v16, DV); | |
| pv_acc1 = | |
| __riscv_vfmacc_vf_f16m2(pv_acc1, (_Float16) kq_tile_buffer[1 * ic_tile_step + tk], v16, DV); | |
| pv_acc2 = | |
| __riscv_vfmacc_vf_f16m2(pv_acc2, (_Float16) kq_tile_buffer[2 * ic_tile_step + tk], v16, DV); | |
| pv_acc3 = | |
| __riscv_vfmacc_vf_f16m2(pv_acc3, (_Float16) kq_tile_buffer[3 * ic_tile_step + tk], v16, DV); | |
| } | |
| __riscv_vse16_v_f16m2(kq16_buffer + 0 * DV, pv_acc0, DV); | |
| __riscv_vse16_v_f16m2(kq16_buffer + 1 * DV, pv_acc1, DV); | |
| __riscv_vse16_v_f16m2(kq16_buffer + 2 * DV, pv_acc2, DV); | |
| __riscv_vse16_v_f16m2(kq16_buffer + 3 * DV, pv_acc3, DV); | |
| } else { | |
| static_assert(QLEN == 2, "unsupported QLEN"); | |
| vfloat16m2_t pv_acc0 = __riscv_vle16_v_f16m2(kq16_buffer + 0 * DV, DV); | |
| vfloat16m2_t pv_acc1 = __riscv_vle16_v_f16m2(kq16_buffer + 1 * DV, DV); | |
| for (int64_t tk = 0; tk < ic_step; ++tk) { | |
| const vfloat16m2_t v16 = __riscv_vle16_v_f16m2(v_tile_pack + tk * DV, DV); | |
| pv_acc0 = | |
| __riscv_vfmacc_vf_f16m2(pv_acc0, (_Float16) kq_tile_buffer[0 * ic_tile_step + tk], v16, DV); | |
| pv_acc1 = | |
| __riscv_vfmacc_vf_f16m2(pv_acc1, (_Float16) kq_tile_buffer[1 * ic_tile_step + tk], v16, DV); | |
| } | |
| __riscv_vse16_v_f16m2(kq16_buffer + 0 * DV, pv_acc0, DV); | |
| __riscv_vse16_v_f16m2(kq16_buffer + 1 * DV, pv_acc1, DV); | |
| } | |
| ic_step = 0; | |
| } | |
| } | |
| } else { | |
| for (int64_t ic = 0; ic < nek1; ++ic) { | |
| const float mv = mp ? slope * ((_Float16 *) mp)[ic] : 0.0f; | |
| const char * k_data = k_data_row + ic * nbk1; | |
| const char * v_data = v_data_row + ic * nbv1; | |
| vfloat16m2_t k_data_v; | |
| vfloat16m2_t v_data_v; | |
| if (mv != -INFINITY) { | |
| k_data_v = __riscv_vle16_v_f16m2((_Float16 *) k_data, DK); | |
| v_data_v = __riscv_vle16_v_f16m2((_Float16 *) v_data, DV); | |
| } else { | |
| continue; | |
| } | |
| for (int i = 0; i < QLEN; ++i) { | |
| vfloat16m2_t Q_q_v = __riscv_vle16_v_f16m2(qv_buffer + i * DK, DK); | |
| vfloat32m4_t qk_acc_v = __riscv_vfwmul_vv_f32m4(k_data_v, Q_q_v, DK); | |
| float s = reduce_sum_f32m4_vlen1024(qk_acc_v, DK); | |
| s = s * scale; | |
| if (logit_softcap != 0.0f) { | |
| s = logit_softcap * tanhf(s); | |
| } | |
| s += mv; | |
| const float Mold = M[i]; | |
| float ms = 1.0f; // upon new higher max val, scale VKQ and KQ sum with this value | |
| float vs = 1.0f; // post-softmax KQ value, expf(s - M) | |
| vfloat16m2_t VKQ16_v = __riscv_vle16_v_f16m2(kq16_buffer + i * DV, DV); | |
| if (s > M[i]) { | |
| // s is new maximum, ms < 1.0f, vs == expf(s - s) == 1.0f | |
| M[i] = s; | |
| ms = expf(Mold - M[i]); | |
| // V = V*expf(Mold - M) | |
| VKQ16_v = __riscv_vfmul_vf_f16m2(VKQ16_v, ms, DV); | |
| } else { | |
| // no new maximum, ms == 1.0f, vs != 1.0f | |
| vs = expf(s - M[i]); | |
| } | |
| VKQ16_v = __riscv_vfmacc_vf_f16m2(VKQ16_v, vs, v_data_v, DV); | |
| __riscv_vse16_v_f16m2(kq16_buffer + i * DV, VKQ16_v, DV); | |
| S[i] = S[i] * ms + vs; // scale and increment sum with partial sum | |
| } | |
| } | |
| } | |
| for (int i = 0; i < QLEN; ++i) { | |
| vfloat16m2_t VKQ16_v = __riscv_vle16_v_f16m2(kq16_buffer + i * DV, DV); | |
| vfloat32m4_t VKQ32_v = __riscv_vfwcvt_f_f_v_f32m4(VKQ16_v, DV); | |
| // sinks | |
| if (sinks[i]) { | |
| const float s = *(sinks[i]); | |
| float ms = 1.0f; | |
| float vs = 1.0f; | |
| if (s > M[i]) { | |
| ms = expf(M[i] - s); | |
| M[i] = s; | |
| VKQ32_v = __riscv_vfmul_vf_f32m4(VKQ32_v, ms, DV); | |
| } else { | |
| vs = expf(s - M[i]); | |
| } | |
| S[i] = S[i] * ms + vs; | |
| } | |
| // V /= S | |
| const float S_inv = S[i] == 0.0f ? 0.0f : 1.0f / S[i]; | |
| VKQ32_v = __riscv_vfmul_vf_f32m4(VKQ32_v, S_inv, DV); | |
| __riscv_vse32_v_f32m4(dst[i], VKQ32_v, DV); | |
| } | |
| } | |
| static void flash_attn_ext_f16_one_chunk_inner_vlen1024_vf16_m1(const float * pq, | |
| const char * k_data_row, | |
| const char * v_data_row, | |
| const ggml_fp16_t * mp, | |
| const float * sinks, | |
| float * dst, | |
| float scale, | |
| float logit_softcap, | |
| float slope, | |
| int64_t nek1, | |
| int64_t nbk1, | |
| int64_t nbv1, | |
| int64_t DV, | |
| int64_t DK) { | |
| GGML_ASSERT(flash_attn_ext_supported_shape_vlen1024_vf16(DK, DV)); | |
| float S = 0.0f; // sum | |
| float M = -INFINITY; // maximum KQ value | |
| vfloat16m2_t VKQ16_v = __riscv_vfmv_v_f_f16m2(0.0f, DV); | |
| vfloat16m2_t Q_q_v = __riscv_vfncvt_f_f_w_f16m2(__riscv_vle32_v_f32m4(pq, DK), DK); | |
| for (int64_t ic = 0; ic < nek1; ++ic) { | |
| const float mv = mp ? slope * ((_Float16 *) mp)[ic] : 0.0f; | |
| if (mv == -INFINITY) { | |
| continue; | |
| } | |
| const char * k_data = k_data_row + ic * nbk1; | |
| vfloat16m2_t k_data_v = __riscv_vle16_v_f16m2((_Float16 *) k_data, DK); | |
| vfloat32m4_t qk_acc_v = __riscv_vfwmul_vv_f32m4(k_data_v, Q_q_v, DK); | |
| float s = reduce_sum_f32m4_vlen1024(qk_acc_v, DK); | |
| s = s * scale; // scale KQ value | |
| if (logit_softcap != 0.0f) { | |
| s = logit_softcap * tanhf(s); | |
| } | |
| s += mv; // apply mask | |
| const float Mold = M; | |
| float ms = 1.0f; // upon new higher max val, scale VKQ and KQ sum with this value | |
| float vs = 1.0f; // post-softmax KQ value, expf(s - M) | |
| const char * v_data = v_data_row + ic * nbv1; | |
| vfloat16m2_t v_data_v = __riscv_vle16_v_f16m2((_Float16 *) v_data, DV); | |
| if (s > M) { | |
| // s is new maximum, ms < 1.0f, vs == expf(s - s) == 1.0f | |
| M = s; | |
| ms = expf(Mold - M); | |
| // V = V*expf(Mold - M) | |
| VKQ16_v = __riscv_vfmul_vf_f16m2(VKQ16_v, ms, DV); | |
| } else { | |
| // no new maximum, ms == 1.0f, vs != 1.0f | |
| vs = expf(s - M); | |
| } | |
| VKQ16_v = __riscv_vfmacc_vf_f16m2(VKQ16_v, vs, v_data_v, DV); | |
| S = S * ms + vs; // scale and increment sum with partial sum | |
| } | |
| vfloat32m4_t VKQ32_v = __riscv_vfwcvt_f_f_v_f32m4(VKQ16_v, DV); | |
| // sinks | |
| if (sinks) { | |
| const float s = *sinks; | |
| float ms = 1.0f; | |
| float vs = 1.0f; | |
| if (s > M) { | |
| ms = expf(M - s); | |
| M = s; | |
| VKQ32_v = __riscv_vfmul_vf_f32m4(VKQ32_v, ms, DV); | |
| } else { | |
| vs = expf(s - M); | |
| } | |
| S = S * ms + vs; | |
| } | |
| // V /= S | |
| const float S_inv = S == 0.0f ? 0.0f : 1.0f / S; | |
| VKQ32_v = __riscv_vfmul_vf_f32m4(VKQ32_v, S_inv, DV); | |
| __riscv_vse32_v_f32m4(dst, VKQ32_v, DV); | |
| } | |
| } // namespace | |
| void memcpy1d(void * dst, const void * src, int64_t size) { | |
| size_t byte_size_all = size; | |
| size_t vlen = __riscv_vlenb() * 8; | |
| if (vlen == 256) { | |
| // 1024 bytes | |
| __asm__ volatile( | |
| // | |
| "srli t0, %[size], 10 \n\t" | |
| "blez t0, memcpy_tail%= \n\t" | |
| "vsetvli t1, x0, e8, m8, tu, mu \n\t" | |
| "memcpy_main_loop%=: \n\t" | |
| "addi t0, t0, -1 \n\t" | |
| "vle8.v v0, (%[s]) \n\t" | |
| "addi %[s], %[s], 256 \n\t" | |
| "vle8.v v8, (%[s]) \n\t" | |
| "addi %[s], %[s], 256 \n\t" | |
| "vle8.v v16, (%[s]) \n\t" | |
| "addi %[s], %[s], 256 \n\t" | |
| "vle8.v v24, (%[s]) \n\t" | |
| "addi %[s], %[s], 256 \n\t" | |
| // | |
| "vse8.v v0, (%[d]) \n\t" | |
| "addi %[d], %[d], 256 \n\t" | |
| "vse8.v v8, (%[d]) \n\t" | |
| "addi %[d], %[d], 256 \n\t" | |
| "vse8.v v16, (%[d]) \n\t" | |
| "addi %[d], %[d], 256 \n\t" | |
| "vse8.v v24, (%[d]) \n\t" | |
| "addi %[d], %[d], 256 \n\t" | |
| // | |
| "bnez t0, memcpy_main_loop%= \n\t" | |
| "memcpy_tail%=: \n\t" | |
| "andi t1, %[size], 1023 \n\t" | |
| "blez t1, out%= \n\t" | |
| "memcpy_tail_loop%=: \n\t" | |
| "vsetvli t0, t1, e8, m8, tu, mu \n\t" | |
| "sub t1, t1, t0 \n\t" | |
| "vle8.v v0, (%[s]) \n\t" | |
| "add %[s], %[s], t0 \n\t" | |
| "vse8.v v0, (%[d]) \n\t" | |
| "add %[d], %[d], t0 \n\t" | |
| "bnez t1, memcpy_tail_loop%= \n\t" | |
| "out%=: \n\t" | |
| : [s] "+r"(src), [d] "+r"(dst) | |
| : [size] "r"(byte_size_all) | |
| : "cc", "t0", "t1"); | |
| } else if (vlen == 1024) { | |
| // 2048 bytes | |
| __asm__ volatile( | |
| // | |
| "srli t0, %[size], 11 \n\t" | |
| "blez t0, memcpy_tail%= \n\t" | |
| "vsetvli t1, x0, e8, m8, tu, mu \n\t" | |
| "addi t2, %[s], 1024 \n\t" | |
| "addi t3, %[d], 1024 \n\t" | |
| "li t5, 2048 \n\t" | |
| "memcpy_main_loop%=: \n\t" | |
| "addi t0, t0, -1 \n\t" | |
| "vle8.v v0, (%[s]) \n\t" | |
| "add %[s], %[s], t5 \n\t" | |
| "vle8.v v8, (t2) \n\t" | |
| "add t2, t2, t5 \n\t" | |
| // | |
| "vse8.v v0, (%[d]) \n\t" | |
| "add %[d], %[d], t5 \n\t" | |
| "vse8.v v8, (t3) \n\t" | |
| "add t3, t3, t5 \n\t" | |
| // | |
| "bnez t0, memcpy_main_loop%= \n\t" | |
| "memcpy_tail%=: \n\t" | |
| "andi t1, %[size], 2047 \n\t" | |
| "blez t1, out%= \n\t" | |
| "memcpy_tail_loop%=: \n\t" | |
| "vsetvli t0, t1, e8, m2, tu, mu \n\t" | |
| "sub t1, t1, t0 \n\t" | |
| "vle8.v v0, (%[s]) \n\t" | |
| "add %[s], %[s], t0 \n\t" | |
| "vse8.v v0, (%[d]) \n\t" | |
| "add %[d], %[d], t0 \n\t" | |
| "bnez t1, memcpy_tail_loop%= \n\t" | |
| "out%=: \n\t" | |
| : [s] "+r"(src), [d] "+r"(dst) | |
| : [size] "r"(byte_size_all) | |
| : "cc", "t0", "t1", "t2", "t3", "t5"); | |
| } else { | |
| __asm__ volatile( | |
| // | |
| "add t1, %[size], zero \n\t" | |
| "memcpy_tail_loop%=: \n\t" | |
| "vsetvli t0, t1, e8, m8, tu, mu \n\t" | |
| "sub t1, t1, t0 \n\t" | |
| "vle8.v v0, (%[s]) \n\t" | |
| "add %[s], %[s], t0 \n\t" | |
| "vse8.v v0, (%[d]) \n\t" | |
| "add %[d], %[d], t0 \n\t" | |
| "bnez t1, memcpy_tail_loop%= \n\t" | |
| : [s] "+r"(src), [d] "+r"(dst) | |
| : [size] "r"(byte_size_all) | |
| : "cc", "t0", "t1", "t2", "t4", "t3"); | |
| } | |
| } | |
| void memcpy2d(void * dst, int64_t dst_stride, const void * src, int64_t src_stride, int64_t tile_rows, int64_t size) { | |
| for (int64_t i = 0; i < tile_rows; ++i) { | |
| memcpy1d((char *) dst + i * dst_stride, (const char *) src + i * src_stride, size); | |
| } | |
| } | |
| void forward_flash_attn_ext_f16_one_chunk_vlen1024_vf16(const ggml_compute_params * params, | |
| ggml_tensor * dst, | |
| int ir0, | |
| int ir1, | |
| void * tcm_buffer, | |
| size_t tcm_buffer_size) { | |
| const ggml_tensor * q = dst->src[0]; | |
| const ggml_tensor * k = dst->src[1]; | |
| const ggml_tensor * v = dst->src[2]; | |
| const ggml_tensor * mask = dst->src[3]; | |
| const ggml_tensor * sinks = dst->src[4]; | |
| GGML_TENSOR_LOCALS(int64_t, neq, q, ne) | |
| GGML_TENSOR_LOCALS(size_t, nbq, q, nb) | |
| GGML_TENSOR_LOCALS(int64_t, nek, k, ne) | |
| GGML_TENSOR_LOCALS(size_t, nbk, k, nb) | |
| GGML_TENSOR_LOCALS(int64_t, nev, v, ne) | |
| GGML_TENSOR_LOCALS(size_t, nbv, v, nb) | |
| GGML_TENSOR_LOCALS(int64_t, ne, dst, ne) | |
| GGML_TENSOR_LOCALS(size_t, nb, dst, nb) | |
| const int64_t DK = nek0; | |
| const int64_t DV = nev0; | |
| const int64_t N = neq1; | |
| GGML_ASSERT(flash_attn_ext_supported_shape_vlen1024_vf16(DK, DV)); | |
| // broadcast factors | |
| const int64_t rk2 = neq2 / nek2; | |
| const int64_t rk3 = neq3 / nek3; | |
| const int64_t rv2 = neq2 / nev2; | |
| const int64_t rv3 = neq3 / nev3; | |
| // parallelize by q rows using ggml_vec_dot_f32 | |
| float scale = *((float *) dst->op_params + 0); | |
| float max_bias = *((float *) dst->op_params + 1); | |
| float logit_softcap = *((float *) dst->op_params + 2); | |
| if (logit_softcap != 0) { | |
| scale /= logit_softcap; | |
| } | |
| const uint32_t n_head = neq2; | |
| const uint32_t n_head_log2 = 1u << (uint32_t) floor(log2(n_head)); | |
| const float m0 = powf(2.0f, -(max_bias) / n_head_log2); | |
| const float m1 = powf(2.0f, -(max_bias / 2.0f) / n_head_log2); | |
| const int KV_row_size = DK * sizeof(_Float16) + DV * sizeof(_Float16); | |
| int ith = params->ith; | |
| int ir_step = 1; | |
| for (int ir = ir0; ir < ir1; ir += ir_step) { | |
| // q indices | |
| const int iq3 = ir / (neq2 * neq1); | |
| const int iq2 = (ir - iq3 * neq2 * neq1) / neq1; | |
| const int iq1 = (ir - iq3 * neq2 * neq1 - iq2 * neq1); | |
| const int iq3_1 = (ir + 1) / (neq2 * neq1); | |
| const int iq2_1 = (ir + 1 - iq3_1 * neq2 * neq1) / neq1; | |
| const int iq1_1 = (ir + 1 - iq3_1 * neq2 * neq1 - iq2_1 * neq1); | |
| const int iq3_2 = (ir + 2) / (neq2 * neq1); | |
| const int iq2_2 = (ir + 2 - iq3_2 * neq2 * neq1) / neq1; | |
| const int iq1_2 = (ir + 2 - iq3_2 * neq2 * neq1 - iq2_2 * neq1); | |
| const int iq3_3 = (ir + 3) / (neq2 * neq1); | |
| const int iq2_3 = (ir + 3 - iq3_3 * neq2 * neq1) / neq1; | |
| const int iq1_3 = (ir + 3 - iq3_3 * neq2 * neq1 - iq2_3 * neq1); | |
| const uint32_t h = iq2; // head index | |
| const float slope = | |
| (max_bias > 0.0f) ? h < n_head_log2 ? powf(m0, h + 1) : powf(m1, 2 * (h - n_head_log2) + 1) : 1.0f; | |
| const ggml_fp16_t * mp = | |
| mask ? (ggml_fp16_t *) ((char *) mask->data + iq1 * mask->nb[1] + (iq2 % mask->ne[2]) * mask->nb[2] + | |
| (iq3 % mask->ne[3]) * mask->nb[3]) : | |
| NULL; | |
| const bool mp_equal_2 = iq1_1 == iq1 && (iq2 % mask->ne[2]) == (iq2_1 % mask->ne[2]) && | |
| (iq3 % mask->ne[3]) == (iq3_1 % mask->ne[3]); | |
| const bool mp_equal_4 = mp_equal_2 && iq1_2 == iq1 && (iq2 % mask->ne[2]) == (iq2_2 % mask->ne[2]) && | |
| (iq3 % mask->ne[3]) == (iq3_2 % mask->ne[3]) && iq1_3 == iq1 && | |
| (iq2 % mask->ne[2]) == (iq2_3 % mask->ne[2]) && | |
| (iq3 % mask->ne[3]) == (iq3_3 % mask->ne[3]); | |
| // k indices | |
| const int ik3 = iq3 / rk3; | |
| const int ik2 = iq2 / rk2; | |
| const int ik3_1 = iq3_1 / rk3; | |
| const int ik2_1 = iq2_1 / rk2; | |
| const int ik3_2 = iq3_2 / rk3; | |
| const int ik2_2 = iq2_2 / rk2; | |
| const int ik3_3 = iq3_3 / rk3; | |
| const int ik2_3 = iq2_3 / rk2; | |
| // v indices | |
| const int iv3 = iq3 / rv3; | |
| const int iv2 = iq2 / rv2; | |
| const int iv3_1 = iq3_1 / rv3; | |
| const int iv2_1 = iq2_1 / rv2; | |
| const int iv3_2 = iq3_2 / rv3; | |
| const int iv2_2 = iq2_2 / rv2; | |
| const int iv3_3 = iq3_3 / rv3; | |
| const int iv2_3 = iq2_3 / rv2; | |
| const float * pq = (const float *) ((char *) q->data + (iq1 * nbq1 + iq2 * nbq2 + iq3 * nbq3)); | |
| std::array<float *, 4> pq_buffer; | |
| std::array<float *, 4> sinks_buffer; | |
| std::array<float *, 4> dst_buffer; | |
| if (tcm_buffer != nullptr && 4 * KV_row_size < tcm_buffer_size && ir < (ir1 - 3) && mp_equal_4 && | |
| ik3_3 == ik3 && ik2_3 == ik2 && iv3_3 == iv3 && iv2_3 == iv2 && ik3_2 == ik3 && ik2_2 == ik2 && | |
| iv3_2 == iv3 && iv2_2 == iv2 && ik3_1 == ik3 && ik2_1 == ik2 && iv3_1 == iv3 && iv2_1 == iv2) { | |
| ir_step = 4; | |
| pq_buffer[0] = (float *) ((char *) q->data + (iq1 * nbq1 + iq2 * nbq2 + iq3 * nbq3)); | |
| pq_buffer[1] = (float *) ((char *) q->data + (iq1_1 * nbq1 + iq2_1 * nbq2 + iq3_1 * nbq3)); | |
| pq_buffer[2] = (float *) ((char *) q->data + (iq1_2 * nbq1 + iq2_2 * nbq2 + iq3_2 * nbq3)); | |
| pq_buffer[3] = (float *) ((char *) q->data + (iq1_3 * nbq1 + iq2_3 * nbq2 + iq3_3 * nbq3)); | |
| sinks_buffer[0] = sinks ? ((float *) ((char *) sinks->data)) + iq2 : nullptr; | |
| sinks_buffer[1] = sinks ? ((float *) ((char *) sinks->data)) + iq2_1 : nullptr; | |
| sinks_buffer[2] = sinks ? ((float *) ((char *) sinks->data)) + iq2_2 : nullptr; | |
| sinks_buffer[3] = sinks ? ((float *) ((char *) sinks->data)) + iq2_3 : nullptr; | |
| dst_buffer[0] = (float *) ((char *) dst->data + (iq3 * ne2 * ne1 + iq2 + iq1 * ne1) * nb1); | |
| dst_buffer[1] = (float *) ((char *) dst->data + (iq3_1 * ne2 * ne1 + iq2_1 + iq1_1 * ne1) * nb1); | |
| dst_buffer[2] = (float *) ((char *) dst->data + (iq3_2 * ne2 * ne1 + iq2_2 + iq1_2 * ne1) * nb1); | |
| dst_buffer[3] = (float *) ((char *) dst->data + (iq3_3 * ne2 * ne1 + iq2_3 + iq1_3 * ne1) * nb1); | |
| flash_attn_ext_f16_one_chunk_inner_vlen1024_vf16_mrow<4>( // | |
| pq_buffer.data(), // | |
| (const char *) k->data + (ik2 * nbk2 + ik3 * nbk3), // | |
| (const char *) v->data + (iv2 * nbv2 + iv3 * nbv3), // | |
| mp, // | |
| sinks_buffer.data(), // | |
| dst_buffer.data(), // | |
| scale, logit_softcap, slope, nek1, nbk1, nbv1, DV, DK, tcm_buffer, tcm_buffer_size); | |
| } else if (tcm_buffer != nullptr && 2 * KV_row_size < tcm_buffer_size && ir < (ir1 - 1) && mp_equal_2 && | |
| ik3_1 == ik3 && ik2_1 == ik2 && iv3_1 == iv3 && iv2_1 == iv2) { | |
| ir_step = 2; | |
| pq_buffer[0] = (float *) ((char *) q->data + (iq1 * nbq1 + iq2 * nbq2 + iq3 * nbq3)); | |
| pq_buffer[1] = (float *) ((char *) q->data + (iq1_1 * nbq1 + iq2_1 * nbq2 + iq3_1 * nbq3)); | |
| sinks_buffer[0] = sinks ? ((float *) ((char *) sinks->data)) + iq2 : nullptr; | |
| sinks_buffer[1] = sinks ? ((float *) ((char *) sinks->data)) + iq2_1 : nullptr; | |
| dst_buffer[0] = (float *) ((char *) dst->data + (iq3 * ne2 * ne1 + iq2 + iq1 * ne1) * nb1); | |
| dst_buffer[1] = (float *) ((char *) dst->data + (iq3_1 * ne2 * ne1 + iq2_1 + iq1_1 * ne1) * nb1); | |
| flash_attn_ext_f16_one_chunk_inner_vlen1024_vf16_mrow<2>( // | |
| pq_buffer.data(), // | |
| (const char *) k->data + (ik2 * nbk2 + ik3 * nbk3), // | |
| (const char *) v->data + (iv2 * nbv2 + iv3 * nbv3), // | |
| mp, // | |
| sinks_buffer.data(), // | |
| dst_buffer.data(), // | |
| scale, logit_softcap, slope, nek1, nbk1, nbv1, DV, DK, tcm_buffer, tcm_buffer_size); | |
| } else { | |
| ir_step = 1; | |
| flash_attn_ext_f16_one_chunk_inner_vlen1024_vf16_m1( // | |
| pq, // | |
| (const char *) k->data + (ik2 * nbk2 + ik3 * nbk3), // | |
| (const char *) v->data + (iv2 * nbv2 + iv3 * nbv3), // | |
| mp, // | |
| sinks ? ((float *) ((char *) sinks->data)) + h : nullptr, // | |
| (float *) ((char *) dst->data + (iq3 * ne2 * ne1 + iq2 + iq1 * ne1) * nb1), // | |
| scale, logit_softcap, slope, nek1, nbk1, nbv1, DV, DK); | |
| } | |
| } | |
| } | |
| void forward_flash_attn_ext_f16_tiled_vlen1024_vf16(const ggml_compute_params * params, | |
| ggml_tensor * dst, | |
| int ir0, | |
| int ir1, | |
| void * tcm_buffer, | |
| size_t tcm_buffer_size) { | |
| const ggml_tensor * q = dst->src[0]; | |
| const ggml_tensor * k = dst->src[1]; | |
| const ggml_tensor * v = dst->src[2]; | |
| const ggml_tensor * mask = dst->src[3]; | |
| const ggml_tensor * sinks = dst->src[4]; | |
| GGML_TENSOR_LOCALS(int64_t, neq, q, ne) | |
| GGML_TENSOR_LOCALS(size_t, nbq, q, nb) | |
| GGML_TENSOR_LOCALS(int64_t, nek, k, ne) | |
| GGML_TENSOR_LOCALS(size_t, nbk, k, nb) | |
| GGML_TENSOR_LOCALS(int64_t, nev, v, ne) | |
| GGML_TENSOR_LOCALS(size_t, nbv, v, nb) | |
| GGML_TENSOR_LOCALS(int64_t, ne, dst, ne) | |
| GGML_TENSOR_LOCALS(size_t, nb, dst, nb) | |
| const int64_t DK = nek0; | |
| const int64_t DV = nev0; | |
| const int64_t N = neq1; | |
| GGML_ASSERT(flash_attn_ext_supported_shape_vlen1024_vf16(DK, DV)); | |
| GGML_ASSERT(ne0 == DV); | |
| GGML_ASSERT(ne2 == N); | |
| // input tensor rows must be contiguous | |
| GGML_ASSERT(nbq0 == ggml_type_size(q->type)); | |
| GGML_ASSERT(nbk0 == ggml_type_size(k->type)); | |
| GGML_ASSERT(nbv0 == ggml_type_size(v->type)); | |
| GGML_ASSERT(neq0 == DK); | |
| GGML_ASSERT(nek0 == DK); | |
| GGML_ASSERT(nev0 == DV); | |
| GGML_ASSERT(neq1 == N); | |
| // dst cannot be transposed or permuted | |
| GGML_ASSERT(nb0 == sizeof(float)); | |
| GGML_ASSERT(nb0 <= nb1); | |
| GGML_ASSERT(nb1 <= nb2); | |
| GGML_ASSERT(nb2 <= nb3); | |
| GGML_ASSERT(k->type == v->type); | |
| const ggml_type kv_type = k->type; | |
| // broadcast factors | |
| const int64_t rk2 = neq2 / nek2; | |
| const int64_t rk3 = neq3 / nek3; | |
| const int64_t rv2 = neq2 / nev2; | |
| const int64_t rv3 = neq3 / nev3; | |
| float * param_list = (float *) dst->op_params; | |
| float scale = param_list[0]; | |
| float max_bias = param_list[1]; | |
| float logit_softcap = param_list[2]; | |
| if (logit_softcap != 0) { | |
| scale /= logit_softcap; | |
| } | |
| const uint32_t n_head = neq2; | |
| const uint32_t n_head_log2 = 1u << (uint32_t) floor(log2(n_head)); | |
| const float m0 = powf(2.0f, -(max_bias) / n_head_log2); | |
| const float m1 = powf(2.0f, -(max_bias / 2.0f) / n_head_log2); | |
| int ith = params->ith; | |
| static constexpr int Q_TILE_SZ = ggml_fa_tile_config::Q; | |
| static constexpr int KV_TILE_SZ = ggml_fa_tile_config::KV; | |
| // Per-thread scratch layout: | |
| // Q_f32: Q_TILE_SZ * DK | |
| // KQ: Q_TILE_SZ * KV_TILE_SZ | |
| // mask32: Q_TILE_SZ * KV_TILE_SZ | |
| // VKQ32: Q_TILE_SZ * DV | |
| // V32: KV_TILE_SZ * DV | |
| // K_f32: DK * KV_TILE_SZ (transposed K tile) | |
| float * base = (float *) params->wdata + ith * (Q_TILE_SZ * DK + 2 * Q_TILE_SZ * KV_TILE_SZ + Q_TILE_SZ * DV + | |
| KV_TILE_SZ * DV + KV_TILE_SZ * DK + CACHE_LINE_SIZE_F32); | |
| const size_t base_size = | |
| (Q_TILE_SZ * DK + 2 * Q_TILE_SZ * KV_TILE_SZ + Q_TILE_SZ * DV + KV_TILE_SZ * DV + KV_TILE_SZ * DK) * | |
| sizeof(float) + | |
| CACHE_LINE_SIZE_F32; | |
| if (base_size <= tcm_buffer_size && tcm_buffer != nullptr) { | |
| base = (float *) tcm_buffer; | |
| } | |
| float S_M_Buf[Q_TILE_SZ * 2]; // buffer to hold S, M, bias for one tile to reduce register pressure in main loop | |
| float * S = S_M_Buf; | |
| float * M = S_M_Buf + Q_TILE_SZ; | |
| int ir = ir0; | |
| while (ir < ir1) { | |
| // q indices for the start of this tile | |
| const int iq3 = ir / (neq2 * neq1); | |
| const int iq2 = (ir - iq3 * neq2 * neq1) / neq1; | |
| const int iq1 = (ir - iq3 * neq2 * neq1 - iq2 * neq1); | |
| // Number of valid rows in this tile: | |
| // - limited by tile size (Q_TILE_SZ) | |
| // - limited by chunk boundary (ir1 - ir) | |
| // - limited by head boundary (neq1 - iq1) to avoid crossing into next head | |
| const int tile_rows = MIN(Q_TILE_SZ, MIN((int) (ir1 - ir), (int) (neq1 - iq1))); | |
| GGML_ASSERT(tile_rows > 0); | |
| const uint32_t h = iq2; // head index | |
| const float slope = | |
| (max_bias > 0.0f) ? h < n_head_log2 ? powf(m0, h + 1) : powf(m1, 2 * (h - n_head_log2) + 1) : 1.0f; | |
| for (int i = 0; i < Q_TILE_SZ; ++i) { | |
| S[i] = 0.; | |
| M[i] = -INFINITY; | |
| } | |
| float * Q_f32 = base; | |
| float * KQ = (float *) ((char *) base + Q_TILE_SZ * DK * sizeof(float)); | |
| float * mask32 = KQ + Q_TILE_SZ * KV_TILE_SZ; | |
| float * VKQ32 = mask32 + Q_TILE_SZ * KV_TILE_SZ; | |
| float * V32 = VKQ32 + Q_TILE_SZ * DV; | |
| float * K_f32 = V32 + KV_TILE_SZ * DV; | |
| _Float16 * Q_f16 = (_Float16 *) Q_f32; | |
| _Float16 * V_f16 = (_Float16 *) V32; | |
| _Float16 * K_f16 = (_Float16 *) K_f32; | |
| rvv_zero_f32(VKQ32, Q_TILE_SZ * DV); | |
| // k indices | |
| const int ik3 = iq3 / rk3; | |
| const int ik2 = iq2 / rk2; | |
| // v indices | |
| const int iv3 = iq3 / rv3; | |
| const int iv2 = iq2 / rv2; | |
| const float * pq = (const float *) ((char *) q->data + (iq1 * nbq1 + iq2 * nbq2 + iq3 * nbq3)); | |
| if (kv_type == GGML_TYPE_F16) { | |
| rvv_pack_f32_as_scaled_f16((uint8_t *) Q_f16, DK * sizeof(_Float16), (uint8_t *) pq, nbq1, tile_rows, DK, | |
| scale); | |
| } else { | |
| memcpy2d(Q_f32, DK * sizeof(float), pq, nbq1, tile_rows, DK * sizeof(float)); | |
| } | |
| for (int64_t ic = 0; ic < nek1; ic += KV_TILE_SZ) { | |
| const int kv_tile = (int) std::min((int64_t) KV_TILE_SZ, nek1 - ic); | |
| rvv_zero_f32(K_f32, DK * KV_TILE_SZ); | |
| rvv_zero_f32(V32, KV_TILE_SZ * DV); | |
| // skip the tile entirely if all the masks are -inf | |
| if (mask) { | |
| bool can_skip = true; | |
| const ggml_fp16_t * mp_row = | |
| (const ggml_fp16_t *) ((const char *) mask->data + iq1 * mask->nb[1] + | |
| (iq2 % mask->ne[2]) * mask->nb[2] + (iq3 % mask->ne[3]) * mask->nb[3]); | |
| rvv_pack_scaled_f16_as_f32(mask32, KV_TILE_SZ * sizeof(float), mp_row + ic, mask->nb[1], tile_rows, | |
| kv_tile, slope); | |
| for (int tq = 0; tq < tile_rows; tq++) { | |
| for (int tk = 0; tk < kv_tile; tk++) { | |
| if (mask32[tq * KV_TILE_SZ + tk] != -INFINITY) { | |
| can_skip = false; | |
| } | |
| } | |
| // Pad remaining mask entries with -inf | |
| for (int tk = kv_tile; tk < KV_TILE_SZ; tk++) { | |
| mask32[tq * KV_TILE_SZ + tk] = -INFINITY; | |
| } | |
| } | |
| if (can_skip) { | |
| continue; | |
| } | |
| } | |
| if (kv_type == GGML_TYPE_F16) { | |
| rvv_transposed_s16_mn_to_nm((int8_t *) K_f16, KV_TILE_SZ * sizeof(_Float16), | |
| (int8_t *) k->data + ic * nbk1 + ik2 * nbk2 + ik3 * nbk3, nbk1, kv_tile, | |
| DK); | |
| int tq = 0; | |
| for (; tq + 3 < tile_rows; tq += 4) { | |
| rvv_qk_dot_tile_f16_x4(KQ + (tq + 0) * KV_TILE_SZ, KQ + (tq + 1) * KV_TILE_SZ, | |
| KQ + (tq + 2) * KV_TILE_SZ, KQ + (tq + 3) * KV_TILE_SZ, | |
| Q_f16 + (tq + 0) * DK, Q_f16 + (tq + 1) * DK, Q_f16 + (tq + 2) * DK, | |
| Q_f16 + (tq + 3) * DK, K_f16, DK, kv_tile); | |
| } | |
| for (; tq < tile_rows; ++tq) { | |
| rvv_qk_dot_tile_f16_x1(KQ + tq * KV_TILE_SZ, Q_f16 + tq * DK, K_f16, DK, kv_tile); | |
| } | |
| } else { | |
| for (int tk = 0; tk < kv_tile; tk++) { | |
| const char * k_data = (const char *) k->data + (ic + tk) * nbk1 + ik2 * nbk2 + ik3 * nbk3; | |
| float * k_col = K_f32 + tk; | |
| const float * k_src = (const float *) k_data; | |
| for (int64_t dk = 0; dk < DK; ++dk) { | |
| k_col[dk * KV_TILE_SZ] = k_src[dk]; | |
| } | |
| } | |
| for (int tq = 0; tq < tile_rows; ++tq) { | |
| rvv_qk_dot_tile(KQ + tq * KV_TILE_SZ, Q_f32 + tq * DK, K_f32, DK, KV_TILE_SZ, scale); | |
| } | |
| } | |
| // Set padded KQ entries to -inf so softmax gives them zero weight | |
| if (kv_tile < KV_TILE_SZ) { | |
| for (int tq = 0; tq < tile_rows; tq++) { | |
| for (int tk = kv_tile; tk < KV_TILE_SZ; tk++) { | |
| KQ[tq * KV_TILE_SZ + tk] = -INFINITY; | |
| } | |
| } | |
| } | |
| if (logit_softcap != 0.0f) { | |
| rvv_softcap_tanh_inplace_f32(KQ, KV_TILE_SZ, tile_rows, KV_TILE_SZ, logit_softcap); | |
| } | |
| if (mask) { | |
| rvv_add_inplace_f32(KQ, KV_TILE_SZ, mask32, KV_TILE_SZ, tile_rows, KV_TILE_SZ); | |
| } | |
| bool skip[Q_TILE_SZ] = {}; | |
| for (int tq = 0; tq < tile_rows; tq++) { | |
| float * kq_row = KQ + tq * KV_TILE_SZ; | |
| const float tile_max = rvv_max_f32(kq_row, KV_TILE_SZ); | |
| if (tile_max == -INFINITY) { | |
| skip[tq] = true; | |
| continue; | |
| } | |
| const float Mold = M[tq]; | |
| const float Mnew = fmaxf(Mold, tile_max); | |
| if (Mnew > Mold) { | |
| const float ms = expf(Mold - Mnew); | |
| rvv_scale_f32(VKQ32 + tq * DV, ms, DV); | |
| S[tq] *= ms; | |
| } | |
| M[tq] = Mnew; | |
| S[tq] += rvv_softmax_exp_inplace_f32(kq_row, KV_TILE_SZ, Mnew); | |
| } | |
| // Pack V as contiguous [KV_TILE_SZ][DV]. | |
| if (kv_type == GGML_TYPE_F16) { | |
| const char * v_data = (const char *) v->data + ic * nbv1 + iv2 * nbv2 + iv3 * nbv3; | |
| memcpy2d(V_f16, DV * sizeof(_Float16), v_data, nbv1, kv_tile, DV * sizeof(_Float16)); | |
| int tq = 0; | |
| for (; tq + 3 < tile_rows; tq += 4) { | |
| if (skip[tq + 0] || skip[tq + 1] || skip[tq + 2] || skip[tq + 3]) { | |
| for (int i = 0; i < 4; ++i) { | |
| if (!skip[tq + i]) { | |
| rvv_pv_accumulate_f16_x1(VKQ32 + (tq + i) * DV, KQ + (tq + i) * KV_TILE_SZ, V_f16, | |
| KV_TILE_SZ, DV); | |
| } | |
| } | |
| continue; | |
| } | |
| rvv_pv_accumulate_f16_x4(VKQ32 + (tq + 0) * DV, VKQ32 + (tq + 1) * DV, VKQ32 + (tq + 2) * DV, | |
| VKQ32 + (tq + 3) * DV, KQ + (tq + 0) * KV_TILE_SZ, | |
| KQ + (tq + 1) * KV_TILE_SZ, KQ + (tq + 2) * KV_TILE_SZ, | |
| KQ + (tq + 3) * KV_TILE_SZ, V_f16, KV_TILE_SZ, DV); | |
| } | |
| for (; tq < tile_rows; ++tq) { | |
| if (!skip[tq]) { | |
| rvv_pv_accumulate_f16_x1(VKQ32 + tq * DV, KQ + tq * KV_TILE_SZ, V_f16, KV_TILE_SZ, DV); | |
| } | |
| } | |
| } else { | |
| const char * v_data = (const char *) v->data + ic * nbv1 + iv2 * nbv2 + iv3 * nbv3; | |
| memcpy2d(V32, DV * sizeof(float), v_data, nbv1, kv_tile, DV * sizeof(float)); | |
| for (int tq = 0; tq < tile_rows; ++tq) { | |
| if (!skip[tq]) { | |
| rvv_pv_accumulate(VKQ32 + tq * DV, KQ + tq * KV_TILE_SZ, V32, KV_TILE_SZ, DV); | |
| } | |
| } | |
| } | |
| } | |
| // sinks (apply only to valid rows in the tile) | |
| if (sinks) { | |
| const float s = ((float *) ((char *) sinks->data))[h]; | |
| for (int tq = 0; tq < tile_rows; tq++) { | |
| float ms = 1.0f; | |
| float vs = 1.0f; | |
| if (s > M[tq]) { | |
| ms = expf(M[tq] - s); | |
| rvv_scale_f32(VKQ32 + tq * DV, ms, DV); | |
| } else { | |
| vs = expf(s - M[tq]); | |
| } | |
| float S_temp = S[tq] * ms + vs; | |
| S[tq] = S_temp == 0.0f ? 0.0f : 1.0f / S_temp; | |
| } | |
| } else { | |
| for (int tq = 0; tq < tile_rows; tq++) { | |
| const float S_inv = S[tq] == 0.0f ? 0.0f : 1.0f / S[tq]; | |
| S[tq] = S_inv; | |
| } | |
| } | |
| float * dst_ptr = (float *) ((char *) dst->data + (iq3 * ne2 * ne1 + iq2 + (iq1) *ne1) * nb1); | |
| rvv_pack_scaled_f32_as_f32(dst_ptr, nb1 * ne1, VKQ32, DV * sizeof(float), tile_rows, DV, S); | |
| ir += tile_rows; | |
| } | |
| } | |
| void forward_rms_norm_f32(ggml_compute_params * params, ggml_tensor * op) { | |
| const ggml_tensor * src0 = op->src[0]; | |
| ggml_tensor * dst = op; | |
| GGML_ASSERT(ggml_are_same_shape(src0, dst)); | |
| GGML_ASSERT(src0->nb[0] == sizeof(float)); | |
| int ith = params->ith; | |
| int nth = params->nth; | |
| GGML_TENSOR_UNARY_OP_LOCALS | |
| float epsilon = *((float *) dst->op_params); | |
| GGML_ASSERT(epsilon > 0.0f); | |
| auto * input = (char *) src0->data; | |
| auto * output = (char *) dst->data; | |
| const auto hidden_size = ne00; | |
| const auto task_count = ne01 * ne02 * ne03; | |
| const auto task_per_thread = (task_count + nth - 1) / nth; | |
| const auto task_begin = ith * task_per_thread; | |
| const auto task_end = std::min((ith + 1) * task_per_thread, task_count); | |
| for (auto task_idx = task_begin; task_idx < task_end; task_idx++) { | |
| int64_t i03 = task_idx / (ne02 * ne01); | |
| int64_t i02 = (task_idx - i03 * ne02 * ne01) / ne01; | |
| int64_t i01 = (task_idx - i03 * ne02 * ne01 - i02 * ne01); | |
| auto * p_input = (float *) (input + i01 * nb01 + i02 * nb02 + i03 * nb03); | |
| auto * p_output = (float *) (output + i01 * nb1 + i02 * nb2 + i03 * nb3); | |
| auto * p_temp_output = p_output; | |
| size_t gvl = __riscv_vsetvlmax_e32m4(); | |
| vfloat32m4_t sum_sq = __riscv_vfmv_v_f_f32m4(0.f, gvl); | |
| int64_t length = hidden_size; | |
| while (length > 0) { | |
| gvl = __riscv_vsetvl_e32m4(length); | |
| vfloat32m4_t src_data = __riscv_vle32_v_f32m4(p_input, gvl); | |
| sum_sq = __riscv_vfmacc_vv_f32m4(sum_sq, src_data, src_data, gvl); | |
| __riscv_vse32_v_f32m4(p_temp_output, src_data, gvl); | |
| p_input += gvl; | |
| p_temp_output += gvl; | |
| length -= gvl; | |
| } | |
| gvl = __riscv_vsetvlmax_e32m1(); | |
| vfloat32m1_t zero_v = __riscv_vfmv_v_f_f32m1(0.f, gvl); | |
| vfloat32m1_t mean_square_v = | |
| __riscv_vfadd_vv_f32m1(__riscv_vget_v_f32m4_f32m1(sum_sq, 0), __riscv_vget_v_f32m4_f32m1(sum_sq, 1), gvl); | |
| mean_square_v = __riscv_vfadd_vv_f32m1(mean_square_v, __riscv_vget_v_f32m4_f32m1(sum_sq, 2), gvl); | |
| mean_square_v = __riscv_vfadd_vv_f32m1(mean_square_v, __riscv_vget_v_f32m4_f32m1(sum_sq, 3), gvl); | |
| mean_square_v = __riscv_vfredusum_vs_f32m1_f32m1(mean_square_v, zero_v, gvl); | |
| float mean_square = __riscv_vfmv_f_s_f32m1_f32(mean_square_v); | |
| mean_square /= hidden_size; | |
| mean_square = sqrt(mean_square + epsilon); | |
| mean_square = 1.0f / mean_square; | |
| length = hidden_size; | |
| p_temp_output = p_output; | |
| while (length > 0) { | |
| gvl = __riscv_vsetvl_e32m4(length); | |
| vfloat32m4_t src_data = __riscv_vle32_v_f32m4(p_temp_output, gvl); | |
| src_data = __riscv_vfmul_vf_f32m4(src_data, mean_square, gvl); | |
| __riscv_vse32_v_f32m4(p_output, src_data, gvl); | |
| p_temp_output += gvl; | |
| p_output += gvl; | |
| length -= gvl; | |
| } | |
| } | |
| } | |
| template <size_t MB_ROWS> | |
| void quantize_a_nrow_i8_ref(size_t blk_len, const float * a_ptr, size_t count_k, uint8_t * quant_a_ptr) { | |
| int64_t a_blk_stride = q8_blk_size(blk_len, true); | |
| int64_t a_nrow_block_stride = a_blk_stride * MB_ROWS; | |
| for (size_t k = 0; k < count_k; k += blk_len, quant_a_ptr += a_nrow_block_stride) { | |
| float * scale_a_ptr = reinterpret_cast<float *>(quant_a_ptr); | |
| int16_t * a_sum_ptr = reinterpret_cast<int16_t *>(quant_a_ptr + sizeof(float) * MB_ROWS); | |
| int8_t * quant_a_blk = | |
| reinterpret_cast<int8_t *>(quant_a_ptr + sizeof(float) * MB_ROWS + sizeof(int16_t) * MB_ROWS); | |
| for (size_t row = 0; row < MB_ROWS; row++) { | |
| float max_abs_a = 0.0f; | |
| for (size_t bk = 0; bk < blk_len; bk++) { | |
| max_abs_a = std::max(max_abs_a, std::abs(a_ptr[row * count_k + k + bk])); | |
| } | |
| float rep_scale_a = ((1 << 7) - 1) / max_abs_a; | |
| scale_a_ptr[row] = 1 / rep_scale_a; | |
| int16_t a_sum = 0; | |
| for (size_t bk = 0; bk < blk_len; bk++) { | |
| const int8_t quantized = static_cast<int8_t>( | |
| std::clamp(std::nearbyintf(a_ptr[row * count_k + k + bk] * rep_scale_a), -128.0f, 127.0f)); | |
| quant_a_blk[row * blk_len + bk] = quantized; | |
| a_sum += quantized; | |
| } | |
| a_sum_ptr[row] = -a_sum; | |
| } | |
| } | |
| } | |
| template <size_t MB_ROWS> | |
| void quantize_a_nrow_i8_hp_ref(size_t blk_len, const float * a_ptr, size_t count_k, uint8_t * quant_a_ptr) { | |
| constexpr size_t k_subblk_len = 32; | |
| const size_t subblk_count = blk_len / k_subblk_len; | |
| GGML_ASSERT(blk_len == 256); | |
| float scale_temp[8] = { 0.0f }; | |
| int64_t a_blk_stride = q8_hp_blk_size(blk_len, true, true); | |
| int64_t a_nrow_block_stride = a_blk_stride * MB_ROWS; | |
| int64_t a_subblk_stride = q8_hp_blk_size(k_subblk_len, false, false) * MB_ROWS; | |
| for (size_t k = 0; k < count_k; k += blk_len, quant_a_ptr += a_nrow_block_stride) { | |
| _Float16 * a_sum_ptr = reinterpret_cast<_Float16 *>(quant_a_ptr + a_subblk_stride * subblk_count); | |
| float scale_avg = 0.0f; | |
| for (size_t kk = 0; kk < subblk_count; kk++) { | |
| float max_abs_a = 0.0f; | |
| for (size_t row = 0; row < MB_ROWS; row++) { | |
| for (size_t bk = 0; bk < k_subblk_len; bk++) { | |
| max_abs_a = std::max(max_abs_a, std::abs(a_ptr[row * count_k + k + bk + kk * k_subblk_len])); | |
| } | |
| } | |
| scale_temp[kk] = max_abs_a / ((1 << 7) - 1); | |
| scale_avg += scale_temp[kk]; | |
| } | |
| scale_avg /= subblk_count; | |
| float scale_factor = 1.0f / scale_avg; | |
| _Float16 * scale_avg_ptr = | |
| reinterpret_cast<_Float16 *>(quant_a_ptr + a_nrow_block_stride - sizeof(_Float16) * MB_ROWS); | |
| scale_avg_ptr[0] = scale_avg; | |
| for (size_t kk = 0; kk < subblk_count; kk++) { | |
| uint8_t * a_subblk_base = quant_a_ptr + kk * a_subblk_stride; | |
| _Float16 * scale_a_ptr = reinterpret_cast<_Float16 *>(a_subblk_base); | |
| int8_t * quant_a_blk = reinterpret_cast<int8_t *>(a_subblk_base + sizeof(_Float16) * MB_ROWS); | |
| scale_a_ptr[0] = static_cast<_Float16>(scale_temp[kk] * scale_factor); | |
| const float rep_scale_a = 1.0f / scale_temp[kk]; | |
| for (size_t row = 0; row < MB_ROWS; row++) { | |
| int16_t a_sum = 0; | |
| for (size_t bk = 0; bk < k_subblk_len; bk++) { | |
| const int8_t quantized = static_cast<int8_t>( | |
| std::clamp(std::nearbyintf(a_ptr[row * count_k + k + bk + kk * k_subblk_len] * rep_scale_a), | |
| -128.0f, 127.0f)); | |
| quant_a_blk[row * k_subblk_len + bk] = quantized; | |
| a_sum += quantized; | |
| } | |
| a_sum_ptr[row * subblk_count + kk] = static_cast<_Float16>(-a_sum) * static_cast<_Float16>(8.0f); | |
| } | |
| } | |
| } | |
| } | |
| template <size_t MB_ROWS> | |
| void quantize_a_nrow_i8k_ref(size_t blk_len, const float * a_ptr, size_t count_k, uint8_t * quant_a_ptr) { | |
| int64_t a_blk_stride = q8k_blk_size(256); | |
| int64_t a_nrow_block_stride = a_blk_stride * MB_ROWS; | |
| int64_t a_sum_size = 256 / 16; | |
| for (size_t k = 0; k < count_k; k += blk_len, quant_a_ptr += a_nrow_block_stride) { | |
| float * scale_a_ptr = reinterpret_cast<float *>(quant_a_ptr); | |
| int16_t * a_sum_ptr = reinterpret_cast<int16_t *>(quant_a_ptr + sizeof(float) * MB_ROWS); | |
| int8_t * quant_a_blk = | |
| reinterpret_cast<int8_t *>(quant_a_ptr + sizeof(float) * MB_ROWS + sizeof(int16_t) * a_sum_size * MB_ROWS); | |
| for (size_t row = 0; row < MB_ROWS; row++) { | |
| float max_a = 0.0f; | |
| float max_abs_a = 0.0f; | |
| for (size_t bk = 0; bk < blk_len; bk++) { | |
| float ax = std::abs(a_ptr[row * count_k + k + bk]); | |
| if (ax > max_abs_a) { | |
| max_abs_a = ax; | |
| max_a = a_ptr[row * count_k + k + bk]; | |
| } | |
| } | |
| if (!max_abs_a) { | |
| scale_a_ptr[row] = 0; | |
| for (size_t bki = 0; bki < a_sum_size; bki++) { | |
| for (size_t bk = bki * 16; bk < (bki + 1) * 16; bk++) { | |
| quant_a_blk[row * blk_len + bk] = 0; | |
| } | |
| a_sum_ptr[row * a_sum_size + bki] = 0; | |
| } | |
| continue; | |
| } | |
| float rep_scale_a = ((1 << 7) - 1) / max_abs_a; | |
| scale_a_ptr[row] = 1 / rep_scale_a; | |
| for (size_t bki = 0; bki < a_sum_size; bki++) { | |
| int16_t a_sum = 0; | |
| for (size_t bk = bki * 16; bk < (bki + 1) * 16; bk++) { | |
| const int8_t quantized = static_cast<int8_t>( | |
| std::clamp(std::nearbyintf(a_ptr[row * count_k + k + bk] * rep_scale_a), -128.0f, 127.0f)); | |
| quant_a_blk[row * blk_len + bk] = quantized; | |
| a_sum += quantized; | |
| } | |
| a_sum_ptr[row * a_sum_size + bki] = -a_sum; | |
| } | |
| } | |
| } | |
| } | |
| void quantize_a_row_i8(size_t blk_len, const float * a_ptr, size_t count_k, uint8_t * quant_a_ptr) { | |
| GGML_ASSERT(blk_len == 32); | |
| int64_t a_blk_stride = q8_blk_size(blk_len, true); | |
| size_t vlenb = __riscv_vlenb(); | |
| if (vlenb == 128) { | |
| for (size_t k = 0; k < count_k; k += blk_len, quant_a_ptr += a_blk_stride) { | |
| float * scale_a_ptr = reinterpret_cast<float *>(quant_a_ptr); | |
| int16_t * a_sum_ptr = reinterpret_cast<int16_t *>(quant_a_ptr + sizeof(float)); | |
| int8_t * quant_a_blk = reinterpret_cast<int8_t *>(quant_a_ptr + sizeof(float) + sizeof(int16_t)); | |
| size_t vl = __riscv_vsetvl_e32m1(blk_len); | |
| vfloat32m1_t v_a = __riscv_vle32_v_f32m1(a_ptr + k, vl); | |
| vfloat32m1_t v_a_abs = __riscv_vfabs_v_f32m1(v_a, vl); | |
| vfloat32m1_t tmp = __riscv_vfmv_v_f_f32m1(0.0f, vl); | |
| vfloat32m1_t v_a_max = __riscv_vfredmax_vs_f32m1_f32m1(v_a_abs, tmp, vl); | |
| float max_abs_a = __riscv_vfmv_f_s_f32m1_f32(v_a_max); | |
| float scale_a = max_abs_a / ((1 << 7) - 1); | |
| float rep_scale_a = scale_a ? 1.0f / scale_a : 0.0f; | |
| scale_a_ptr[0] = scale_a; | |
| vfloat32m1_t v_a_scale = __riscv_vfmul_vf_f32m1(v_a, rep_scale_a, vl); | |
| vint16mf2_t v_a_quant = __riscv_vfncvt_x_f_w_i16mf2(v_a_scale, vl); | |
| vint8mf4_t v_a_quant_i8 = __riscv_vncvt_x_x_w_i8mf4(v_a_quant, vl); | |
| vint16m1_t tmp_sum = __riscv_vmv_v_x_i16m1(0, vl); | |
| vint16m1_t v_a_sum = __riscv_vwredsum_vs_i8mf4_i16m1(v_a_quant_i8, tmp_sum, vl); | |
| int16_t a_sum = __riscv_vmv_x_s_i16m1_i16(v_a_sum); | |
| a_sum_ptr[0] = -a_sum; | |
| __riscv_vse8_v_i8mf4(quant_a_blk, v_a_quant_i8, vl); | |
| } | |
| } else if (vlenb == 32) { | |
| for (size_t k = 0; k < count_k; k += blk_len, quant_a_ptr += a_blk_stride) { | |
| float * scale_a_ptr = reinterpret_cast<float *>(quant_a_ptr); | |
| int16_t * a_sum_ptr = reinterpret_cast<int16_t *>(quant_a_ptr + sizeof(float)); | |
| int8_t * quant_a_blk = reinterpret_cast<int8_t *>(quant_a_ptr + sizeof(float) + sizeof(int16_t)); | |
| size_t vl = __riscv_vsetvl_e32m4(blk_len); | |
| vfloat32m4_t v_a = __riscv_vle32_v_f32m4(a_ptr + k, vl); | |
| vfloat32m4_t v_a_abs = __riscv_vfabs_v_f32m4(v_a, vl); | |
| vfloat32m1_t tmp = __riscv_vfmv_v_f_f32m1(0.0f, vl); | |
| vfloat32m1_t v_a_max = __riscv_vfredmax_vs_f32m4_f32m1(v_a_abs, tmp, vl); | |
| float max_abs_a = __riscv_vfmv_f_s_f32m1_f32(v_a_max); | |
| float scale_a = max_abs_a / ((1 << 7) - 1); | |
| float rep_scale_a = scale_a ? 1.0f / scale_a : 0.0f; | |
| scale_a_ptr[0] = scale_a; | |
| vfloat32m4_t v_a_scale = __riscv_vfmul_vf_f32m4(v_a, rep_scale_a, vl); | |
| vint16m2_t v_a_quant = __riscv_vfncvt_x_f_w_i16m2(v_a_scale, vl); | |
| vint8m1_t v_a_quant_i8 = __riscv_vncvt_x_x_w_i8m1(v_a_quant, vl); | |
| vint16m1_t tmp_sum = __riscv_vmv_v_x_i16m1(0, vl); | |
| vint16m1_t v_a_sum = __riscv_vwredsum_vs_i8m1_i16m1(v_a_quant_i8, tmp_sum, vl); | |
| int16_t a_sum = __riscv_vmv_x_s_i16m1_i16(v_a_sum); | |
| a_sum_ptr[0] = -a_sum; | |
| __riscv_vse8_v_i8m1(quant_a_blk, v_a_quant_i8, vl); | |
| } | |
| } else { | |
| quantize_a_nrow_i8_ref<1>(blk_len, a_ptr, count_k, quant_a_ptr); | |
| } | |
| } | |
| void quantize_a_4row_i8(size_t blk_len, const float * a_ptr, size_t count_k, uint8_t * quant_a_ptr) { | |
| GGML_ASSERT(blk_len == 32); | |
| int64_t a_blk_stride = q8_blk_size(blk_len, true); | |
| int64_t a_nrow_block_stride = a_blk_stride * 4; | |
| size_t vlenb = __riscv_vlenb(); | |
| if (vlenb == 128) { | |
| for (size_t k = 0; k < count_k; k += blk_len, quant_a_ptr += a_nrow_block_stride) { | |
| float * scale_a_ptr = reinterpret_cast<float *>(quant_a_ptr); | |
| int16_t * a_sum_ptr = reinterpret_cast<int16_t *>(quant_a_ptr + sizeof(float) * 4); | |
| int8_t * quant_a_blk = reinterpret_cast<int8_t *>(quant_a_ptr + sizeof(float) * 4 + sizeof(int16_t) * 4); | |
| for (size_t mi = 0; mi < 4; mi++) { | |
| size_t vl = __riscv_vsetvl_e32m1(blk_len); | |
| vfloat32m1_t v_a = __riscv_vle32_v_f32m1(a_ptr + mi * count_k + k, vl); | |
| vfloat32m1_t v_a_abs = __riscv_vfabs_v_f32m1(v_a, vl); | |
| vfloat32m1_t tmp = __riscv_vfmv_v_f_f32m1(0.0f, vl); | |
| vfloat32m1_t v_a_max = __riscv_vfredmax_vs_f32m1_f32m1(v_a_abs, tmp, vl); | |
| float max_abs_a = __riscv_vfmv_f_s_f32m1_f32(v_a_max); | |
| float scale_a = max_abs_a / ((1 << 7) - 1); | |
| float rep_scale_a = scale_a ? 1.0f / scale_a : 0.0f; | |
| scale_a_ptr[mi] = scale_a; | |
| vfloat32m1_t v_a_scale = __riscv_vfmul_vf_f32m1(v_a, rep_scale_a, vl); | |
| vint16mf2_t v_a_quant = __riscv_vfncvt_x_f_w_i16mf2(v_a_scale, vl); | |
| vint8mf4_t v_a_quant_i8 = __riscv_vncvt_x_x_w_i8mf4(v_a_quant, vl); | |
| vint16m1_t tmp_sum = __riscv_vmv_v_x_i16m1(0, vl); | |
| vint16m1_t v_a_sum = __riscv_vwredsum_vs_i8mf4_i16m1(v_a_quant_i8, tmp_sum, vl); | |
| int16_t a_sum = __riscv_vmv_x_s_i16m1_i16(v_a_sum); | |
| a_sum_ptr[mi] = -a_sum; | |
| __riscv_vse8_v_i8mf4(quant_a_blk + mi * blk_len, v_a_quant_i8, vl); | |
| } | |
| } | |
| } else if (vlenb == 32) { | |
| for (size_t k = 0; k < count_k; k += blk_len, quant_a_ptr += a_nrow_block_stride) { | |
| float * scale_a_ptr = reinterpret_cast<float *>(quant_a_ptr); | |
| int16_t * a_sum_ptr = reinterpret_cast<int16_t *>(quant_a_ptr + sizeof(float) * 4); | |
| int8_t * quant_a_blk = reinterpret_cast<int8_t *>(quant_a_ptr + sizeof(float) * 4 + sizeof(int16_t) * 4); | |
| for (size_t mi = 0; mi < 4; mi++) { | |
| size_t vl = __riscv_vsetvl_e32m4(blk_len); | |
| vfloat32m4_t v_a = __riscv_vle32_v_f32m4(a_ptr + mi * count_k + k, vl); | |
| vfloat32m4_t v_a_abs = __riscv_vfabs_v_f32m4(v_a, vl); | |
| vfloat32m1_t tmp = __riscv_vfmv_v_f_f32m1(0.0f, vl); | |
| vfloat32m1_t v_a_max = __riscv_vfredmax_vs_f32m4_f32m1(v_a_abs, tmp, vl); | |
| float max_abs_a = __riscv_vfmv_f_s_f32m1_f32(v_a_max); | |
| float scale_a = max_abs_a / ((1 << 7) - 1); | |
| float rep_scale_a = scale_a ? 1.0f / scale_a : 0.0f; | |
| scale_a_ptr[mi] = scale_a; | |
| vfloat32m4_t v_a_scale = __riscv_vfmul_vf_f32m4(v_a, rep_scale_a, vl); | |
| vint16m2_t v_a_quant = __riscv_vfncvt_x_f_w_i16m2(v_a_scale, vl); | |
| vint8m1_t v_a_quant_i8 = __riscv_vncvt_x_x_w_i8m1(v_a_quant, vl); | |
| vint16m1_t tmp_sum = __riscv_vmv_v_x_i16m1(0, vl); | |
| vint16m1_t v_a_sum = __riscv_vwredsum_vs_i8m1_i16m1(v_a_quant_i8, tmp_sum, vl); | |
| int16_t a_sum = __riscv_vmv_x_s_i16m1_i16(v_a_sum); | |
| a_sum_ptr[mi] = -a_sum; | |
| __riscv_vse8_v_i8m1(quant_a_blk + mi * blk_len, v_a_quant_i8, vl); | |
| } | |
| } | |
| } else { | |
| quantize_a_nrow_i8_ref<4>(blk_len, a_ptr, count_k, quant_a_ptr); | |
| } | |
| } | |
| void quantize_a_row_i8_hp(size_t blk_len, const float * a_ptr, size_t count_k, uint8_t * quant_a_ptr) { | |
| constexpr size_t k_subblk_len = 32; | |
| GGML_ASSERT(blk_len == 256); | |
| constexpr size_t subblk_count = 256 / k_subblk_len; | |
| int64_t a_blk_stride = q8_hp_blk_size(blk_len, true, true); | |
| int64_t a_subblk_stride = q8_hp_blk_size(k_subblk_len, false, false); | |
| size_t vlenb = __riscv_vlenb(); | |
| float scale_temp[subblk_count] = { 0.0f }; | |
| if (vlenb == 128) { | |
| for (size_t k = 0; k < count_k; k += blk_len, quant_a_ptr += a_blk_stride) { | |
| _Float16 * a_sum_ptr = reinterpret_cast<_Float16 *>(quant_a_ptr + a_subblk_stride * subblk_count); | |
| _Float16 * scale_avg_ptr = reinterpret_cast<_Float16 *>(quant_a_ptr + a_blk_stride - sizeof(_Float16)); | |
| float scale_avg = 0.0f; | |
| for (size_t kk = 0; kk < subblk_count; ++kk) { | |
| const float * a_src_ptr = a_ptr + k + kk * k_subblk_len; | |
| size_t vl = __riscv_vsetvl_e32m1(k_subblk_len); | |
| vfloat32m1_t v_a = __riscv_vle32_v_f32m1(a_src_ptr, vl); | |
| vfloat32m1_t v_a_abs = __riscv_vfabs_v_f32m1(v_a, vl); | |
| vfloat32m1_t tmp = __riscv_vfmv_v_f_f32m1(0.0f, vl); | |
| vfloat32m1_t v_a_max = __riscv_vfredmax_vs_f32m1_f32m1(v_a_abs, tmp, vl); | |
| float max_abs_a = __riscv_vfmv_f_s_f32m1_f32(v_a_max); | |
| scale_temp[kk] = max_abs_a / ((1 << 7) - 1); | |
| scale_avg += scale_temp[kk]; | |
| } | |
| scale_avg /= subblk_count; | |
| const float scale_factor = scale_avg ? 1.0f / scale_avg : 0.0f; | |
| scale_avg_ptr[0] = static_cast<_Float16>(scale_avg); | |
| for (size_t kk = 0; kk < subblk_count; ++kk) { | |
| uint8_t * a_subblk_base = quant_a_ptr + kk * a_subblk_stride; | |
| _Float16 * scale_a_ptr = reinterpret_cast<_Float16 *>(a_subblk_base); | |
| int8_t * quant_a_blk = reinterpret_cast<int8_t *>(a_subblk_base + sizeof(_Float16)); | |
| const float * a_src_ptr = a_ptr + k + kk * k_subblk_len; | |
| size_t vl = __riscv_vsetvl_e32m1(k_subblk_len); | |
| vfloat32m1_t v_a = __riscv_vle32_v_f32m1(a_src_ptr, vl); | |
| float rep_scale_a = scale_temp[kk] ? 1.0f / scale_temp[kk] : 0.0f; | |
| scale_a_ptr[0] = static_cast<_Float16>(scale_temp[kk] * scale_factor); | |
| vfloat32m1_t v_a_scale = __riscv_vfmul_vf_f32m1(v_a, rep_scale_a, vl); | |
| vint16mf2_t v_a_quant = __riscv_vfncvt_x_f_w_i16mf2(v_a_scale, vl); | |
| vint8mf4_t v_a_quant_i8 = __riscv_vncvt_x_x_w_i8mf4(v_a_quant, vl); | |
| vint16m1_t tmp_sum = __riscv_vmv_v_x_i16m1(0, vl); | |
| vint16m1_t v_a_sum = __riscv_vwredsum_vs_i8mf4_i16m1(v_a_quant_i8, tmp_sum, vl); | |
| int16_t a_sum = __riscv_vmv_x_s_i16m1_i16(v_a_sum); | |
| a_sum_ptr[kk] = static_cast<_Float16>(-a_sum) * static_cast<_Float16>(8.0f); | |
| __riscv_vse8_v_i8mf4(quant_a_blk, v_a_quant_i8, vl); | |
| } | |
| } | |
| } else if (vlenb == 32) { | |
| for (size_t k = 0; k < count_k; k += blk_len, quant_a_ptr += a_blk_stride) { | |
| _Float16 * a_sum_ptr = reinterpret_cast<_Float16 *>(quant_a_ptr + a_subblk_stride * subblk_count); | |
| _Float16 * scale_avg_ptr = reinterpret_cast<_Float16 *>(quant_a_ptr + a_blk_stride - sizeof(_Float16)); | |
| float scale_avg = 0.0f; | |
| for (size_t kk = 0; kk < subblk_count; ++kk) { | |
| const float * a_src_ptr = a_ptr + k + kk * k_subblk_len; | |
| size_t vl = __riscv_vsetvl_e32m4(k_subblk_len); | |
| vfloat32m4_t v_a = __riscv_vle32_v_f32m4(a_src_ptr, vl); | |
| vfloat32m4_t v_a_abs = __riscv_vfabs_v_f32m4(v_a, vl); | |
| vfloat32m1_t tmp = __riscv_vfmv_v_f_f32m1(0.0f, vl); | |
| vfloat32m1_t v_a_max = __riscv_vfredmax_vs_f32m4_f32m1(v_a_abs, tmp, vl); | |
| float max_abs_a = __riscv_vfmv_f_s_f32m1_f32(v_a_max); | |
| scale_temp[kk] = max_abs_a / ((1 << 7) - 1); | |
| scale_avg += scale_temp[kk]; | |
| } | |
| scale_avg /= subblk_count; | |
| const float scale_factor = scale_avg ? 1.0f / scale_avg : 0.0f; | |
| scale_avg_ptr[0] = static_cast<_Float16>(scale_avg); | |
| for (size_t kk = 0; kk < subblk_count; ++kk) { | |
| uint8_t * a_subblk_base = quant_a_ptr + kk * a_subblk_stride; | |
| _Float16 * scale_a_ptr = reinterpret_cast<_Float16 *>(a_subblk_base); | |
| int8_t * quant_a_blk = reinterpret_cast<int8_t *>(a_subblk_base + sizeof(_Float16)); | |
| const float * a_src_ptr = a_ptr + k + kk * k_subblk_len; | |
| size_t vl = __riscv_vsetvl_e32m4(k_subblk_len); | |
| vfloat32m4_t v_a = __riscv_vle32_v_f32m4(a_src_ptr, vl); | |
| float rep_scale_a = scale_temp[kk] ? 1.0f / scale_temp[kk] : 0.0f; | |
| scale_a_ptr[0] = static_cast<_Float16>(scale_temp[kk] * scale_factor); | |
| vfloat32m4_t v_a_scale = __riscv_vfmul_vf_f32m4(v_a, rep_scale_a, vl); | |
| vint16m2_t v_a_quant = __riscv_vfncvt_x_f_w_i16m2(v_a_scale, vl); | |
| vint8m1_t v_a_quant_i8 = __riscv_vncvt_x_x_w_i8m1(v_a_quant, vl); | |
| vint16m1_t tmp_sum = __riscv_vmv_v_x_i16m1(0, vl); | |
| vint16m1_t v_a_sum = __riscv_vwredsum_vs_i8m1_i16m1(v_a_quant_i8, tmp_sum, vl); | |
| int16_t a_sum = __riscv_vmv_x_s_i16m1_i16(v_a_sum); | |
| a_sum_ptr[kk] = static_cast<_Float16>(-a_sum) * static_cast<_Float16>(8.0f); | |
| __riscv_vse8_v_i8m1(quant_a_blk, v_a_quant_i8, vl); | |
| } | |
| } | |
| } else { | |
| quantize_a_nrow_i8_hp_ref<1>(blk_len, a_ptr, count_k, quant_a_ptr); | |
| } | |
| } | |
| void quantize_a_4row_i8_hp(size_t blk_len, const float * a_ptr, size_t count_k, uint8_t * quant_a_ptr) { | |
| constexpr size_t k_subblk_len = 32; | |
| GGML_ASSERT(blk_len == 256); | |
| constexpr size_t subblk_count = 256 / k_subblk_len; | |
| int64_t a_blk_stride = q8_hp_blk_size(blk_len, true, true); | |
| int64_t a_nrow_block_stride = a_blk_stride * 4; | |
| int64_t a_subblk_stride = q8_hp_blk_size(k_subblk_len, false, false) * 4; | |
| size_t vlenb = __riscv_vlenb(); | |
| float scale_temp[subblk_count] = { 0.0f }; | |
| if (vlenb == 128) { | |
| for (size_t k = 0; k < count_k; k += blk_len, quant_a_ptr += a_nrow_block_stride) { | |
| _Float16 * a_sum_ptr = reinterpret_cast<_Float16 *>(quant_a_ptr + a_subblk_stride * subblk_count); | |
| _Float16 * scale_avg_ptr = | |
| reinterpret_cast<_Float16 *>(quant_a_ptr + a_nrow_block_stride - sizeof(_Float16) * 4); | |
| float scale_avg = 0.0f; | |
| for (size_t kk = 0; kk < subblk_count; ++kk) { | |
| const float * a_src_ptr0 = a_ptr + 0 * count_k + k + kk * k_subblk_len; | |
| const float * a_src_ptr1 = a_ptr + 1 * count_k + k + kk * k_subblk_len; | |
| const float * a_src_ptr2 = a_ptr + 2 * count_k + k + kk * k_subblk_len; | |
| const float * a_src_ptr3 = a_ptr + 3 * count_k + k + kk * k_subblk_len; | |
| size_t vl = __riscv_vsetvl_e32m1(k_subblk_len); | |
| vfloat32m1_t v_a0 = __riscv_vle32_v_f32m1(a_src_ptr0, vl); | |
| vfloat32m1_t v_a1 = __riscv_vle32_v_f32m1(a_src_ptr1, vl); | |
| vfloat32m1_t v_a2 = __riscv_vle32_v_f32m1(a_src_ptr2, vl); | |
| vfloat32m1_t v_a3 = __riscv_vle32_v_f32m1(a_src_ptr3, vl); | |
| vfloat32m1_t v_a0_abs = __riscv_vfabs_v_f32m1(v_a0, vl); | |
| vfloat32m1_t v_a1_abs = __riscv_vfabs_v_f32m1(v_a1, vl); | |
| vfloat32m1_t v_a2_abs = __riscv_vfabs_v_f32m1(v_a2, vl); | |
| vfloat32m1_t v_a3_abs = __riscv_vfabs_v_f32m1(v_a3, vl); | |
| vfloat32m1_t v_max_abs = __riscv_vfmax_vv_f32m1(v_a0_abs, v_a1_abs, vl); | |
| v_max_abs = __riscv_vfmax_vv_f32m1(v_max_abs, v_a2_abs, vl); | |
| v_max_abs = __riscv_vfmax_vv_f32m1(v_max_abs, v_a3_abs, vl); | |
| vfloat32m1_t tmp = __riscv_vfmv_v_f_f32m1(0.0f, vl); | |
| vfloat32m1_t v_a_max = __riscv_vfredmax_vs_f32m1_f32m1(v_max_abs, tmp, vl); | |
| float max_abs_a = __riscv_vfmv_f_s_f32m1_f32(v_a_max); | |
| scale_temp[kk] = max_abs_a / ((1 << 7) - 1); | |
| scale_avg += scale_temp[kk]; | |
| } | |
| scale_avg /= subblk_count; | |
| const float scale_factor = scale_avg ? 1.0f / scale_avg : 0.0f; | |
| scale_avg_ptr[0] = static_cast<_Float16>(scale_avg); | |
| for (size_t kk = 0; kk < subblk_count; ++kk) { | |
| uint8_t * a_subblk_base = quant_a_ptr + kk * a_subblk_stride; | |
| _Float16 * scale_a_ptr = reinterpret_cast<_Float16 *>(a_subblk_base); | |
| int8_t * quant_a_blk = reinterpret_cast<int8_t *>(a_subblk_base + sizeof(_Float16) * 4); | |
| const float * a_src_ptr0 = a_ptr + 0 * count_k + k + kk * k_subblk_len; | |
| const float * a_src_ptr1 = a_ptr + 1 * count_k + k + kk * k_subblk_len; | |
| const float * a_src_ptr2 = a_ptr + 2 * count_k + k + kk * k_subblk_len; | |
| const float * a_src_ptr3 = a_ptr + 3 * count_k + k + kk * k_subblk_len; | |
| size_t vl = __riscv_vsetvl_e32m1(k_subblk_len); | |
| vfloat32m1_t v_a0 = __riscv_vle32_v_f32m1(a_src_ptr0, vl); | |
| vfloat32m1_t v_a1 = __riscv_vle32_v_f32m1(a_src_ptr1, vl); | |
| vfloat32m1_t v_a2 = __riscv_vle32_v_f32m1(a_src_ptr2, vl); | |
| vfloat32m1_t v_a3 = __riscv_vle32_v_f32m1(a_src_ptr3, vl); | |
| float rep_scale_a = scale_temp[kk] ? 1.0f / scale_temp[kk] : 0.0f; | |
| scale_a_ptr[0] = static_cast<_Float16>(scale_temp[kk] * scale_factor); | |
| vfloat32m1_t v_a0_scale = __riscv_vfmul_vf_f32m1(v_a0, rep_scale_a, vl); | |
| vfloat32m1_t v_a1_scale = __riscv_vfmul_vf_f32m1(v_a1, rep_scale_a, vl); | |
| vfloat32m1_t v_a2_scale = __riscv_vfmul_vf_f32m1(v_a2, rep_scale_a, vl); | |
| vfloat32m1_t v_a3_scale = __riscv_vfmul_vf_f32m1(v_a3, rep_scale_a, vl); | |
| vint16mf2_t v_a0_quant = __riscv_vfncvt_x_f_w_i16mf2(v_a0_scale, vl); | |
| vint16mf2_t v_a1_quant = __riscv_vfncvt_x_f_w_i16mf2(v_a1_scale, vl); | |
| vint16mf2_t v_a2_quant = __riscv_vfncvt_x_f_w_i16mf2(v_a2_scale, vl); | |
| vint16mf2_t v_a3_quant = __riscv_vfncvt_x_f_w_i16mf2(v_a3_scale, vl); | |
| vint8mf4_t v_a0_quant_i8 = __riscv_vncvt_x_x_w_i8mf4(v_a0_quant, vl); | |
| vint8mf4_t v_a1_quant_i8 = __riscv_vncvt_x_x_w_i8mf4(v_a1_quant, vl); | |
| vint8mf4_t v_a2_quant_i8 = __riscv_vncvt_x_x_w_i8mf4(v_a2_quant, vl); | |
| vint8mf4_t v_a3_quant_i8 = __riscv_vncvt_x_x_w_i8mf4(v_a3_quant, vl); | |
| vint16m1_t tmp_sum0 = __riscv_vmv_v_x_i16m1(0, vl); | |
| vint16m1_t tmp_sum1 = __riscv_vmv_v_x_i16m1(0, vl); | |
| vint16m1_t tmp_sum2 = __riscv_vmv_v_x_i16m1(0, vl); | |
| vint16m1_t tmp_sum3 = __riscv_vmv_v_x_i16m1(0, vl); | |
| vint16m1_t v_a0_sum = __riscv_vwredsum_vs_i8mf4_i16m1(v_a0_quant_i8, tmp_sum0, vl); | |
| vint16m1_t v_a1_sum = __riscv_vwredsum_vs_i8mf4_i16m1(v_a1_quant_i8, tmp_sum1, vl); | |
| vint16m1_t v_a2_sum = __riscv_vwredsum_vs_i8mf4_i16m1(v_a2_quant_i8, tmp_sum2, vl); | |
| vint16m1_t v_a3_sum = __riscv_vwredsum_vs_i8mf4_i16m1(v_a3_quant_i8, tmp_sum3, vl); | |
| a_sum_ptr[0 * subblk_count + kk] = | |
| static_cast<_Float16>(-__riscv_vmv_x_s_i16m1_i16(v_a0_sum)) * static_cast<_Float16>(8.0f); | |
| a_sum_ptr[1 * subblk_count + kk] = | |
| static_cast<_Float16>(-__riscv_vmv_x_s_i16m1_i16(v_a1_sum)) * static_cast<_Float16>(8.0f); | |
| a_sum_ptr[2 * subblk_count + kk] = | |
| static_cast<_Float16>(-__riscv_vmv_x_s_i16m1_i16(v_a2_sum)) * static_cast<_Float16>(8.0f); | |
| a_sum_ptr[3 * subblk_count + kk] = | |
| static_cast<_Float16>(-__riscv_vmv_x_s_i16m1_i16(v_a3_sum)) * static_cast<_Float16>(8.0f); | |
| __riscv_vse8_v_i8mf4(quant_a_blk + 0 * k_subblk_len, v_a0_quant_i8, vl); | |
| __riscv_vse8_v_i8mf4(quant_a_blk + 1 * k_subblk_len, v_a1_quant_i8, vl); | |
| __riscv_vse8_v_i8mf4(quant_a_blk + 2 * k_subblk_len, v_a2_quant_i8, vl); | |
| __riscv_vse8_v_i8mf4(quant_a_blk + 3 * k_subblk_len, v_a3_quant_i8, vl); | |
| } | |
| } | |
| } else if (vlenb == 32) { | |
| for (size_t k = 0; k < count_k; k += blk_len, quant_a_ptr += a_nrow_block_stride) { | |
| _Float16 * a_sum_ptr = reinterpret_cast<_Float16 *>(quant_a_ptr + a_subblk_stride * subblk_count); | |
| _Float16 * scale_avg_ptr = | |
| reinterpret_cast<_Float16 *>(quant_a_ptr + a_nrow_block_stride - sizeof(_Float16) * 4); | |
| float scale_avg = 0.0f; | |
| for (size_t kk = 0; kk < subblk_count; ++kk) { | |
| const float * a_src_ptr0 = a_ptr + 0 * count_k + k + kk * k_subblk_len; | |
| const float * a_src_ptr1 = a_ptr + 1 * count_k + k + kk * k_subblk_len; | |
| const float * a_src_ptr2 = a_ptr + 2 * count_k + k + kk * k_subblk_len; | |
| const float * a_src_ptr3 = a_ptr + 3 * count_k + k + kk * k_subblk_len; | |
| size_t vl = __riscv_vsetvl_e32m4(k_subblk_len); | |
| vfloat32m4_t v_a0 = __riscv_vle32_v_f32m4(a_src_ptr0, vl); | |
| vfloat32m4_t v_a1 = __riscv_vle32_v_f32m4(a_src_ptr1, vl); | |
| vfloat32m4_t v_a2 = __riscv_vle32_v_f32m4(a_src_ptr2, vl); | |
| vfloat32m4_t v_a3 = __riscv_vle32_v_f32m4(a_src_ptr3, vl); | |
| vfloat32m4_t v_a0_abs = __riscv_vfabs_v_f32m4(v_a0, vl); | |
| vfloat32m4_t v_a1_abs = __riscv_vfabs_v_f32m4(v_a1, vl); | |
| vfloat32m4_t v_a2_abs = __riscv_vfabs_v_f32m4(v_a2, vl); | |
| vfloat32m4_t v_a3_abs = __riscv_vfabs_v_f32m4(v_a3, vl); | |
| vfloat32m4_t v_max_abs = __riscv_vfmax_vv_f32m4(v_a0_abs, v_a1_abs, vl); | |
| v_max_abs = __riscv_vfmax_vv_f32m4(v_max_abs, v_a2_abs, vl); | |
| v_max_abs = __riscv_vfmax_vv_f32m4(v_max_abs, v_a3_abs, vl); | |
| vfloat32m1_t tmp = __riscv_vfmv_v_f_f32m1(0.0f, vl); | |
| vfloat32m1_t v_a_max = __riscv_vfredmax_vs_f32m4_f32m1(v_max_abs, tmp, vl); | |
| float max_abs_a = __riscv_vfmv_f_s_f32m1_f32(v_a_max); | |
| scale_temp[kk] = max_abs_a / ((1 << 7) - 1); | |
| scale_avg += scale_temp[kk]; | |
| } | |
| scale_avg /= subblk_count; | |
| const float scale_factor = scale_avg ? 1.0f / scale_avg : 0.0f; | |
| scale_avg_ptr[0] = static_cast<_Float16>(scale_avg); | |
| for (size_t kk = 0; kk < subblk_count; ++kk) { | |
| uint8_t * a_subblk_base = quant_a_ptr + kk * a_subblk_stride; | |
| _Float16 * scale_a_ptr = reinterpret_cast<_Float16 *>(a_subblk_base); | |
| int8_t * quant_a_blk = reinterpret_cast<int8_t *>(a_subblk_base + sizeof(_Float16) * 4); | |
| const float * a_src_ptr0 = a_ptr + 0 * count_k + k + kk * k_subblk_len; | |
| const float * a_src_ptr1 = a_ptr + 1 * count_k + k + kk * k_subblk_len; | |
| const float * a_src_ptr2 = a_ptr + 2 * count_k + k + kk * k_subblk_len; | |
| const float * a_src_ptr3 = a_ptr + 3 * count_k + k + kk * k_subblk_len; | |
| size_t vl = __riscv_vsetvl_e32m4(k_subblk_len); | |
| vfloat32m4_t v_a0 = __riscv_vle32_v_f32m4(a_src_ptr0, vl); | |
| vfloat32m4_t v_a1 = __riscv_vle32_v_f32m4(a_src_ptr1, vl); | |
| vfloat32m4_t v_a2 = __riscv_vle32_v_f32m4(a_src_ptr2, vl); | |
| vfloat32m4_t v_a3 = __riscv_vle32_v_f32m4(a_src_ptr3, vl); | |
| float rep_scale_a = scale_temp[kk] ? 1.0f / scale_temp[kk] : 0.0f; | |
| scale_a_ptr[0] = static_cast<_Float16>(scale_temp[kk] * scale_factor); | |
| vfloat32m4_t v_a0_scale = __riscv_vfmul_vf_f32m4(v_a0, rep_scale_a, vl); | |
| vfloat32m4_t v_a1_scale = __riscv_vfmul_vf_f32m4(v_a1, rep_scale_a, vl); | |
| vfloat32m4_t v_a2_scale = __riscv_vfmul_vf_f32m4(v_a2, rep_scale_a, vl); | |
| vfloat32m4_t v_a3_scale = __riscv_vfmul_vf_f32m4(v_a3, rep_scale_a, vl); | |
| vint16m2_t v_a0_quant = __riscv_vfncvt_x_f_w_i16m2(v_a0_scale, vl); | |
| vint16m2_t v_a1_quant = __riscv_vfncvt_x_f_w_i16m2(v_a1_scale, vl); | |
| vint16m2_t v_a2_quant = __riscv_vfncvt_x_f_w_i16m2(v_a2_scale, vl); | |
| vint16m2_t v_a3_quant = __riscv_vfncvt_x_f_w_i16m2(v_a3_scale, vl); | |
| vint8m1_t v_a0_quant_i8 = __riscv_vncvt_x_x_w_i8m1(v_a0_quant, vl); | |
| vint8m1_t v_a1_quant_i8 = __riscv_vncvt_x_x_w_i8m1(v_a1_quant, vl); | |
| vint8m1_t v_a2_quant_i8 = __riscv_vncvt_x_x_w_i8m1(v_a2_quant, vl); | |
| vint8m1_t v_a3_quant_i8 = __riscv_vncvt_x_x_w_i8m1(v_a3_quant, vl); | |
| vint16m1_t tmp_sum0 = __riscv_vmv_v_x_i16m1(0, vl); | |
| vint16m1_t tmp_sum1 = __riscv_vmv_v_x_i16m1(0, vl); | |
| vint16m1_t tmp_sum2 = __riscv_vmv_v_x_i16m1(0, vl); | |
| vint16m1_t tmp_sum3 = __riscv_vmv_v_x_i16m1(0, vl); | |
| vint16m1_t v_a0_sum = __riscv_vwredsum_vs_i8m1_i16m1(v_a0_quant_i8, tmp_sum0, vl); | |
| vint16m1_t v_a1_sum = __riscv_vwredsum_vs_i8m1_i16m1(v_a1_quant_i8, tmp_sum1, vl); | |
| vint16m1_t v_a2_sum = __riscv_vwredsum_vs_i8m1_i16m1(v_a2_quant_i8, tmp_sum2, vl); | |
| vint16m1_t v_a3_sum = __riscv_vwredsum_vs_i8m1_i16m1(v_a3_quant_i8, tmp_sum3, vl); | |
| a_sum_ptr[0 * subblk_count + kk] = | |
| static_cast<_Float16>(-__riscv_vmv_x_s_i16m1_i16(v_a0_sum)) * static_cast<_Float16>(8.0f); | |
| a_sum_ptr[1 * subblk_count + kk] = | |
| static_cast<_Float16>(-__riscv_vmv_x_s_i16m1_i16(v_a1_sum)) * static_cast<_Float16>(8.0f); | |
| a_sum_ptr[2 * subblk_count + kk] = | |
| static_cast<_Float16>(-__riscv_vmv_x_s_i16m1_i16(v_a2_sum)) * static_cast<_Float16>(8.0f); | |
| a_sum_ptr[3 * subblk_count + kk] = | |
| static_cast<_Float16>(-__riscv_vmv_x_s_i16m1_i16(v_a3_sum)) * static_cast<_Float16>(8.0f); | |
| __riscv_vse8_v_i8m1(quant_a_blk + 0 * k_subblk_len, v_a0_quant_i8, vl); | |
| __riscv_vse8_v_i8m1(quant_a_blk + 1 * k_subblk_len, v_a1_quant_i8, vl); | |
| __riscv_vse8_v_i8m1(quant_a_blk + 2 * k_subblk_len, v_a2_quant_i8, vl); | |
| __riscv_vse8_v_i8m1(quant_a_blk + 3 * k_subblk_len, v_a3_quant_i8, vl); | |
| } | |
| } | |
| } else { | |
| quantize_a_nrow_i8_hp_ref<4>(blk_len, a_ptr, count_k, quant_a_ptr); | |
| } | |
| } | |
| void quantize_a_row_i8k(size_t blk_len, const float * a_ptr, size_t count_k, uint8_t * quant_a_ptr) { | |
| GGML_ASSERT(blk_len == 256); | |
| constexpr int64_t a_blk_stride = q8k_blk_size(256); | |
| constexpr int64_t a_sum_size = 256 / 16; | |
| size_t vlenb = __riscv_vlenb(); | |
| if (vlenb == 128) { | |
| // vlen = 1024 bits, can process 32 float32 elements with m1 | |
| for (size_t k = 0; k < count_k; k += blk_len, quant_a_ptr += a_blk_stride) { | |
| float * scale_a_ptr = reinterpret_cast<float *>(quant_a_ptr); | |
| int16_t * a_sum_ptr = reinterpret_cast<int16_t *>(quant_a_ptr + sizeof(float)); | |
| int8_t * quant_a_blk = | |
| reinterpret_cast<int8_t *>(quant_a_ptr + sizeof(float) + sizeof(int16_t) * a_sum_size); | |
| // Find max absolute value across all 256 elements | |
| size_t vl = __riscv_vsetvl_e32m1(16); | |
| vfloat32m1_t v_max_abs = __riscv_vfmv_v_f_f32m1(0.0f, vl); | |
| for (size_t bki = 0; bki < a_sum_size; bki++) { | |
| vfloat32m1_t v_a = __riscv_vle32_v_f32m1(a_ptr + k + bki * 16, vl); | |
| vfloat32m1_t v_a_abs = __riscv_vfabs_v_f32m1(v_a, vl); | |
| v_max_abs = __riscv_vfmax_vv_f32m1(v_a_abs, v_max_abs, vl); | |
| } | |
| vfloat32m1_t tmp = __riscv_vfmv_v_f_f32m1(0.0f, vl); | |
| vfloat32m1_t v_local_max = __riscv_vfredmax_vs_f32m1_f32m1(v_max_abs, tmp, vl); | |
| float max_abs_a = __riscv_vfmv_f_s_f32m1_f32(v_local_max); | |
| float scale_a = max_abs_a / ((1 << 7) - 1); | |
| float rep_scale_a = scale_a ? 1.0f / scale_a : 0.0f; | |
| scale_a_ptr[0] = scale_a; | |
| // Quantize and compute sums for each 16-element group | |
| for (size_t bki = 0; bki < a_sum_size; bki++) { | |
| vfloat32m1_t v_a = __riscv_vle32_v_f32m1(a_ptr + k + bki * 16, vl); | |
| vfloat32m1_t v_a_scale = __riscv_vfmul_vf_f32m1(v_a, rep_scale_a, vl); | |
| vint16mf2_t v_a_quant = __riscv_vfncvt_x_f_w_i16mf2(v_a_scale, vl); | |
| vint8mf4_t v_a_quant_i8 = __riscv_vncvt_x_x_w_i8mf4(v_a_quant, vl); | |
| vint16m1_t tmp_sum = __riscv_vmv_v_x_i16m1(0, vl); | |
| vint16m1_t v_a_sum = __riscv_vwredsum_vs_i8mf4_i16m1(v_a_quant_i8, tmp_sum, vl); | |
| int16_t a_sum = __riscv_vmv_x_s_i16m1_i16(v_a_sum); | |
| a_sum_ptr[bki] = -a_sum; | |
| __riscv_vse8_v_i8mf4(quant_a_blk + bki * 16, v_a_quant_i8, vl); | |
| } | |
| } | |
| } else if (vlenb == 32) { | |
| // vlen = 256 bits, can process 8 float32 elements with m1 | |
| for (size_t k = 0; k < count_k; k += blk_len, quant_a_ptr += a_blk_stride) { | |
| float * scale_a_ptr = reinterpret_cast<float *>(quant_a_ptr); | |
| int16_t * a_sum_ptr = reinterpret_cast<int16_t *>(quant_a_ptr + sizeof(float)); | |
| int8_t * quant_a_blk = | |
| reinterpret_cast<int8_t *>(quant_a_ptr + sizeof(float) + sizeof(int16_t) * a_sum_size); | |
| // Find max absolute value across all 256 elements | |
| size_t vl = __riscv_vsetvl_e32m2(16); | |
| vfloat32m2_t v_max_abs = __riscv_vfmv_v_f_f32m2(0.0f, vl); | |
| for (size_t bki = 0; bki < a_sum_size; bki++) { | |
| vfloat32m2_t v_a = __riscv_vle32_v_f32m2(a_ptr + k + bki * 16, vl); | |
| vfloat32m2_t v_a_abs = __riscv_vfabs_v_f32m2(v_a, vl); | |
| v_max_abs = __riscv_vfmax_vv_f32m2(v_a_abs, v_max_abs, vl); | |
| } | |
| vfloat32m1_t tmp = __riscv_vfmv_v_f_f32m1(0.0f, vl); | |
| vfloat32m1_t v_local_max = __riscv_vfredmax_vs_f32m2_f32m1(v_max_abs, tmp, vl); | |
| float max_abs_a = __riscv_vfmv_f_s_f32m1_f32(v_local_max); | |
| float scale_a = max_abs_a / ((1 << 7) - 1); | |
| float rep_scale_a = scale_a ? 1.0f / scale_a : 0.0f; | |
| scale_a_ptr[0] = scale_a; | |
| // Quantize and compute sums for each 16-element group | |
| for (size_t bki = 0; bki < a_sum_size; bki++) { | |
| vfloat32m2_t v_a = __riscv_vle32_v_f32m2(a_ptr + k + bki * 16, vl); | |
| vfloat32m2_t v_a_scale = __riscv_vfmul_vf_f32m2(v_a, rep_scale_a, vl); | |
| vint16m1_t v_a_quant = __riscv_vfncvt_x_f_w_i16m1(v_a_scale, vl); | |
| vint8mf2_t v_a_quant_i8 = __riscv_vncvt_x_x_w_i8mf2(v_a_quant, vl); | |
| vint16m1_t tmp_sum = __riscv_vmv_v_x_i16m1(0, vl); | |
| vint16m1_t v_a_sum = __riscv_vwredsum_vs_i8mf2_i16m1(v_a_quant_i8, tmp_sum, vl); | |
| int16_t a_sum = __riscv_vmv_x_s_i16m1_i16(v_a_sum); | |
| a_sum_ptr[bki] = -a_sum; | |
| __riscv_vse8_v_i8mf2(quant_a_blk + bki * 16, v_a_quant_i8, vl); | |
| } | |
| } | |
| } else { | |
| quantize_a_nrow_i8k_ref<1>(blk_len, a_ptr, count_k, quant_a_ptr); | |
| } | |
| } | |
| void quantize_a_4row_i8k(size_t blk_len, const float * a_ptr, size_t count_k, uint8_t * quant_a_ptr) { | |
| GGML_ASSERT(blk_len == 256); | |
| constexpr int64_t a_blk_stride = q8k_blk_size(256); | |
| constexpr int64_t a_nrow_block_stride = a_blk_stride * 4; | |
| constexpr int64_t a_sum_size = 256 / 16; | |
| size_t vlenb = __riscv_vlenb(); | |
| if (vlenb == 128) { | |
| // vlen = 1024 bits | |
| for (size_t k = 0; k < count_k; k += blk_len, quant_a_ptr += a_nrow_block_stride) { | |
| float * scale_a_ptr = reinterpret_cast<float *>(quant_a_ptr); | |
| int16_t * a_sum_ptr = reinterpret_cast<int16_t *>(quant_a_ptr + sizeof(float) * 4); | |
| int8_t * quant_a_blk = | |
| reinterpret_cast<int8_t *>(quant_a_ptr + sizeof(float) * 4 + sizeof(int16_t) * a_sum_size * 4); | |
| for (size_t mi = 0; mi < 4; mi++) { | |
| // Find max absolute value across all 256 elements for this row | |
| size_t vl = __riscv_vsetvl_e32m1(16); | |
| vfloat32m1_t v_max_abs = __riscv_vfmv_v_f_f32m1(0.0f, vl); | |
| for (size_t bki = 0; bki < a_sum_size; bki++) { | |
| vfloat32m1_t v_a = __riscv_vle32_v_f32m1(a_ptr + mi * count_k + k + bki * 16, vl); | |
| vfloat32m1_t v_a_abs = __riscv_vfabs_v_f32m1(v_a, vl); | |
| v_max_abs = __riscv_vfmax_vv_f32m1(v_a_abs, v_max_abs, vl); | |
| } | |
| vfloat32m1_t tmp = __riscv_vfmv_v_f_f32m1(0.0f, vl); | |
| vfloat32m1_t v_local_max = __riscv_vfredmax_vs_f32m1_f32m1(v_max_abs, tmp, vl); | |
| float max_abs_a = __riscv_vfmv_f_s_f32m1_f32(v_local_max); | |
| float scale_a = max_abs_a / ((1 << 7) - 1); | |
| float rep_scale_a = scale_a ? 1.0f / scale_a : 0.0f; | |
| scale_a_ptr[mi] = scale_a; | |
| // Quantize and compute sums for each 16-element group | |
| for (size_t bki = 0; bki < a_sum_size; bki++) { | |
| vfloat32m1_t v_a = __riscv_vle32_v_f32m1(a_ptr + mi * count_k + k + bki * 16, vl); | |
| vfloat32m1_t v_a_scale = __riscv_vfmul_vf_f32m1(v_a, rep_scale_a, vl); | |
| vint16mf2_t v_a_quant = __riscv_vfncvt_x_f_w_i16mf2(v_a_scale, vl); | |
| vint8mf4_t v_a_quant_i8 = __riscv_vncvt_x_x_w_i8mf4(v_a_quant, vl); | |
| vint16m1_t tmp_sum = __riscv_vmv_v_x_i16m1(0, vl); | |
| vint16m1_t v_a_sum = __riscv_vwredsum_vs_i8mf4_i16m1(v_a_quant_i8, tmp_sum, vl); | |
| int16_t a_sum = __riscv_vmv_x_s_i16m1_i16(v_a_sum); | |
| a_sum_ptr[mi * a_sum_size + bki] = -a_sum; | |
| __riscv_vse8_v_i8mf4(quant_a_blk + mi * blk_len + bki * 16, v_a_quant_i8, vl); | |
| } | |
| } | |
| } | |
| } else if (vlenb == 32) { | |
| // vlen = 256 bits | |
| for (size_t k = 0; k < count_k; k += blk_len, quant_a_ptr += a_nrow_block_stride) { | |
| float * scale_a_ptr = reinterpret_cast<float *>(quant_a_ptr); | |
| int16_t * a_sum_ptr = reinterpret_cast<int16_t *>(quant_a_ptr + sizeof(float) * 4); | |
| int8_t * quant_a_blk = | |
| reinterpret_cast<int8_t *>(quant_a_ptr + sizeof(float) * 4 + sizeof(int16_t) * a_sum_size * 4); | |
| for (size_t mi = 0; mi < 4; mi++) { | |
| // Find max absolute value across all 256 elements for this row | |
| size_t vl = __riscv_vsetvl_e32m2(16); | |
| vfloat32m2_t v_max_abs = __riscv_vfmv_v_f_f32m2(0.0f, vl); | |
| for (size_t bki = 0; bki < a_sum_size; bki++) { | |
| vfloat32m2_t v_a = __riscv_vle32_v_f32m2(a_ptr + mi * count_k + k + bki * 16, vl); | |
| vfloat32m2_t v_a_abs = __riscv_vfabs_v_f32m2(v_a, vl); | |
| v_max_abs = __riscv_vfmax_vv_f32m2(v_a_abs, v_max_abs, vl); | |
| } | |
| vfloat32m1_t tmp = __riscv_vfmv_v_f_f32m1(0.0f, vl); | |
| vfloat32m1_t v_local_max = __riscv_vfredmax_vs_f32m2_f32m1(v_max_abs, tmp, vl); | |
| float max_abs_a = __riscv_vfmv_f_s_f32m1_f32(v_local_max); | |
| float scale_a = max_abs_a / ((1 << 7) - 1); | |
| float rep_scale_a = scale_a ? 1.0f / scale_a : 0.0f; | |
| scale_a_ptr[mi] = scale_a; | |
| // Quantize and compute sums for each 16-element group | |
| for (size_t bki = 0; bki < a_sum_size; bki++) { | |
| vfloat32m2_t v_a = __riscv_vle32_v_f32m2(a_ptr + mi * count_k + k + bki * 16, vl); | |
| vfloat32m2_t v_a_scale = __riscv_vfmul_vf_f32m2(v_a, rep_scale_a, vl); | |
| vint16m1_t v_a_quant = __riscv_vfncvt_x_f_w_i16m1(v_a_scale, vl); | |
| vint8mf2_t v_a_quant_i8 = __riscv_vncvt_x_x_w_i8mf2(v_a_quant, vl); | |
| vint16m1_t tmp_sum = __riscv_vmv_v_x_i16m1(0, vl); | |
| vint16m1_t v_a_sum = __riscv_vwredsum_vs_i8mf2_i16m1(v_a_quant_i8, tmp_sum, vl); | |
| int16_t a_sum = __riscv_vmv_x_s_i16m1_i16(v_a_sum); | |
| a_sum_ptr[mi * a_sum_size + bki] = -a_sum; | |
| __riscv_vse8_v_i8mf2(quant_a_blk + mi * blk_len + bki * 16, v_a_quant_i8, vl); | |
| } | |
| } | |
| } | |
| } else { | |
| quantize_a_nrow_i8k_ref<4>(blk_len, a_ptr, count_k, quant_a_ptr); | |
| } | |
| } | |
| void forward_cpy_with_permute(ggml_compute_params * params, ggml_tensor * op) { | |
| const ggml_tensor * src0 = op->src[0]; | |
| ggml_tensor * dst = op; | |
| const int ith = params->ith; | |
| const int nth = params->nth; | |
| // [batch, m, n] -> [batch, n, m] | |
| int64_t batch = src0->ne[2] * src0->ne[3]; | |
| int64_t m = src0->ne[1]; | |
| int64_t n = src0->ne[0]; | |
| int64_t batch_stride = src0->nb[2]; | |
| int64_t m_src_stride = src0->nb[0]; | |
| int64_t n_src_stride = src0->nb[1]; | |
| int64_t n_dst_stride = n_src_stride * m; | |
| permute_transpose_impl(src0, dst, batch, m, n, batch_stride, m_src_stride, n_src_stride, n_dst_stride, ith, nth); | |
| } | |
| void forward_cont_with_permute(ggml_compute_params * params, ggml_tensor * op) { | |
| const ggml_tensor * src0 = op->src[0]; | |
| ggml_tensor * dst = op; | |
| const int ith = params->ith; | |
| const int nth = params->nth; | |
| // [batch, m, n] -> [batch, n, m] | |
| int64_t batch = dst->ne[2] * dst->ne[3]; | |
| int64_t n = dst->ne[1]; | |
| int64_t m = dst->ne[0]; | |
| int64_t batch_stride = dst->nb[2]; | |
| int64_t m_src_stride = src0->nb[0]; | |
| int64_t n_src_stride = src0->nb[1]; | |
| int64_t n_dst_stride = dst->nb[1]; | |
| permute_transpose_impl(src0, dst, batch, m, n, batch_stride, m_src_stride, n_src_stride, n_dst_stride, ith, nth); | |
| } | |
| void forward_norm_f32(ggml_compute_params * params, ggml_tensor * op) { | |
| const ggml_tensor * src0 = op->src[0]; | |
| ggml_tensor * dst = op; | |
| GGML_ASSERT(ggml_are_same_shape(src0, dst)); | |
| GGML_ASSERT(src0->nb[0] == sizeof(float)); | |
| int ith = params->ith; | |
| int nth = params->nth; | |
| GGML_TENSOR_UNARY_OP_LOCALS | |
| float epsilon = *((float *) dst->op_params); | |
| GGML_ASSERT(epsilon > 0.0f); | |
| auto * input = (char *) src0->data; | |
| auto * output = (char *) dst->data; | |
| const auto hidden_size = ne00; | |
| const auto task_count = ne01 * ne02 * ne03; | |
| const auto task_per_thread = (task_count + nth - 1) / nth; | |
| const auto task_begin = ith * task_per_thread; | |
| const auto task_end = std::min((ith + 1) * task_per_thread, task_count); | |
| for (auto task_idx = task_begin; task_idx < task_end; task_idx++) { | |
| int64_t i03 = task_idx / (ne02 * ne01); | |
| int64_t i02 = (task_idx - i03 * ne02 * ne01) / ne01; | |
| int64_t i01 = (task_idx - i03 * ne02 * ne01 - i02 * ne01); | |
| auto * p_input = (float *) (input + i01 * nb01 + i02 * nb02 + i03 * nb03); | |
| auto * p_output = (float *) (output + i01 * nb1 + i02 * nb2 + i03 * nb3); | |
| auto * p_temp_output = p_output; | |
| size_t gvl = __riscv_vsetvlmax_e32m4(); | |
| vfloat32m4_t sum = __riscv_vfmv_v_f_f32m4(0.f, gvl); | |
| vfloat32m4_t sum_sq = __riscv_vfmv_v_f_f32m4(0.f, gvl); | |
| int64_t length = hidden_size; | |
| while (length > 0) { | |
| gvl = __riscv_vsetvl_e32m4(length); | |
| // load data | |
| vfloat32m4_t src_data = __riscv_vle32_v_f32m4(p_input, gvl); | |
| sum = __riscv_vfadd_vv_f32m4(sum, src_data, gvl); | |
| sum_sq = __riscv_vfmacc_vv_f32m4(sum_sq, src_data, src_data, gvl); | |
| __riscv_vse32_v_f32m4(p_temp_output, src_data, gvl); | |
| p_input += gvl; | |
| p_temp_output += gvl; | |
| length -= gvl; | |
| } | |
| gvl = __riscv_vsetvlmax_e32m1(); | |
| float mean = 0.f; | |
| vfloat32m1_t zero_v = __riscv_vfmv_v_f_f32m1(0.f, gvl); | |
| vfloat32m1_t mean_v = | |
| __riscv_vfadd_vv_f32m1(__riscv_vget_v_f32m4_f32m1(sum, 0), __riscv_vget_v_f32m4_f32m1(sum, 1), gvl); | |
| mean_v = __riscv_vfadd_vv_f32m1(mean_v, __riscv_vget_v_f32m4_f32m1(sum, 2), gvl); | |
| mean_v = __riscv_vfadd_vv_f32m1(mean_v, __riscv_vget_v_f32m4_f32m1(sum, 3), gvl); | |
| mean_v = __riscv_vfredusum_vs_f32m1_f32m1(mean_v, zero_v, gvl); | |
| mean = __riscv_vfmv_f_s_f32m1_f32(mean_v); | |
| mean /= hidden_size; | |
| vfloat32m1_t mean_square_v = | |
| __riscv_vfadd_vv_f32m1(__riscv_vget_v_f32m4_f32m1(sum_sq, 0), __riscv_vget_v_f32m4_f32m1(sum_sq, 1), gvl); | |
| mean_square_v = __riscv_vfadd_vv_f32m1(mean_square_v, __riscv_vget_v_f32m4_f32m1(sum_sq, 2), gvl); | |
| mean_square_v = __riscv_vfadd_vv_f32m1(mean_square_v, __riscv_vget_v_f32m4_f32m1(sum_sq, 3), gvl); | |
| mean_square_v = __riscv_vfredusum_vs_f32m1_f32m1(mean_square_v, zero_v, gvl); | |
| float mean_square = __riscv_vfmv_f_s_f32m1_f32(mean_square_v); | |
| mean_square /= hidden_size; | |
| mean_square = sqrt(mean_square - mean * mean + epsilon); | |
| mean_square = 1.0f / mean_square; | |
| length = hidden_size; | |
| p_temp_output = p_output; | |
| while (length > 0) { | |
| gvl = __riscv_vsetvl_e32m4(length); | |
| vfloat32m4_t src_data = __riscv_vle32_v_f32m4(p_temp_output, gvl); | |
| src_data = __riscv_vfsub_vf_f32m4(src_data, mean, gvl); | |
| src_data = __riscv_vfmul_vf_f32m4(src_data, mean_square, gvl); | |
| __riscv_vse32_v_f32m4(p_output, src_data, gvl); | |
| p_temp_output += gvl; | |
| p_output += gvl; | |
| length -= gvl; | |
| } | |
| } | |
| } | |
| template <ggml_op op_type, typename T> void forward_binary(ggml_compute_params * params, ggml_tensor * op) { | |
| const ggml_tensor * src0 = op->src[0]; | |
| const ggml_tensor * src1 = op->src[1]; | |
| ggml_tensor * dst = op; | |
| GGML_ASSERT(ggml_can_repeat(src1, src0) && ggml_are_same_shape(src0, dst)); | |
| auto src0_rows = ggml_nrows(src0); | |
| auto src1_rows = ggml_nrows(src1); | |
| int ith = params->ith; | |
| int nth = params->nth; | |
| GGML_TENSOR_BINARY_OP_LOCALS | |
| GGML_ASSERT(nb0 == sizeof(T)); | |
| GGML_ASSERT(nb00 == sizeof(T)); | |
| const auto [ir0, ir1] = get_thread_range(params, src0); | |
| auto compute_func_vv = [&](int64_t blk_len, int64_t r, T * src0_ptr, T * src1_ptr, T * dst_ptr) { | |
| int64_t idx = 0; | |
| if constexpr (op_type == GGML_OP_ADD) { | |
| if constexpr (std::is_same_v<T, float>) { | |
| for (size_t vl; blk_len > 0; blk_len -= vl, idx += vl) { | |
| vl = __riscv_vsetvl_e32m4(blk_len); | |
| vfloat32m4_t lhs = __riscv_vle32_v_f32m4(src0_ptr + idx + r, vl); | |
| vfloat32m4_t rhs = __riscv_vle32_v_f32m4(src1_ptr + idx, vl); | |
| vfloat32m4_t res = __riscv_vfadd_vv_f32m4(lhs, rhs, vl); | |
| __riscv_vse32_v_f32m4(dst_ptr + idx + r, res, vl); | |
| } | |
| } else if constexpr (std::is_same_v<T, _Float16>) { | |
| for (size_t vl; blk_len > 0; blk_len -= vl, idx += vl) { | |
| vl = __riscv_vsetvl_e16m4(blk_len); | |
| vfloat16m4_t lhs = __riscv_vle16_v_f16m4((src0_ptr + idx + r), vl); | |
| vfloat16m4_t rhs = __riscv_vle16_v_f16m4((src1_ptr + idx), vl); | |
| vfloat16m4_t res = __riscv_vfadd_vv_f16m4(lhs, rhs, vl); | |
| __riscv_vse16_v_f16m4((dst_ptr + idx + r), res, vl); | |
| } | |
| } else { | |
| GGML_ABORT("fatal error"); | |
| } | |
| } else if constexpr (op_type == GGML_OP_SUB) { | |
| if constexpr (std::is_same_v<T, float>) { | |
| for (size_t vl; blk_len > 0; blk_len -= vl, idx += vl) { | |
| vl = __riscv_vsetvl_e32m4(blk_len); | |
| vfloat32m4_t lhs = __riscv_vle32_v_f32m4(src0_ptr + idx + r, vl); | |
| vfloat32m4_t rhs = __riscv_vle32_v_f32m4(src1_ptr + idx, vl); | |
| vfloat32m4_t res = __riscv_vfsub_vv_f32m4(lhs, rhs, vl); | |
| __riscv_vse32_v_f32m4(dst_ptr + idx + r, res, vl); | |
| } | |
| } else if constexpr (std::is_same_v<T, _Float16>) { | |
| for (size_t vl; blk_len > 0; blk_len -= vl, idx += vl) { | |
| vl = __riscv_vsetvl_e16m4(blk_len); | |
| vfloat16m4_t lhs = __riscv_vle16_v_f16m4((src0_ptr + idx + r), vl); | |
| vfloat16m4_t rhs = __riscv_vle16_v_f16m4((src1_ptr + idx), vl); | |
| vfloat16m4_t res = __riscv_vfsub_vv_f16m4(lhs, rhs, vl); | |
| __riscv_vse16_v_f16m4((dst_ptr + idx + r), res, vl); | |
| } | |
| } else { | |
| GGML_ABORT("fatal error"); | |
| } | |
| } else if constexpr (op_type == GGML_OP_MUL) { | |
| if constexpr (std::is_same_v<T, float>) { | |
| for (size_t vl; blk_len > 0; blk_len -= vl, idx += vl) { | |
| vl = __riscv_vsetvl_e32m4(blk_len); | |
| vfloat32m4_t lhs = __riscv_vle32_v_f32m4(src0_ptr + idx + r, vl); | |
| vfloat32m4_t rhs = __riscv_vle32_v_f32m4(src1_ptr + idx, vl); | |
| vfloat32m4_t res = __riscv_vfmul_vv_f32m4(lhs, rhs, vl); | |
| __riscv_vse32_v_f32m4(dst_ptr + idx + r, res, vl); | |
| } | |
| } else if constexpr (std::is_same_v<T, _Float16>) { | |
| for (size_t vl; blk_len > 0; blk_len -= vl, idx += vl) { | |
| vl = __riscv_vsetvl_e16m4(blk_len); | |
| vfloat16m4_t lhs = __riscv_vle16_v_f16m4((src0_ptr + idx + r), vl); | |
| vfloat16m4_t rhs = __riscv_vle16_v_f16m4((src1_ptr + idx), vl); | |
| vfloat16m4_t res = __riscv_vfmul_vv_f16m4(lhs, rhs, vl); | |
| __riscv_vse16_v_f16m4((dst_ptr + idx + r), res, vl); | |
| } | |
| } else { | |
| GGML_ABORT("fatal error"); | |
| } | |
| } else if constexpr (op_type == GGML_OP_DIV) { | |
| if constexpr (std::is_same_v<T, float>) { | |
| for (size_t vl; blk_len > 0; blk_len -= vl, idx += vl) { | |
| vl = __riscv_vsetvl_e32m4(blk_len); | |
| vfloat32m4_t lhs = __riscv_vle32_v_f32m4(src0_ptr + idx + r, vl); | |
| vfloat32m4_t rhs = __riscv_vle32_v_f32m4(src1_ptr + idx, vl); | |
| vfloat32m4_t res = __riscv_vfdiv_vv_f32m4(lhs, rhs, vl); | |
| __riscv_vse32_v_f32m4(dst_ptr + idx + r, res, vl); | |
| } | |
| } else if constexpr (std::is_same_v<T, _Float16>) { | |
| for (size_t vl; blk_len > 0; blk_len -= vl, idx += vl) { | |
| vl = __riscv_vsetvl_e16m4(blk_len); | |
| vfloat16m4_t lhs = __riscv_vle16_v_f16m4((src0_ptr + idx + r), vl); | |
| vfloat16m4_t rhs = __riscv_vle16_v_f16m4((src1_ptr + idx), vl); | |
| vfloat16m4_t res = __riscv_vfdiv_vv_f16m4(lhs, rhs, vl); | |
| __riscv_vse16_v_f16m4((dst_ptr + idx + r), res, vl); | |
| } | |
| } else { | |
| GGML_ABORT("fatal error"); | |
| } | |
| } else { | |
| GGML_ABORT("fatal error"); | |
| } | |
| }; | |
| if (src0_rows == src1_rows && src0_rows == 1 && ne00 == ne10) { | |
| int64_t task_per_thread = (ne00 + nth - 1) / nth; | |
| int64_t task_begin = ith * task_per_thread; | |
| int64_t task_end = std::min((ith + 1) * task_per_thread, ne00); | |
| T * dst_ptr = ((T *) dst->data) + task_begin; | |
| T * src0_ptr = ((T *) src0->data) + task_begin; | |
| T * src1_ptr = ((T *) src1->data) + task_begin; | |
| compute_func_vv(task_end - task_begin, 0, src0_ptr, src1_ptr, dst_ptr); | |
| } else if (ne10 > 1) { | |
| for (int64_t ir = ir0; ir < ir1; ++ir) { | |
| const int64_t i03 = ir / (ne02 * ne01); | |
| const int64_t i02 = (ir - i03 * ne02 * ne01) / ne01; | |
| const int64_t i01 = (ir - i03 * ne02 * ne01 - i02 * ne01); | |
| const int64_t i13 = i03 % ne13; | |
| const int64_t i12 = i02 % ne12; | |
| const int64_t i11 = i01 % ne11; | |
| T * dst_ptr = (T *) ((char *) dst->data + i03 * nb3 + i02 * nb2 + i01 * nb1); | |
| T * src0_ptr = (T *) ((char *) src0->data + i03 * nb03 + i02 * nb02 + i01 * nb01); | |
| T * src1_ptr = (T *) ((char *) src1->data + i13 * nb13 + i12 * nb12 + i11 * nb11); | |
| // src1 is broadcastable across src0 and dst in i1, i2, i3 | |
| for (int64_t r = 0; r < ne00; r += ne10) { | |
| compute_func_vv(ne10, r, src0_ptr, src1_ptr, dst_ptr); | |
| } | |
| } | |
| } else { | |
| for (int64_t ir = ir0; ir < ir1; ++ir) { | |
| const int64_t i03 = ir / (ne02 * ne01); | |
| const int64_t i02 = (ir - i03 * ne02 * ne01) / ne01; | |
| const int64_t i01 = (ir - i03 * ne02 * ne01 - i02 * ne01); | |
| const int64_t i13 = i03 % ne13; | |
| const int64_t i12 = i02 % ne12; | |
| const int64_t i11 = i01 % ne11; | |
| T * dst_ptr = (T *) ((char *) dst->data + i03 * nb3 + i02 * nb2 + i01 * nb1); | |
| T * src0_ptr = (T *) ((char *) src0->data + i03 * nb03 + i02 * nb02 + i01 * nb01); | |
| T * src1_ptr = (T *) ((char *) src1->data + i13 * nb13 + i12 * nb12 + i11 * nb11); | |
| T rhs_scalar = src1_ptr[0]; | |
| int64_t blk_len = ne00; | |
| int64_t r = 0; | |
| for (size_t vl; blk_len > 0; blk_len -= vl, r += vl) { | |
| if constexpr (op_type == GGML_OP_ADD) { | |
| if constexpr (std::is_same_v<T, float>) { | |
| vl = __riscv_vsetvl_e32m4(blk_len); | |
| vfloat32m4_t lhs = __riscv_vle32_v_f32m4(src0_ptr + r, vl); | |
| vfloat32m4_t res = __riscv_vfadd_vf_f32m4(lhs, rhs_scalar, vl); | |
| __riscv_vse32_v_f32m4(dst_ptr + r, res, vl); | |
| } else if constexpr (std::is_same_v<T, _Float16>) { | |
| vl = __riscv_vsetvl_e16m4(blk_len); | |
| vfloat16m4_t lhs = __riscv_vle16_v_f16m4((src0_ptr + r), vl); | |
| vfloat16m4_t res = __riscv_vfadd_vf_f16m4(lhs, rhs_scalar, vl); | |
| __riscv_vse16_v_f16m4((dst_ptr + r), res, vl); | |
| } else { | |
| GGML_ABORT("fatal error"); | |
| } | |
| } else if constexpr (op_type == GGML_OP_SUB) { | |
| if constexpr (std::is_same_v<T, float>) { | |
| vl = __riscv_vsetvl_e32m4(blk_len); | |
| vfloat32m4_t lhs = __riscv_vle32_v_f32m4(src0_ptr + r, vl); | |
| vfloat32m4_t res = __riscv_vfsub_vf_f32m4(lhs, rhs_scalar, vl); | |
| __riscv_vse32_v_f32m4(dst_ptr + r, res, vl); | |
| } else if constexpr (std::is_same_v<T, _Float16>) { | |
| vl = __riscv_vsetvl_e16m4(blk_len); | |
| vfloat16m4_t lhs = __riscv_vle16_v_f16m4((src0_ptr + r), vl); | |
| vfloat16m4_t res = __riscv_vfsub_vf_f16m4(lhs, rhs_scalar, vl); | |
| __riscv_vse16_v_f16m4((dst_ptr + r), res, vl); | |
| } else { | |
| GGML_ABORT("fatal error"); | |
| } | |
| } else if constexpr (op_type == GGML_OP_MUL) { | |
| if constexpr (std::is_same_v<T, float>) { | |
| vl = __riscv_vsetvl_e32m4(blk_len); | |
| vfloat32m4_t lhs = __riscv_vle32_v_f32m4(src0_ptr + r, vl); | |
| vfloat32m4_t res = __riscv_vfmul_vf_f32m4(lhs, rhs_scalar, vl); | |
| __riscv_vse32_v_f32m4(dst_ptr + r, res, vl); | |
| } else if constexpr (std::is_same_v<T, _Float16>) { | |
| vl = __riscv_vsetvl_e16m4(blk_len); | |
| vfloat16m4_t lhs = __riscv_vle16_v_f16m4((src0_ptr + r), vl); | |
| vfloat16m4_t res = __riscv_vfmul_vf_f16m4(lhs, rhs_scalar, vl); | |
| __riscv_vse16_v_f16m4((dst_ptr + r), res, vl); | |
| } else { | |
| GGML_ABORT("fatal error"); | |
| } | |
| } else if constexpr (op_type == GGML_OP_DIV) { | |
| if constexpr (std::is_same_v<T, float>) { | |
| vl = __riscv_vsetvl_e32m4(blk_len); | |
| vfloat32m4_t lhs = __riscv_vle32_v_f32m4(src0_ptr + r, vl); | |
| vfloat32m4_t res = __riscv_vfdiv_vf_f32m4(lhs, rhs_scalar, vl); | |
| __riscv_vse32_v_f32m4(dst_ptr + r, res, vl); | |
| } else if constexpr (std::is_same_v<T, _Float16>) { | |
| vl = __riscv_vsetvl_e16m4(blk_len); | |
| vfloat16m4_t lhs = __riscv_vle16_v_f16m4((src0_ptr + r), vl); | |
| vfloat16m4_t res = __riscv_vfdiv_vf_f16m4(lhs, rhs_scalar, vl); | |
| __riscv_vse16_v_f16m4((dst_ptr + r), res, vl); | |
| } else { | |
| GGML_ABORT("fatal error"); | |
| } | |
| } else { | |
| GGML_ABORT("fatal error"); | |
| } | |
| } | |
| } | |
| } | |
| } | |
| template <typename T> void forward_sum_rows(const ggml_compute_params * params, ggml_tensor * op) { | |
| const ggml_tensor * src0 = op->src[0]; | |
| ggml_tensor * dst = op; | |
| const int ith = params->ith; | |
| const int nth = params->nth; | |
| GGML_TENSOR_UNARY_OP_LOCALS | |
| GGML_ASSERT(ne0 == 1); | |
| GGML_ASSERT(ne1 == ne01); | |
| GGML_ASSERT(ne2 == ne02); | |
| GGML_ASSERT(ne3 == ne03); | |
| int64_t n_task = ne01 * ne02 * ne03; | |
| int64_t task_per_thread = (n_task + nth - 1) / nth; | |
| int64_t ir_start = ith * task_per_thread; | |
| int64_t ir_end = std::min(ir_start + task_per_thread, n_task); | |
| for (int64_t ir = ir_start; ir < ir_end; ir++) { | |
| const int64_t i3 = ir / (ne02 * ne01); | |
| const int64_t i2 = (ir - i3 * ne02 * ne01) / ne01; | |
| const int64_t i1 = (ir - i3 * ne02 * ne01 - i2 * ne01); | |
| T * src_row = (T *) ((char *) src0->data + i1 * nb01 + i2 * nb02 + i3 * nb03); | |
| T * dst_row = (T *) ((char *) op->data + i1 * nb1 + i2 * nb2 + i3 * nb3); | |
| float row_sum = 0; | |
| if constexpr (std::is_same_v<T, float>) { | |
| size_t gvl = __riscv_vsetvlmax_e32m4(); | |
| vfloat32m4_t acc_vec = __riscv_vfmv_v_f_f32m4(0.0f, gvl); | |
| int64_t length = ne00; | |
| const float * p_data = src_row; | |
| while (length > 0) { | |
| size_t vl = __riscv_vsetvl_e32m4(length); | |
| vfloat32m4_t vec = __riscv_vle32_v_f32m4(p_data, vl); | |
| acc_vec = __riscv_vfadd_vv_f32m4(acc_vec, vec, vl); | |
| p_data += vl; | |
| length -= vl; | |
| } | |
| gvl = __riscv_vsetvlmax_e32m1(); | |
| vfloat32m1_t zero_v = __riscv_vfmv_v_f_f32m1(0.0f, gvl); | |
| vfloat32m1_t sum_v = __riscv_vfadd_vv_f32m1(__riscv_vget_v_f32m4_f32m1(acc_vec, 0), | |
| __riscv_vget_v_f32m4_f32m1(acc_vec, 1), gvl); | |
| sum_v = __riscv_vfadd_vv_f32m1(sum_v, __riscv_vget_v_f32m4_f32m1(acc_vec, 2), gvl); | |
| sum_v = __riscv_vfadd_vv_f32m1(sum_v, __riscv_vget_v_f32m4_f32m1(acc_vec, 3), gvl); | |
| sum_v = __riscv_vfredusum_vs_f32m1_f32m1(sum_v, zero_v, gvl); | |
| row_sum = __riscv_vfmv_f_s_f32m1_f32(sum_v); | |
| } else if constexpr (std::is_same_v<T, _Float16>) { | |
| size_t gvl = __riscv_vsetvlmax_e16m2(); | |
| vfloat32m4_t acc_vec = __riscv_vfmv_v_f_f32m4(0.0f, gvl); | |
| int64_t length = ne00; | |
| const _Float16 * p_data = src_row; | |
| while (length > 0) { | |
| size_t vl = __riscv_vsetvl_e16m2(length); | |
| vfloat16m2_t vec_f16 = __riscv_vle16_v_f16m2(p_data, vl); | |
| vfloat32m4_t vec_f32 = __riscv_vfwcvt_f_f_v_f32m4(vec_f16, vl); | |
| acc_vec = __riscv_vfadd_vv_f32m4(acc_vec, vec_f32, vl); | |
| p_data += vl; | |
| length -= vl; | |
| } | |
| gvl = __riscv_vsetvlmax_e32m1(); | |
| vfloat32m1_t zero_v = __riscv_vfmv_v_f_f32m1(0.0f, gvl); | |
| vfloat32m1_t sum_v = __riscv_vfadd_vv_f32m1(__riscv_vget_v_f32m4_f32m1(acc_vec, 0), | |
| __riscv_vget_v_f32m4_f32m1(acc_vec, 1), gvl); | |
| sum_v = __riscv_vfadd_vv_f32m1(sum_v, __riscv_vget_v_f32m4_f32m1(acc_vec, 2), gvl); | |
| sum_v = __riscv_vfadd_vv_f32m1(sum_v, __riscv_vget_v_f32m4_f32m1(acc_vec, 3), gvl); | |
| sum_v = __riscv_vfredusum_vs_f32m1_f32m1(sum_v, zero_v, gvl); | |
| row_sum = __riscv_vfmv_f_s_f32m1_f32(sum_v); | |
| } else { | |
| GGML_ABORT("fatal error"); | |
| } | |
| dst_row[0] = row_sum; | |
| } | |
| } | |
| template <typename T> void forward_repeat_nrows(ggml_compute_params * params, ggml_tensor * op) { | |
| const ggml_tensor * src0 = op->src[0]; | |
| ggml_tensor * dst = op; | |
| const int ith = params->ith; | |
| const int nth = params->nth; | |
| int64_t nrows = ggml_nrows(src0); | |
| int64_t nrows_per_thread = (nrows + nth - 1) / nth; | |
| int64_t ir_start = ith * nrows_per_thread; | |
| int64_t ir_end = std::min(ir_start + nrows_per_thread, nrows); | |
| if (src0->ne[0] == 1) { | |
| for (int64_t ir = ir_start; ir < ir_end; ir++) { | |
| T * src_row = (T *) ((char *) src0->data + ir * src0->nb[1]); | |
| T * dst_row = (T *) ((char *) dst->data + ir * dst->nb[1]); | |
| T src_scalar = src_row[0]; | |
| int64_t length = dst->ne[0]; | |
| int64_t idx = 0; | |
| size_t vl = 0; | |
| while (length > 0) { | |
| if constexpr (std::is_same_v<T, int32_t>) { | |
| vl = __riscv_vsetvl_e32m4(length); | |
| vint32m4_t vec = __riscv_vmv_v_x_i32m4(src_scalar, vl); | |
| __riscv_vse32_v_i32m4(dst_row + idx, vec, vl); | |
| } else if constexpr (std::is_same_v<T, int16_t>) { | |
| vl = __riscv_vsetvl_e16m4(length); | |
| vint16m4_t vec = __riscv_vmv_v_x_i16m4(src_scalar, vl); | |
| __riscv_vse16_v_i16m4((dst_row + idx), vec, vl); | |
| } else { | |
| GGML_ABORT("fatal error"); | |
| } | |
| idx += vl; | |
| length -= vl; | |
| } | |
| } | |
| } else if (src0->ne[0] == dst->ne[0]) { | |
| for (int64_t ir = ir_start; ir < ir_end; ir++) { | |
| T * src_row = (T *) ((char *) src0->data + ir * src0->nb[1]); | |
| T * dst_row = (T *) ((char *) dst->data + ir * dst->nb[1]); | |
| int64_t length = dst->ne[0]; | |
| int64_t idx = 0; | |
| size_t vl = 0; | |
| while (length > 0) { | |
| if constexpr (std::is_same_v<T, int32_t>) { | |
| vl = __riscv_vsetvl_e32m4(length); | |
| vint32m4_t vec = __riscv_vle32_v_i32m4(src_row + idx, vl); | |
| __riscv_vse32_v_i32m4(dst_row + idx, vec, vl); | |
| } else if constexpr (std::is_same_v<T, int16_t>) { | |
| vl = __riscv_vsetvl_e16m4(length); | |
| vint16m4_t vec = __riscv_vle16_v_i16m4((src_row + idx), vl); | |
| __riscv_vse16_v_i16m4((dst_row + idx), vec, vl); | |
| } else { | |
| GGML_ABORT("fatal error"); | |
| } | |
| idx += vl; | |
| length -= vl; | |
| } | |
| } | |
| } else { | |
| GGML_ABORT("fatal error"); | |
| } | |
| } | |
| template <typename T> void forward_repeat_dim1(ggml_compute_params * params, ggml_tensor * op) { | |
| const ggml_tensor * src0 = op->src[0]; | |
| ggml_tensor * dst = op; | |
| const int ith = params->ith; | |
| const int nth = params->nth; | |
| const int64_t ne0 = dst->ne[0]; | |
| const int64_t ne1 = dst->ne[1]; | |
| const int64_t ne2 = dst->ne[2]; | |
| const int64_t ne3 = dst->ne[3]; | |
| const int64_t total_batches = ne2 * ne3; | |
| const int64_t batches_per_thread = (total_batches + nth - 1) / nth; | |
| const int64_t batch_start = ith * batches_per_thread; | |
| const int64_t batch_end = std::min(batch_start + batches_per_thread, total_batches); | |
| for (int64_t b = batch_start; b < batch_end; b++) { | |
| const int64_t i3 = b / ne2; | |
| const int64_t i2 = b % ne2; | |
| T * src_base = (T *) ((char *) src0->data + i2 * src0->nb[2] + i3 * src0->nb[3]); | |
| T * dst_batch = (T *) ((char *) dst->data + i2 * dst->nb[2] + i3 * dst->nb[3]); | |
| for (int64_t i1 = 0; i1 < ne1; i1++) { | |
| T * dst_ptr = (T *) ((char *) dst_batch + i1 * dst->nb[1]); | |
| int64_t length = ne0; | |
| int64_t idx = 0; | |
| while (length > 0) { | |
| if constexpr (std::is_same_v<T, int32_t>) { | |
| size_t vl = __riscv_vsetvl_e32m4(length); | |
| vint32m4_t vec = __riscv_vle32_v_i32m4(src_base + idx, vl); | |
| __riscv_vse32_v_i32m4(dst_ptr + idx, vec, vl); | |
| idx += vl; | |
| length -= vl; | |
| } else if constexpr (std::is_same_v<T, int16_t>) { | |
| size_t vl = __riscv_vsetvl_e16m4(length); | |
| vint16m4_t vec = __riscv_vle16_v_i16m4((src_base + idx), vl); | |
| __riscv_vse16_v_i16m4((dst_ptr + idx), vec, vl); | |
| idx += vl; | |
| length -= vl; | |
| } else { | |
| GGML_ABORT("fatal error"); | |
| } | |
| } | |
| } | |
| } | |
| } | |
| template <typename T> void forward_get_rows(ggml_compute_params * params, ggml_tensor * op) { | |
| const ggml_tensor * src0 = op->src[0]; | |
| const ggml_tensor * src1 = op->src[1]; | |
| ggml_tensor * dst = op; | |
| GGML_TENSOR_BINARY_OP_LOCALS | |
| const int64_t nc = ne00; | |
| const int64_t nr = ggml_nelements(src1); | |
| assert(ne0 == nc); | |
| assert(ne02 == ne11); | |
| assert(nb00 == sizeof(float)); | |
| assert(ggml_nrows(op) == nr); | |
| const int ith = params->ith; | |
| const int nth = params->nth; | |
| int rows_nth = nth; | |
| int cols_nth = 1; | |
| if (nr == 1) { | |
| rows_nth = 1; | |
| cols_nth = nth; | |
| } | |
| // rows per thread | |
| const int dr = (nr + rows_nth - 1) / rows_nth; | |
| const int dc = (nc + cols_nth - 1) / cols_nth; | |
| int rows_ith = ith % rows_nth; | |
| int cols_ith = ith % cols_nth; | |
| // row range for this thread | |
| const int ir0 = dr * rows_ith; | |
| const int ir1 = MIN(ir0 + dr, nr); | |
| const int cr0 = dc * cols_ith; | |
| const int cr1 = MIN(cr0 + dc, nc); | |
| for (int64_t i = ir0; i < ir1; ++i) { | |
| const int64_t i12 = i / (ne11 * ne10); | |
| const int64_t i11 = (i - i12 * ne11 * ne10) / ne10; | |
| const int64_t i10 = (i - i12 * ne11 * ne10 - i11 * ne10); | |
| const int64_t i01 = *(int32_t *) ((char *) src1->data + i10 * nb10 + i11 * nb11 + i12 * nb12); | |
| GGML_ASSERT(i01 >= 0 && i01 < ne01); | |
| memcpy1d(((char *) dst->data + i10 * nb1 + i11 * nb2 + i12 * nb3) + cr0 * sizeof(T), | |
| ((char *) src0->data + i01 * nb01 + i11 * nb02 + i12 * nb03) + cr0 * sizeof(T), | |
| (cr1 - cr0) * sizeof(T)); | |
| } | |
| } | |
| template <typename T> void forward_concat(ggml_compute_params * params, ggml_tensor * op) { | |
| const ggml_tensor * src0 = op->src[0]; | |
| const ggml_tensor * src1 = op->src[1]; | |
| ggml_tensor * dst = op; | |
| GGML_ASSERT(ggml_type_size(src0->type) == sizeof(float)); | |
| GGML_TENSOR_BINARY_OP_LOCALS | |
| const int32_t dim = ggml_get_op_params_i32(dst, 0); | |
| GGML_ASSERT(dim == 0 && nb0 == sizeof(float) && nb1 == sizeof(float) * (ne00 + ne10)); | |
| const int64_t nr = ggml_nrows(dst); | |
| const int64_t nc = ne0; | |
| const int ith = params->ith; | |
| const int nth = params->nth; | |
| int rows_nth = nth; | |
| int cols_nth = 1; | |
| if (nr == 1) { | |
| rows_nth = 1; | |
| cols_nth = nth; | |
| } | |
| const int dr = (nr + rows_nth - 1) / rows_nth; | |
| const int dc = (nc + cols_nth - 1) / cols_nth; | |
| int rows_ith = ith % rows_nth; | |
| int cols_ith = ith % cols_nth; | |
| // row range for this thread | |
| const int ir0 = dr * rows_ith; | |
| const int ir1 = MIN(ir0 + dr, nr); | |
| const int cr0 = dc * cols_ith; | |
| const int cr1 = MIN(cr0 + dc, nc); | |
| int64_t o[4] = { 0, 0, 0, 0 }; | |
| o[dim] = src0->ne[dim]; | |
| const float * x; | |
| for (int64_t i = ir0; i < ir1; ++i) { | |
| const int64_t i3 = i / (ne02 * ne01); | |
| const int64_t i2 = (i - i3 * ne02 * ne01) / ne01; | |
| const int64_t i1 = (i - i3 * ne02 * ne01 - i2 * ne01); | |
| for (int i0 = cr0; i0 < cr1; i0++) { | |
| if (i0 < ne00 && i1 < ne01 && i2 < ne02 && i3 < ne03) { | |
| x = (const float *) ((const char *) src0->data + (i0) *nb00 + (i1) *nb01 + (i2) *nb02 + (i3) *nb03); | |
| } else { | |
| x = (const float *) ((const char *) src1->data + (i0 - o[0]) * nb10 + (i1 - o[1]) * nb11 + | |
| (i2 - o[2]) * nb12 + (i3 - o[3]) * nb13); | |
| } | |
| float * y = (float *) ((char *) dst->data + i0 * nb0 + i1 * nb1 + i2 * nb2 + i3 * nb3); | |
| *y = *x; | |
| } | |
| } | |
| } | |
| template void forward_binary<GGML_OP_ADD, float>(ggml_compute_params * params, ggml_tensor * op); | |
| template void forward_binary<GGML_OP_SUB, float>(ggml_compute_params * params, ggml_tensor * op); | |
| template void forward_binary<GGML_OP_MUL, float>(ggml_compute_params * params, ggml_tensor * op); | |
| template void forward_binary<GGML_OP_DIV, float>(ggml_compute_params * params, ggml_tensor * op); | |
| template void forward_binary<GGML_OP_ADD, _Float16>(ggml_compute_params * params, ggml_tensor * op); | |
| template void forward_binary<GGML_OP_SUB, _Float16>(ggml_compute_params * params, ggml_tensor * op); | |
| template void forward_binary<GGML_OP_MUL, _Float16>(ggml_compute_params * params, ggml_tensor * op); | |
| template void forward_binary<GGML_OP_DIV, _Float16>(ggml_compute_params * params, ggml_tensor * op); | |
| template void forward_sum_rows<float>(const ggml_compute_params * params, ggml_tensor * op); | |
| template void forward_sum_rows<_Float16>(const ggml_compute_params * params, ggml_tensor * op); | |
| template void forward_repeat_nrows<int32_t>(ggml_compute_params * params, ggml_tensor * op); | |
| template void forward_repeat_nrows<int16_t>(ggml_compute_params * params, ggml_tensor * op); | |
| template void forward_repeat_dim1<int32_t>(ggml_compute_params * params, ggml_tensor * op); | |
| template void forward_repeat_dim1<int16_t>(ggml_compute_params * params, ggml_tensor * op); | |
| template void forward_get_rows<int32_t>(ggml_compute_params * params, ggml_tensor * op); | |
| template void forward_get_rows<int16_t>(ggml_compute_params * params, ggml_tensor * op); | |
| template void forward_concat<int32_t>(ggml_compute_params * params, ggml_tensor * op); | |
| template void forward_concat<int16_t>(ggml_compute_params * params, ggml_tensor * op); | |
| } // namespace spacemit_kernels::rvv | |