From 69bf64d373fe7b62fdd0f501e075445312b3b8a9 Mon Sep 17 00:00:00 2001 From: isHuangXin Date: Sun, 24 May 2026 21:53:47 +0800 Subject: [PATCH 1/8] Add bitnet-embeddings-0.6b model adaptation with F16 and I2_S GGUF conversion - Add GGUF conversion tool for bitnet-embeddings-0.6b (safetensors -> F16/I2_S GGUF) - Add Qwen3 architecture support in llama.cpp submodule with per-projection RMSNorm - Add I2_S ternary quantization (2-bit packed -1/0/+1) for lossless precision - Add f16 norm weight support for correct embedding inference - Guard bitnet-lut-kernels.h include with TL1/TL2 preprocessor checks - Update llama.cpp submodule to dev-bitnet-embedding-0.6b branch - Document F16 (from multilingual-e5-0.6b) and I2_S (from bitnet-embeddings-0.6b) conversion process --- 3rdparty/llama.cpp | 2 +- ...bitnet-embeddings-qwen3-gguf-conversion.md | 302 +++++++++++ src/ggml-bitnet-lut.cpp | 7 + src/ggml-bitnet-mad.cpp | 2 +- utils/convert-bitnet-embedding-to-gguf.py | 502 ++++++++++++++++++ 5 files changed, 813 insertions(+), 2 deletions(-) create mode 100644 docs/bitnet-embeddings-qwen3-gguf-conversion.md create mode 100644 utils/convert-bitnet-embedding-to-gguf.py diff --git a/3rdparty/llama.cpp b/3rdparty/llama.cpp index 1f86f058d..078c16664 160000 --- a/3rdparty/llama.cpp +++ b/3rdparty/llama.cpp @@ -1 +1 @@ -Subproject commit 1f86f058de0c3f4098dedae2ae8653c335c868a1 +Subproject commit 078c166647bf0f96964dc2614becefbdf9ca2207 diff --git a/docs/bitnet-embeddings-qwen3-gguf-conversion.md b/docs/bitnet-embeddings-qwen3-gguf-conversion.md new file mode 100644 index 000000000..9d63c9300 --- /dev/null +++ b/docs/bitnet-embeddings-qwen3-gguf-conversion.md @@ -0,0 +1,302 @@ +# BitNet Embeddings (Qwen3) GGUF Conversion Implementation + +## 1. Background + +`bitnet-embeddings-0.6b` is a Qwen3-based embedding model with BitNet per-projection RMSNorm (`BitLinear`). Each linear projection (q/k/v/o/gate/up/down) has a `.norm.weight` that applies RMSNorm to the input **before** the matmul: + +``` +x → RMSNorm(x, norm.weight) → activation_quant(8bit) → matmul(weight_quant(ternary)) +``` + +This pattern does **not** exist in any standard llama.cpp architecture: +- Standard Qwen3: no per-projection norms +- Standard BitNet: has `attn_sub_norm`/`ffn_sub_norm` at different positions (after attention/gate*up, not before each projection) + +### Model Config + +- Architecture: `Qwen3Model` +- hidden_size: 1024, num_attention_heads: 16, num_key_value_heads: 8 +- head_dim: 128 (note: != hidden_size/num_heads = 64) +- intermediate_size: 3072, num_hidden_layers: 28 +- tie_word_embeddings: true +- rope_theta: 1000000, rms_norm_eps: 1e-06 + +### Per-Layer Tensors (7 extra norm tensors per layer) + +| Tensor | Shape | +|--------|-------| +| `self_attn.q_proj.norm.weight` | [1024] | +| `self_attn.k_proj.norm.weight` | [1024] | +| `self_attn.v_proj.norm.weight` | [1024] | +| `self_attn.o_proj.norm.weight` | [2048] | +| `mlp.gate_proj.norm.weight` | [1024] | +| `mlp.up_proj.norm.weight` | [1024] | +| `mlp.down_proj.norm.weight` | [3072] | + +--- + +## 2. GGUF Tensor Name Mapping + +| HF Name | GGUF Name | Notes | +|----------|-----------|-------| +| `embed_tokens.weight` | `token_embd.weight` | | +| `norm.weight` | `output_norm.weight` | | +| `layers.{i}.input_layernorm.weight` | `blk.{i}.attn_norm.weight` | | +| `layers.{i}.post_attention_layernorm.weight` | `blk.{i}.ffn_norm.weight` | | +| `layers.{i}.self_attn.q_proj.weight` | `blk.{i}.attn_q.weight` | | +| `layers.{i}.self_attn.k_proj.weight` | `blk.{i}.attn_k.weight` | | +| `layers.{i}.self_attn.v_proj.weight` | `blk.{i}.attn_v.weight` | | +| `layers.{i}.self_attn.o_proj.weight` | `blk.{i}.attn_output.weight` | | +| `layers.{i}.self_attn.q_norm.weight` | `blk.{i}.attn_q_norm.weight` | QK head norm | +| `layers.{i}.self_attn.k_norm.weight` | `blk.{i}.attn_k_norm.weight` | QK head norm | +| `layers.{i}.self_attn.q_proj.norm.weight` | `blk.{i}.attn_q_norm_in.weight` | BitNet per-projection | +| `layers.{i}.self_attn.k_proj.norm.weight` | `blk.{i}.attn_k_norm_in.weight` | BitNet per-projection | +| `layers.{i}.self_attn.v_proj.norm.weight` | `blk.{i}.attn_v_norm_in.weight` | BitNet per-projection | +| `layers.{i}.self_attn.o_proj.norm.weight` | `blk.{i}.attn_output_norm_in.weight` | BitNet per-projection | +| `layers.{i}.mlp.gate_proj.weight` | `blk.{i}.ffn_gate.weight` | | +| `layers.{i}.mlp.up_proj.weight` | `blk.{i}.ffn_up.weight` | | +| `layers.{i}.mlp.down_proj.weight` | `blk.{i}.ffn_down.weight` | | +| `layers.{i}.mlp.gate_proj.norm.weight` | `blk.{i}.ffn_gate_norm_in.weight` | BitNet per-projection | +| `layers.{i}.mlp.up_proj.norm.weight` | `blk.{i}.ffn_up_norm_in.weight` | BitNet per-projection | +| `layers.{i}.mlp.down_proj.norm.weight` | `blk.{i}.ffn_down_norm_in.weight` | BitNet per-projection | + +--- + +## 3. Conversion Script + +### `utils/convert-bitnet-embedding-to-gguf.py` + +Standalone conversion script (safetensors → GGUF). Key features: + +- Hardcoded HF→GGUF tensor name mapping (no dependency on llama.cpp's Python converter) +- Supports three output types: + - `--outtype f32`: all weights in float32 + - `--outtype f16`: 2D weights and embeddings as float16, norms as float16 + - `--outtype i2_s`: ternary weights packed in I2_S layout, non-ternary weights as float16 +- Writes `key_length` and `value_length` metadata for head_dim=128 (critical: default calculation would give wrong value 64) +- GPT-2 BPE tokenizer handling with pre-tokenizer hash verification +- Pooling type auto-detection from `modules.json` / `1_Pooling/config.json` (sentence-transformers convention) +- EOS token override: uses `<|endoftext|>` (151643) for correct last-token pooling +- Architecture string: `"qwen3"` + +### I2_S Ternary Packing + +The I2_S format packs ternary weights {-1, 0, +1} into 2-bit representation: + +- Quantization: `scale = 1/mean(|w|)`, `q = round(w * scale).clamp(-1, 1)` +- Encoding: `-1 → 0`, `0 → 1`, `+1 → 2` +- Every 128 values form a block, packed into 32 bytes +- Each byte stores 4 values: `byte = (c0 << 6) | (c1 << 4) | (c2 << 2) | c3` +- Scale (float32) is appended at the end of the packed data buffer + +### Tensor Type Assignment + +| Tensor Type | f16 mode | i2_s mode | +|-------------|----------|-----------| +| 2D linear weights | float16 | I2_S ternary packed | +| Embedding weights | float16 | float16 | +| Norm weights (1D) | float16 | float16 | + +Note: `output.weight` (lm_head) is skipped for embedding models — it is not needed (no token generation). + +--- + +## 4. C++ Modifications (`3rdparty/llama.cpp/src/llama.cpp`) + +### 4.1 New Tensor Enums + +Added 7 new entries after `LLM_TENSOR_FFN_SUB_NORM`: + +```cpp +LLM_TENSOR_ATTN_Q_NORM_IN, +LLM_TENSOR_ATTN_K_NORM_IN, +LLM_TENSOR_ATTN_V_NORM_IN, +LLM_TENSOR_ATTN_OUT_NORM_IN, +LLM_TENSOR_FFN_GATE_NORM_IN, +LLM_TENSOR_FFN_UP_NORM_IN, +LLM_TENSOR_FFN_DOWN_NORM_IN, +``` + +### 4.2 Tensor Name Mappings + +Added to `LLM_ARCH_QWEN3` tensor name map: + +```cpp +{ LLM_TENSOR_ATTN_Q_NORM_IN, "blk.%d.attn_q_norm_in" }, +{ LLM_TENSOR_ATTN_K_NORM_IN, "blk.%d.attn_k_norm_in" }, +{ LLM_TENSOR_ATTN_V_NORM_IN, "blk.%d.attn_v_norm_in" }, +{ LLM_TENSOR_ATTN_OUT_NORM_IN, "blk.%d.attn_output_norm_in" }, +{ LLM_TENSOR_FFN_GATE_NORM_IN, "blk.%d.ffn_gate_norm_in" }, +{ LLM_TENSOR_FFN_UP_NORM_IN, "blk.%d.ffn_up_norm_in" }, +{ LLM_TENSOR_FFN_DOWN_NORM_IN, "blk.%d.ffn_down_norm_in" }, +``` + +### 4.3 Layer Struct Fields + +Added to `struct llama_layer`: + +```cpp +struct ggml_tensor * attn_q_norm_in; +struct ggml_tensor * attn_k_norm_in; +struct ggml_tensor * attn_v_norm_in; +struct ggml_tensor * attn_out_norm_in; +struct ggml_tensor * ffn_gate_norm_in; +struct ggml_tensor * ffn_up_norm_in; +struct ggml_tensor * ffn_down_norm_in; +``` + +### 4.4 load_tensors (LLM_ARCH_QWEN3) + +Added optional loading with `TENSOR_NOT_REQUIRED`: + +```cpp +layer.attn_q_norm_in = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM_IN, "weight", i), {n_embd}, TENSOR_NOT_REQUIRED); +layer.attn_k_norm_in = create_tensor(tn(LLM_TENSOR_ATTN_K_NORM_IN, "weight", i), {n_embd}, TENSOR_NOT_REQUIRED); +layer.attn_v_norm_in = create_tensor(tn(LLM_TENSOR_ATTN_V_NORM_IN, "weight", i), {n_embd}, TENSOR_NOT_REQUIRED); +layer.attn_out_norm_in = create_tensor(tn(LLM_TENSOR_ATTN_OUT_NORM_IN, "weight", i), {n_embd_head_k * n_head}, TENSOR_NOT_REQUIRED); +layer.ffn_gate_norm_in = create_tensor(tn(LLM_TENSOR_FFN_GATE_NORM_IN, "weight", i), {n_embd}, TENSOR_NOT_REQUIRED); +layer.ffn_up_norm_in = create_tensor(tn(LLM_TENSOR_FFN_UP_NORM_IN, "weight", i), {n_embd}, TENSOR_NOT_REQUIRED); +layer.ffn_down_norm_in = create_tensor(tn(LLM_TENSOR_FFN_DOWN_NORM_IN, "weight", i), {n_ff}, TENSOR_NOT_REQUIRED); +``` + +Note: `o_proj.norm` input dimension is `n_embd_head_k * n_head` (=2048), `down_proj.norm` input dimension is `n_ff` (=3072). + +### 4.5 build_qwen3() Graph Modifications + +The `build_qwen3()` function was modified to conditionally apply per-projection RMSNorm. The logic is fully backward compatible — when no `*_norm_in` tensors exist, behavior is identical to original. + +**Attention per-projection norms:** +``` +// Before Q/K/V matmul: +if (layer.attn_q_norm_in) { + cur_q = ggml_rms_norm(ctx, cur, hparams.f_norm_rms_eps); + cur_q = ggml_mul(ctx, cur_q, layer.attn_q_norm_in); +} else { + cur_q = cur; +} +Qcur = ggml_mul_mat(ctx, layer.wq, cur_q); +// Similarly for K, V +``` + +**O_proj norm** requires special handling because `llm_build_kv()` normally applies `wo` internally. Solution: pass `wo=NULL` to `llm_build_kv()`, then apply norm + wo manually: + +``` +cur = llm_build_kv(..., wo=NULL, ...); // returns attention output without o_proj +if (layer.attn_out_norm_in) { + cur = ggml_rms_norm(ctx, cur, hparams.f_norm_rms_eps); + cur = ggml_mul(ctx, cur, layer.attn_out_norm_in); +} +cur = ggml_mul_mat(ctx, layer.wo, cur); +``` + +**FFN per-projection norms:** +``` +// Instead of llm_build_ffn(), manually: +if (layer.ffn_gate_norm_in) { + tmp_gate = rms_norm(cur) * gate_norm_in; +} else { + tmp_gate = cur; +} +tmp_gate = matmul(gate_proj, tmp_gate); +// Similarly for up_proj +tmp = silu(tmp_gate) * tmp_up; + +if (layer.ffn_down_norm_in) { + tmp = rms_norm(tmp) * down_norm_in; +} +cur = matmul(down_proj, tmp); +``` + +--- + +## 5. GGUF Conversion Process + +There are two GGUF files to produce, from **two different source models**: + +| GGUF Output | Source Model | Description | +|-------------|-------------|-------------| +| `embeddings-0.6b-f16.gguf` | `multilingual-e5-0.6b` (standard Qwen3) | F16 baseline, standard float16 weights | +| `bitnet-embeddings-0.6b-f16-i2_s.gguf` | `bitnet-embeddings-0.6b` (BitNet ternary) | I2_S ternary packed weights | + +### 5.1 F16 GGUF: from multilingual-e5-0.6b + +The F16 GGUF is converted from the **standard (non-BitNet) model** `multilingual-e5-0.6b`, which has normal float weights and no per-projection RMSNorm. This uses llama.cpp's standard converter since it is a vanilla Qwen3 model: + +```bash +python3 /path/to/llama.cpp/convert_hf_to_gguf.py \ + /path/to/multilingual-e5-0.6b \ + --outtype f16 \ + --outfile embeddings-0.6b-f16.gguf +``` + +**What happens:** +1. Load `model.safetensors` (standard Qwen3 weights, bfloat16) +2. Convert all 2D weights (projections, embeddings) to float16 +3. Convert norm weights to float32 +4. Write GGUF with `qwen3` architecture metadata and tokenizer + +**Output:** ~1.11 GiB (595.78M params) + +### 5.2 I2_S GGUF: from bitnet-embeddings-0.6b + +The I2_S GGUF is converted from the **BitNet ternary model** `bitnet-embeddings-0.6b`, which has ternary weights {-1, 0, +1} and 7 extra per-projection RMSNorm tensors per layer. This uses the custom converter because the standard llama.cpp converter does not handle per-projection norms or I2_S quantization: + +```bash +python3 utils/convert-bitnet-embedding-to-gguf.py \ + /path/to/bitnet-embeddings-0.6b \ + --outfile bitnet-embeddings-0.6b-f16-i2_s.gguf --outtype i2_s +``` + +**What happens:** +1. Load `model.safetensors` (BitNet ternary weights, bfloat16) +2. Map HF tensor names to GGUF names, including 7 extra `*_norm_in` tensors per layer (see Section 2) +3. For each 2D linear weight (q/k/v/o/gate/up/down projections): + - Compute scale: `scale = 1 / mean(|w|)` + - Quantize: `q = round(w * scale).clamp(-1, 1)` + - Encode: `-1 -> 0`, `0 -> 1`, `+1 -> 2` + - Pack every 128 values into 32 bytes (4 values per byte, 2 bits each) + - Append per-row float32 scale +4. Keep embeddings (`token_embd.weight`) in float16 (not ternary) +5. Keep all norm weights in float16 +6. Skip `output.weight` (lm_head, not needed for embedding models) +7. Write GGUF with `I2_S` type tag for quantized tensors + +**Output:** ~699 MiB (~50% of F16 size) + +### 5.3 Why Two Different Source Models? + +- `multilingual-e5-0.6b` is the **teacher/baseline model** with standard float weights, used as the F16 performance reference +- `bitnet-embeddings-0.6b` is the **1-bit quantized student model** with ternary weights and per-projection BitLinear norms, converted to I2_S for efficient CPU inference +- Benchmarking compares both to measure the throughput gain and quality trade-off of ternary quantization + +### 5.4 Tensor Type Summary + +| Tensor | F16 (from e5-0.6b) | I2_S (from bitnet-0.6b) | +|--------|---------------------|-------------------------| +| Linear projections (q/k/v/o/gate/up/down) | float16 | I2_S (2-bit packed + float32 scale) | +| Embedding (`token_embd.weight`) | float16 | float16 | +| Per-projection norms (`*_norm_in`) | N/A (not present) | float16 | +| Layer norms (`attn_norm`, `ffn_norm`) | float32 | float16 | +| QK head norms (`attn_q_norm`, `attn_k_norm`) | float32 | float32 | +| `output.weight` (lm_head) | present | skipped | + +--- + +## 6. Build and Run + +```bash +# Build with BitNet repo (includes I2_S support) +cmake -S /path/to/BitNet -B build -DCMAKE_BUILD_TYPE=Release +cmake --build build --target llama-embedding llama-bench -j$(nproc) + +# Run embedding inference +build/bin/llama-embedding -m bitnet-embeddings-0.6b-f16-i2_s.gguf \ + -p "hello world" --embd-normalize 2 --embd-output-format array + +# Benchmark: F16 vs I2_S +build/bin/llama-bench -m embeddings-0.6b-f16.gguf \ + -t 8 -p 128,256,512,1024,2048 -n 32,64 -r 3 -ngl 0 + +build/bin/llama-bench -m bitnet-embeddings-0.6b-f16-i2_s.gguf \ + -t 8 -p 128,256,512,1024,2048 -n 32,64 -r 3 -ngl 0 +``` diff --git a/src/ggml-bitnet-lut.cpp b/src/ggml-bitnet-lut.cpp index 59422d548..beef726f7 100644 --- a/src/ggml-bitnet-lut.cpp +++ b/src/ggml-bitnet-lut.cpp @@ -5,9 +5,16 @@ #include #include +#ifdef __x86_64__ +#include +#endif + #include "ggml-bitnet.h" #include "ggml-quants.h" + +#if defined(GGML_BITNET_ARM_TL1) || defined(GGML_BITNET_X86_TL2) #include "bitnet-lut-kernels.h" +#endif #if defined(GGML_BITNET_ARM_TL1) diff --git a/src/ggml-bitnet-mad.cpp b/src/ggml-bitnet-mad.cpp index 4ba9d6509..ad18bac04 100644 --- a/src/ggml-bitnet-mad.cpp +++ b/src/ggml-bitnet-mad.cpp @@ -808,7 +808,7 @@ void ggml_vec_dot_i2_i8_s_Nx1(int n, float * s, size_t bs, const void * vx, size accu[iy] = _mm256_setzero_si256(); } - int8_t * y_col = y + col * by; + const int8_t * y_col = y + col * by; for (int i = 0; i < group32_num; i++) { const uint8_t *px = x + i * 1024; diff --git a/utils/convert-bitnet-embedding-to-gguf.py b/utils/convert-bitnet-embedding-to-gguf.py new file mode 100644 index 000000000..3a4340734 --- /dev/null +++ b/utils/convert-bitnet-embedding-to-gguf.py @@ -0,0 +1,502 @@ +#!/usr/bin/env python3 + +from __future__ import annotations + +import argparse +import json +import logging +import os +import sys +from hashlib import sha256 +from pathlib import Path +from typing import Any, Iterator + +import numpy as np +import torch + +# Allow using the local gguf-py if present +if "NO_LOCAL_GGUF" not in os.environ: + _local_gguf = Path(__file__).parent / "gguf-py" + if _local_gguf.exists(): + sys.path.insert(1, str(_local_gguf)) +import gguf + +logger = logging.getLogger("convert-bitnet-embedding") + +# --------------------------------------------------------------------------- +# Tensor name mapping: HuggingFace -> GGUF +# --------------------------------------------------------------------------- + +def build_tensor_name_map(n_layers: int) -> dict[str, str]: + """Build HF tensor name -> GGUF tensor name mapping.""" + mapping: dict[str, str] = { + "embed_tokens.weight": "token_embd.weight", + "norm.weight": "output_norm.weight", + } + + for i in range(n_layers): + pfx = f"layers.{i}" + blk = f"blk.{i}" + + mapping.update({ + # Layer norms + f"{pfx}.input_layernorm.weight": f"{blk}.attn_norm.weight", + f"{pfx}.post_attention_layernorm.weight": f"{blk}.ffn_norm.weight", + + # Self-attention projections + f"{pfx}.self_attn.q_proj.weight": f"{blk}.attn_q.weight", + f"{pfx}.self_attn.k_proj.weight": f"{blk}.attn_k.weight", + f"{pfx}.self_attn.v_proj.weight": f"{blk}.attn_v.weight", + f"{pfx}.self_attn.o_proj.weight": f"{blk}.attn_output.weight", + + # QK head norms (standard Qwen3) + f"{pfx}.self_attn.q_norm.weight": f"{blk}.attn_q_norm.weight", + f"{pfx}.self_attn.k_norm.weight": f"{blk}.attn_k_norm.weight", + + # Per-projection input norms (BitNet-specific) + f"{pfx}.self_attn.q_proj.norm.weight": f"{blk}.attn_q_norm_in.weight", + f"{pfx}.self_attn.k_proj.norm.weight": f"{blk}.attn_k_norm_in.weight", + f"{pfx}.self_attn.v_proj.norm.weight": f"{blk}.attn_v_norm_in.weight", + f"{pfx}.self_attn.o_proj.norm.weight": f"{blk}.attn_output_norm_in.weight", + + # MLP projections + f"{pfx}.mlp.gate_proj.weight": f"{blk}.ffn_gate.weight", + f"{pfx}.mlp.up_proj.weight": f"{blk}.ffn_up.weight", + f"{pfx}.mlp.down_proj.weight": f"{blk}.ffn_down.weight", + + # Per-projection input norms for MLP (BitNet-specific) + f"{pfx}.mlp.gate_proj.norm.weight": f"{blk}.ffn_gate_norm_in.weight", + f"{pfx}.mlp.up_proj.norm.weight": f"{blk}.ffn_up_norm_in.weight", + f"{pfx}.mlp.down_proj.norm.weight": f"{blk}.ffn_down_norm_in.weight", + }) + + return mapping + + +# --------------------------------------------------------------------------- +# Tokenizer handling (GPT-2 / BPE for Qwen3) +# --------------------------------------------------------------------------- + +def get_vocab_base_pre(tokenizer) -> str: + # encoding this string and hashing the resulting tokens would (hopefully) give us a unique identifier that + # is specific for the BPE pre-tokenizer used by the model + # we will use this unique identifier to write a "tokenizer.ggml.pre" entry in the GGUF file which we can + # use in llama.cpp to implement the same pre-tokenizer + + chktxt = '\n \n\n \n\n\n \t \t\t \t\n \n \n \n \n\U0001f680 (normal) \U0001f636‍\U0001f32b️ (multiple emojis concatenated) ✅ \U0001f999\U0001f999 3 33 333 3333 33333 333333 3333333 33333333 3.3 3..3 3...3 កាន់តែពិសេសអាច\U0001f601 ?我想在apple工作1314151天～ ------======= нещо на Български \'\'\'\'\'\'```````""""......!!!!!!?????? I\'ve been \'told he\'s there, \'RE you sure? \'M not sure I\'ll make it, \'D you like some tea? We\'Ve a\'lL' + + chktok = tokenizer.encode(chktxt) + chkhsh = sha256(str(chktok).encode()).hexdigest() + + logger.debug(f"chktok: {chktok}") + logger.debug(f"chkhsh: {chkhsh}") + + res = None + + # NOTE: if you get an error here, you need to update the convert_hf_to_gguf_update.py script + # or pull the latest version of the model from Huggingface + # don't edit the hashes manually! + if chkhsh == "0ef9807a4087ebef797fc749390439009c3b9eda9ad1a097abbe738f486c01e5": + # ref: https://huggingface.co/meta-llama/Meta-Llama-3-8B + res = "llama-bpe" + if chkhsh == "049ecf7629871e3041641907f3de7c733e4dbfdc736f57d882ba0b0845599754": + # ref: https://huggingface.co/deepseek-ai/deepseek-llm-7b-base + res = "deepseek-llm" + if chkhsh == "347715f544604f9118bb75ed199f68779f423cabb20db6de6f31b908d04d7821": + # ref: https://huggingface.co/deepseek-ai/deepseek-coder-6.7b-base + res = "deepseek-coder" + if chkhsh == "8aeee3860c56296a157a1fe2fad249ec40aa59b1bb5709f4ade11c4e6fe652ed": + # ref: https://huggingface.co/tiiuae/falcon-7b + res = "falcon" + if chkhsh == "3ce83efda5659b07b1ad37ca97ca5797ea4285d9b9ab0dc679e4a720c9da7454": + # ref: https://huggingface.co/openai-community/gpt2 + res = "gpt-2" + if chkhsh == "d4540891389ea895b53b399da6ac824becc30f2fba0e9ddbb98f92e55ca0e97c": + # ref: https://huggingface.co/Qwen/Qwen3-Embedding-0.6B + res = "qwen2" + + if res is None: + logger.warning("\n") + logger.warning("**************************************************************************************") + logger.warning("** WARNING: The BPE pre-tokenizer was not recognized!") + logger.warning("** There are 2 possible reasons for this:") + logger.warning("** - the model has not been added to convert_hf_to_gguf_update.py yet") + logger.warning("** - the pre-tokenization config has changed upstream") + logger.warning("** Check your model files and convert_hf_to_gguf_update.py and update them accordingly.") + logger.warning("** ref: https://github.com/ggml-org/llama.cpp/pull/6920") + logger.warning("**") + logger.warning(f"** chkhsh: {chkhsh}") + logger.warning("**************************************************************************************") + logger.warning("\n") + raise NotImplementedError("BPE pre-tokenizer was not recognized - update get_vocab_base_pre()") + + logger.debug(f"tokenizer.ggml.pre: {repr(res)}") + logger.debug(f"chkhsh: {chkhsh}") + + return res + + +def _does_token_look_special(token: str) -> bool: + """Check if a token looks like a special token (e.g., <|...|>, <...>).""" + if not token: + return False + # Matches patterns like <|endoftext|>, , , [CLS], [SEP], etc. + if token.startswith(("<|", "<", "[")) and token.endswith(("|>", ">", "]")): + return True + return False + + +def set_vocab(gguf_writer: gguf.GGUFWriter, dir_model: Path, hparams: dict): + """Set GPT-2 BPE vocab for Qwen3.""" + from transformers import AutoTokenizer + tokenizer = AutoTokenizer.from_pretrained(dir_model) + vocab_size = hparams.get("vocab_size", len(tokenizer.vocab)) + + tokpre = get_vocab_base_pre(tokenizer) + + tokens: list[str] = [] + toktypes: list[int] = [] + + reverse_vocab = {id_: tok for tok, id_ in tokenizer.vocab.items()} + added_vocab = tokenizer.get_added_vocab() + + added_tokens_decoder = tokenizer.added_tokens_decoder + + for i in range(vocab_size): + if i not in reverse_vocab: + tokens.append(f"[PAD{i}]") + toktypes.append(gguf.TokenType.UNUSED) + elif reverse_vocab[i] in added_vocab: + token = reverse_vocab[i] + + # Only encode-decode non-normalized tokens (matching llama.cpp upstream) + if not added_tokens_decoder[i].normalized: + token = tokenizer.decode(tokenizer.encode(token, add_special_tokens=False)) + + if added_tokens_decoder[i].special or _does_token_look_special(token): + toktypes.append(gguf.TokenType.CONTROL) + else: + # Pre-normalize user-defined spaces (for Gemma-style tokenizers) + token = token.replace(b"\xe2\x96\x81".decode("utf-8"), " ") + toktypes.append(gguf.TokenType.USER_DEFINED) + + tokens.append(token) + else: + tokens.append(reverse_vocab[i]) + toktypes.append(gguf.TokenType.NORMAL) + + gguf_writer.add_tokenizer_model("gpt2") + gguf_writer.add_tokenizer_pre(tokpre) + gguf_writer.add_token_list(tokens) + gguf_writer.add_token_types(toktypes) + + special_vocab = gguf.SpecialVocab(dir_model, load_merges=True) + # Override EOS token: PyTorch tokenizer appends <|endoftext|> (151643) as the + # sentence-end marker, not <|im_end|> (151645). For last-token pooling to work + # correctly, llama.cpp must append the same token. + special_vocab.special_token_ids["eos"] = 151643 + special_vocab.add_to_gguf(gguf_writer) + + # Embedding models need EOS token appended for last-token pooling + gguf_writer.add_add_eos_token(True) + + +# --------------------------------------------------------------------------- +# GGUF metadata +# --------------------------------------------------------------------------- + +def set_gguf_parameters(gguf_writer: gguf.GGUFWriter, hparams: dict, dir_model: Path, ftype: int): + gguf_writer.add_name(dir_model.name) + + n_layers = hparams["num_hidden_layers"] + n_embd = hparams["hidden_size"] + n_head = hparams["num_attention_heads"] + n_head_kv = hparams.get("num_key_value_heads", n_head) + n_ff = hparams["intermediate_size"] + + gguf_writer.add_block_count(n_layers) + gguf_writer.add_context_length(hparams.get("max_position_embeddings", 32768)) + gguf_writer.add_embedding_length(n_embd) + gguf_writer.add_feed_forward_length(n_ff) + gguf_writer.add_head_count(n_head) + gguf_writer.add_head_count_kv(n_head_kv) + gguf_writer.add_vocab_size(hparams["vocab_size"]) + + head_dim = hparams.get("head_dim", n_embd // n_head) + gguf_writer.add_rope_dimension_count(head_dim) + gguf_writer.add_key_length(head_dim) + gguf_writer.add_value_length(head_dim) + + if hparams.get("rope_theta") is not None: + gguf_writer.add_rope_freq_base(hparams["rope_theta"]) + if hparams.get("rms_norm_eps") is not None: + gguf_writer.add_layer_norm_rms_eps(hparams["rms_norm_eps"]) + + gguf_writer.add_file_type(ftype) + + # Pooling type for embedding models + # Try to read from modules.json / 1_Pooling/config.json (sentence-transformers convention) + pooling_type = None + module_path = dir_model / "modules.json" + if module_path.is_file(): + with open(module_path, encoding="utf-8") as f: + modules = json.load(f) + for mod in modules: + if mod["type"].endswith("Pooling"): + pooling_path = dir_model / mod["path"] / "config.json" + if pooling_path.is_file(): + with open(pooling_path, encoding="utf-8") as f: + pooling = json.load(f) + if pooling.get("pooling_mode_mean_tokens"): + pooling_type = gguf.PoolingType.MEAN + elif pooling.get("pooling_mode_cls_token"): + pooling_type = gguf.PoolingType.CLS + elif pooling.get("pooling_mode_lasttoken"): + pooling_type = gguf.PoolingType.LAST + break + if pooling_type is None: + # Default to MEAN pooling for embedding models + logger.info(" No pooling config found, defaulting to MEAN pooling") + pooling_type = gguf.PoolingType.MEAN + gguf_writer.add_pooling_type(pooling_type) + + logger.info(f" n_layers={n_layers}, n_embd={n_embd}, n_head={n_head}, n_head_kv={n_head_kv}, n_ff={n_ff}") + + +# --------------------------------------------------------------------------- +# Tensor iteration from safetensors +# --------------------------------------------------------------------------- + +def iter_tensors(dir_model: Path) -> Iterator[tuple[str, torch.Tensor]]: + """Yield (name, tensor) from safetensors files.""" + from safetensors import safe_open + + safetensor_files = sorted(dir_model.glob("*.safetensors")) + if not safetensor_files: + raise FileNotFoundError(f"No .safetensors files in {dir_model}") + + for sf_path in safetensor_files: + logger.info(f"Loading {sf_path.name}") + with safe_open(str(sf_path), framework="pt", device="cpu") as f: + for name in f.keys(): + yield name, f.get_tensor(name) + + +# --------------------------------------------------------------------------- +# I2_S ternary packing (platform-independent) +# --------------------------------------------------------------------------- +# +# I2_S format (from dequantize_row_i2_s in ggml-quants.c): +# - Every 128 values form a block, packed into 32 bytes +# - Each byte stores 4 values at positions [0*32+gp, 1*32+gp, 2*32+gp, 3*32+gp] +# where gp is the byte index within the 32-byte group +# - Encoding per byte: c0=(b>>6)&3, c1=(b>>4)&3, c2=(b>>2)&3, c3=(b>>0)&3 +# - Value mapping: 0 -> -1, 1 -> 0, 2 -> +1, 3 -> 0 +# - Scale is stored as a separate tensor (tensor_name + "_scale") + +def quantize_to_i2_s(w: np.ndarray) -> np.ndarray: + """Quantize float weights to ternary and pack into I2_S layout. + + Uses the same quantization as BitLinear weight_quant_minmax(): + scale = 1.0 / mean(|w|) + q = round(w * scale).clamp(-1, 1) + dequant = q / scale = q * mean(|w|) + + The I2_S format is self-contained: packed ternary bytes followed by a f32 scale + appended at the end of the data buffer. + + Args: + w: float weight tensor of shape (M, K) + + Returns: + packed_data: uint8 array containing I2_S packed bytes + scale (as 4 trailing bytes) + """ + M, K = w.shape + n = M * K + w_flat = w.flatten().astype(np.float32) + + # BitLinear weight_quant_minmax: scale = 1/mean(|w|), then round & clamp + abs_mean = np.mean(np.abs(w_flat)) + abs_mean = max(abs_mean, 1e-5) + inv_scale = 1.0 / abs_mean + q_float = np.round(w_flat * inv_scale).clip(-1, 1) # ternary: {-1, 0, 1} + + # scale for dequantization = abs_mean (i.e., dequant = q * abs_mean) + scale = np.float32(abs_mean) + + # Map ternary {-1, 0, 1} -> I2_S encoding {0, 1, 2} + # -1 -> 0, 0 -> 1, +1 -> 2 + q = np.ones(n, dtype=np.uint8) # default to 1 (zero) + q[q_float > 0.5] = 2 # +1 -> 2 + q[q_float < -0.5] = 0 # -1 -> 0 + + # Pack into I2_S layout: 128-value blocks, interleaved into 32 bytes + # Pad to multiple of 128 + pad_len = (128 - n % 128) % 128 + if pad_len: + q = np.pad(q, (0, pad_len), constant_values=1) + + n_padded = len(q) + n_blocks = n_padded // 128 + + q = q.reshape(n_blocks, 4, 32) + + # Pack: byte = (c0 << 6) | (c1 << 4) | (c2 << 2) | c3 + packed = (q[:, 0, :].astype(np.uint8) << 6) | \ + (q[:, 1, :].astype(np.uint8) << 4) | \ + (q[:, 2, :].astype(np.uint8) << 2) | \ + (q[:, 3, :].astype(np.uint8)) + + packed = packed.reshape(-1).astype(np.uint8) + + # I2_S format: packed_bytes + 32-byte aligned tail (scale in first 4 bytes of tail) + # Total size = n_elements / 4 + 32 (as defined in ggml.c) + packed_size = n // 4 + total_size = packed_size + 32 + result = np.zeros(total_size, dtype=np.uint8) + result[:len(packed)] = packed[:packed_size] + # Write scale as float32 at offset packed_size + result[packed_size:packed_size+4] = np.frombuffer(scale.tobytes(), dtype=np.uint8) + + return result + + +# --------------------------------------------------------------------------- +# Main conversion +# --------------------------------------------------------------------------- + +def main(): + parser = argparse.ArgumentParser(description="Convert bitnet-embeddings to GGUF") + parser.add_argument("model", type=Path, help="Model directory") + parser.add_argument("--outfile", type=Path, default=None, help="Output GGUF file") + parser.add_argument("--outtype", choices=["f32", "f16", "i2_s"], default="f16", + help="Output type: f32, f16, or i2_s (ternary quantized)") + parser.add_argument("--verbose", action="store_true") + args = parser.parse_args() + + logging.basicConfig(level=logging.DEBUG if args.verbose else logging.INFO) + + dir_model = args.model + if not dir_model.is_dir(): + logger.error(f"{dir_model} is not a directory") + sys.exit(1) + + # Default output filename + if args.outfile is None: + suffix = {"f32": "-f32", "f16": "-f16", "i2_s": "-f16-new-i2_s"}[args.outtype] + args.outfile = dir_model / f"{dir_model.name}{suffix}.gguf" + + # Load config + with open(dir_model / "config.json") as f: + hparams = json.load(f) + + arch = hparams.get("model_type", "qwen3") + assert arch == "qwen3", f"Expected qwen3 architecture, got {arch}" + + n_layers = hparams["num_hidden_layers"] + + # Determine ftype + if args.outtype == "f32": + ftype = 0 # GGML F32 + elif args.outtype == "f16": + ftype = 1 # GGML F16 + else: # i2_s + ftype = 40 # LLAMA_FTYPE_MOSTLY_I2_S + + logger.info(f"Converting {dir_model.name} to GGUF ({args.outtype})") + + # Create GGUF writer + gguf_writer = gguf.GGUFWriter(str(args.outfile), "qwen3") + + # Set parameters + set_gguf_parameters(gguf_writer, hparams, dir_model, ftype) + + # Set vocab + logger.info("Setting tokenizer/vocab...") + set_vocab(gguf_writer, dir_model, hparams) + + # Build tensor name map + tensor_map = build_tensor_name_map(n_layers) + + # Process tensors + logger.info("Processing tensors...") + tensor_count = 0 + for hf_name, data_torch in iter_tensors(dir_model): + # Skip tensors we don't need + if hf_name.endswith((".attention.masked_bias", ".attention.bias", ".rotary_emb.inv_freq")): + continue + + # Strip "model." prefix if present + name = hf_name + if name.startswith("model."): + name = name[len("model."):] + + # Look up GGUF name + gguf_name = tensor_map.get(name) + if gguf_name is None: + logger.warning(f"Skipping unmapped tensor: {hf_name}") + continue + + old_dtype = data_torch.dtype + + # Convert bf16 -> f32 first (bf16 not directly supported by gguf) + if data_torch.dtype == torch.bfloat16: + data_torch = data_torch.to(torch.float32) + + data = data_torch.squeeze().numpy() + n_dims = len(data.shape) + data_shape = data.shape + + # Determine if this is a linear weight suitable for ternary quantization + is_norm = gguf_name.endswith("_norm.weight") or gguf_name.endswith("_norm_in.weight") + is_embed = gguf_name == "token_embd.weight" + is_linear_weight = n_dims == 2 and not is_norm and not is_embed + suit_i2 = is_linear_weight + + if args.outtype == "i2_s" and suit_i2: + # --- I2_S ternary packing (scale embedded in data) --- + packed = quantize_to_i2_s(data) + data_qtype = gguf.GGMLQuantizationType.I2_S + + shape_str = f"{{{', '.join(str(n) for n in reversed(data_shape))}}}" + logger.info(f" {gguf_name}: {list(data_shape)} {old_dtype} -> I2_S, shape = {shape_str}") + + gguf_writer.add_tensor(gguf_name, packed, raw_shape=data_shape, raw_dtype=data_qtype) + tensor_count += 1 + + elif args.outtype in ("f16", "i2_s") and (is_linear_weight or is_embed): + # 2D weight tensors (linear + embedding) -> f16 + data = data.astype(np.float16) + logger.info(f" {gguf_name}: {list(data_torch.shape)} {old_dtype} -> float16") + gguf_writer.add_tensor(gguf_name, data) + tensor_count += 1 + + else: + # norms, 1D tensors + if args.outtype in ("f16", "i2_s"): + data = data.astype(np.float16) + logger.info(f" {gguf_name}: {list(data_torch.shape)} {old_dtype} -> float16") + else: + if data.dtype != np.float32: + data = data.astype(np.float32) + logger.info(f" {gguf_name}: {list(data_torch.shape)} {old_dtype} -> float32") + gguf_writer.add_tensor(gguf_name, data) + tensor_count += 1 + + logger.info(f"Total tensors written: {tensor_count}") + + # Note: output.weight (lm_head) is skipped for embedding models — + # it is not needed (no token generation) and saves ~297MB for this model. + + # Write GGUF + logger.info(f"Writing to {args.outfile}...") + gguf_writer.write_header_to_file() + gguf_writer.write_kv_data_to_file() + gguf_writer.write_tensors_to_file() + gguf_writer.close() + + logger.info("Done!") + + +if __name__ == "__main__": + main() From 3b04140a54c4e5e14183c607121230b2d6da85dc Mon Sep 17 00:00:00 2001 From: isHuangXin Date: Mon, 13 Jul 2026 06:26:16 +0200 Subject: [PATCH 2/8] feat: add I2_S GGUF conversion for bitnet-b1.58-2B-4T and refactor T-MAC LUT path - Add quantize_to_i2_s() for direct ternary-to-I2_S packing in conversion script - Support offline-quantized models (uint8 packed weights + weight_scale) - Fix weight_quant double-quantization bug for offline-quantized models - Fix I2_S scale computation to use first nonzero absolute value - Add I2_S ftype mapping and BitNetForCausalLM registration - Refactor ggml-bitnet-lut T-MAC wrapper with proper mul_mat implementation - Update llama.cpp submodule with I2_S ftype and 2B model type support --- 3rdparty/llama.cpp | 2 +- include/bitnet-lut-kernels.h | 1171 ++++++++++++++++++++++++++++ include/ggml-bitnet.h | 3 + include/kernel_config.ini | 21 + src/ggml-bitnet-lut.cpp | 90 ++- utils/convert-hf-to-gguf-bitnet.py | 163 +++- 6 files changed, 1395 insertions(+), 55 deletions(-) create mode 100644 include/bitnet-lut-kernels.h create mode 100644 include/kernel_config.ini diff --git a/3rdparty/llama.cpp b/3rdparty/llama.cpp index 078c16664..355c0c4d1 160000 --- a/3rdparty/llama.cpp +++ b/3rdparty/llama.cpp @@ -1 +1 @@ -Subproject commit 078c166647bf0f96964dc2614becefbdf9ca2207 +Subproject commit 355c0c4d1870f0679334cd73727d1e05c1863fbd diff --git a/include/bitnet-lut-kernels.h b/include/bitnet-lut-kernels.h new file mode 100644 index 000000000..f7fb5625f --- /dev/null +++ b/include/bitnet-lut-kernels.h @@ -0,0 +1,1171 @@ +#if defined(GGML_BITNET_X86_TL2) +#include "ggml-bitnet.h" +#include +#include +#define GGML_BITNET_MAX_NODES 8192 +static bool initialized = false; +static bitnet_tensor_extra * bitnet_tensor_extras = nullptr; +static size_t bitnet_tensor_extras_index = 0; +static void * aligned_malloc(size_t size) { +#if defined(_WIN32) + return _aligned_malloc(size, 64); +#else + void * ptr = nullptr; + posix_memalign(&ptr, 64, size); + return ptr; +#endif +} + +static void aligned_free(void * ptr) { +#if defined(_WIN32) + _aligned_free(ptr); +#else + free(ptr); +#endif +} +#define BK2 32 +#if defined __AVX2__ +inline void _mm256_merge_epi32(const __m256i v0, const __m256i v1, __m256i *vl, __m256i *vh) +{ + __m256i va = _mm256_permute4x64_epi64(v0, _MM_SHUFFLE(3, 1, 2, 0)); + __m256i vb = _mm256_permute4x64_epi64(v1, _MM_SHUFFLE(3, 1, 2, 0)); + *vl = _mm256_unpacklo_epi32(va, vb); + *vh = _mm256_unpackhi_epi32(va, vb); +} +inline void _mm256_merge_epi64(const __m256i v0, const __m256i v1, __m256i *vl, __m256i *vh) +{ + __m256i va = _mm256_permute4x64_epi64(v0, _MM_SHUFFLE(3, 1, 2, 0)); + __m256i vb = _mm256_permute4x64_epi64(v1, _MM_SHUFFLE(3, 1, 2, 0)); + *vl = _mm256_unpacklo_epi64(va, vb); + *vh = _mm256_unpackhi_epi64(va, vb); +} +inline void _mm256_merge_si128(const __m256i v0, const __m256i v1, __m256i *vl, __m256i *vh) +{ + *vl = _mm256_permute2x128_si256(v0, v1, _MM_SHUFFLE(0, 2, 0, 0)); + *vh = _mm256_permute2x128_si256(v0, v1, _MM_SHUFFLE(0, 3, 0, 1)); +} +inline void Transpose_8_8( + __m256i *v0, + __m256i *v1, + __m256i *v2, + __m256i *v3, + __m256i *v4, + __m256i *v5, + __m256i *v6, + __m256i *v7) +{ + __m256i w0, w1, w2, w3, w4, w5, w6, w7; + __m256i x0, x1, x2, x3, x4, x5, x6, x7; + _mm256_merge_epi32(*v0, *v1, &w0, &w1); + _mm256_merge_epi32(*v2, *v3, &w2, &w3); + _mm256_merge_epi32(*v4, *v5, &w4, &w5); + _mm256_merge_epi32(*v6, *v7, &w6, &w7); + _mm256_merge_epi64(w0, w2, &x0, &x1); + _mm256_merge_epi64(w1, w3, &x2, &x3); + _mm256_merge_epi64(w4, w6, &x4, &x5); + _mm256_merge_epi64(w5, w7, &x6, &x7); + _mm256_merge_si128(x0, x4, v0, v1); + _mm256_merge_si128(x1, x5, v2, v3); + _mm256_merge_si128(x2, x6, v4, v5); + _mm256_merge_si128(x3, x7, v6, v7); +} +#endif +inline int32_t per_tensor_quant(int k, void* lut_scales_, void* b_) { + bitnet_float_type* lut_scales = (bitnet_float_type*)lut_scales_; + bitnet_float_type* b = (bitnet_float_type*)b_; +#if defined __AVX2__ + __m256 max_vec = _mm256_set1_ps(0.f); + const __m256 vec_sign = _mm256_set1_ps(-0.0f); + for (int i = 0; i < k / 8; i++) { + __m256 vec_b = _mm256_loadu_ps(b + i * 8); + __m256 vec_babs = _mm256_andnot_ps(vec_sign, vec_b); + max_vec = _mm256_max_ps(vec_babs, max_vec); + } + __m128 max1 = _mm_max_ps(_mm256_extractf128_ps(max_vec, 1), _mm256_castps256_ps128(max_vec)); + max1 = _mm_max_ps(max1, _mm_movehl_ps(max1, max1)); + max1 = _mm_max_ss(max1, _mm_movehdup_ps(max1)); + float scales = 127 / _mm_cvtss_f32(max1); + *lut_scales = scales; +#endif + return 0; +} +inline int32_t partial_max_reset(int32_t bs, void* lut_scales_) { + bitnet_float_type* lut_scales = (bitnet_float_type*)lut_scales_; + #pragma unroll + for (int i=0; i< bs; i++) { + lut_scales[i] = 0.0; + } + return 0; +} +template +inline int32_t three_lut_ctor(int8_t* qlut, bitnet_float_type* b, bitnet_float_type* lut_scales) { +#if defined __AVX2__ + __m256i vec_lut[16]; + const __m256i vec_bi = _mm256_set_epi32(84, 72, 60, 48, 36, 24, 12, 0); + float scales = *lut_scales; + __m256i shuffle_mask = _mm256_set_epi8( + 0x0f, 0x0d, 0x0b, 0x09, 0x07, 0x05, 0x03, 0x01, + 0x0e, 0x0c, 0x0a, 0x08, 0x06, 0x04, 0x02, 0x00, + 0x0f, 0x0d, 0x0b, 0x09, 0x07, 0x05, 0x03, 0x01, + 0x0e, 0x0c, 0x0a, 0x08, 0x06, 0x04, 0x02, 0x00 + ); +#pragma unroll + for (int k = 0; k < act_k / 24; ++k) { + __m256 vec_b0 = _mm256_i32gather_ps(b + k * 24 + 0, vec_bi, 1); + __m256 vec_b1 = _mm256_i32gather_ps(b + k * 24 + 1, vec_bi, 1); + __m256 vec_b2 = _mm256_i32gather_ps(b + k * 24 + 2, vec_bi, 1); + + __m256i vec_b0i = _mm256_cvtps_epi32(_mm256_round_ps(_mm256_mul_ps(vec_b0, _mm256_set1_ps(scales)), _MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC)); + __m256i vec_b1i = _mm256_cvtps_epi32(_mm256_round_ps(_mm256_mul_ps(vec_b1, _mm256_set1_ps(scales)), _MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC)); + __m256i vec_b2i = _mm256_cvtps_epi32(_mm256_round_ps(_mm256_mul_ps(vec_b2, _mm256_set1_ps(scales)), _MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC)); + + vec_lut[15] = _mm256_setzero_si256(); + vec_lut[14] = _mm256_setzero_si256(); + vec_lut[13] = vec_b0i; + vec_lut[13] = _mm256_add_epi32(vec_lut[13], vec_b1i); + vec_lut[13] = _mm256_add_epi32(vec_lut[13], vec_b2i); + vec_lut[12] = vec_b0i; + vec_lut[12] = _mm256_add_epi32(vec_lut[12], vec_b1i); + vec_lut[11] = vec_b0i; + vec_lut[11] = _mm256_add_epi32(vec_lut[11], vec_b1i); + vec_lut[11] = _mm256_sub_epi32(vec_lut[11], vec_b2i); + vec_lut[10] = vec_b0i; + vec_lut[10] = _mm256_add_epi32(vec_lut[10], vec_b2i); + vec_lut[9] = vec_b0i; + vec_lut[8] = vec_b0i; + vec_lut[8] = _mm256_sub_epi32(vec_lut[8], vec_b2i); + vec_lut[7] = vec_b0i; + vec_lut[7] = _mm256_sub_epi32(vec_lut[7], vec_b1i); + vec_lut[7] = _mm256_add_epi32(vec_lut[7], vec_b2i); + vec_lut[6] = vec_b0i; + vec_lut[6] = _mm256_sub_epi32(vec_lut[6], vec_b1i); + vec_lut[5] = vec_b0i; + vec_lut[5] = _mm256_sub_epi32(vec_lut[5], vec_b1i); + vec_lut[5] = _mm256_sub_epi32(vec_lut[5], vec_b2i); + vec_lut[4] = vec_b1i; + vec_lut[4] = _mm256_add_epi32(vec_lut[4], vec_b2i); + vec_lut[3] = vec_b1i; + vec_lut[2] = vec_b1i; + vec_lut[2] = _mm256_sub_epi32(vec_lut[2], vec_b2i); + vec_lut[1] = vec_b2i; + vec_lut[0] = _mm256_setzero_si256(); + __m256i ix[16]; + +#pragma unroll + for (int g = 0; g < 16; ++g) { + ix[g] = vec_lut[g]; + } + + Transpose_8_8(&(ix[0]), &(ix[1]), &(ix[2]), &(ix[3]), &(ix[4]), &(ix[5]),&(ix[6]), &(ix[7])); + Transpose_8_8(&(ix[8]), &(ix[9]), &(ix[10]), &(ix[11]), &(ix[12]), &(ix[13]),&(ix[14]), &(ix[15])); + +#pragma unroll + for (int g = 0; g < 8; ++g) { + ix[g] = _mm256_packs_epi32(ix[g], ix[g + 8]); + ix[g] = _mm256_permute4x64_epi64(ix[g], _MM_SHUFFLE(3, 1, 2, 0)); + ix[g] = _mm256_shuffle_epi8(ix[g], shuffle_mask); + ix[g] = _mm256_permute4x64_epi64(ix[g], _MM_SHUFFLE(3, 1, 2, 0)); + } + int8_t* qlut_i8 = reinterpret_cast(qlut); + _mm256_storeu_si256(reinterpret_cast<__m256i*>(qlut_i8 + k * 256 + 0 * 32 + 0), ix[0]); + _mm256_storeu_si256(reinterpret_cast<__m256i*>(qlut_i8 + k * 256 + 1 * 32 + 0), ix[1]); + _mm256_storeu_si256(reinterpret_cast<__m256i*>(qlut_i8 + k * 256 + 2 * 32 + 0), ix[2]); + _mm256_storeu_si256(reinterpret_cast<__m256i*>(qlut_i8 + k * 256 + 3 * 32 + 0), ix[3]); + _mm256_storeu_si256(reinterpret_cast<__m256i*>(qlut_i8 + k * 256 + 4 * 32 + 0), ix[4]); + _mm256_storeu_si256(reinterpret_cast<__m256i*>(qlut_i8 + k * 256 + 5 * 32 + 0), ix[5]); + _mm256_storeu_si256(reinterpret_cast<__m256i*>(qlut_i8 + k * 256 + 6 * 32 + 0), ix[6]); + _mm256_storeu_si256(reinterpret_cast<__m256i*>(qlut_i8 + k * 256 + 7 * 32 + 0), ix[7]); + + } + + *lut_scales = scales; +#endif + return 0; +} + +template +inline int32_t two_lut_ctor(int8_t* qlut, bitnet_float_type* b, bitnet_float_type* lut_scales) { +#if defined __AVX2__ + __m256i vec_lut[16]; + const __m256i vec_bi = _mm256_set_epi32(56, 48, 40, 32, 24, 16, 8, 0); + float scales = *lut_scales; + __m256i shuffle_mask = _mm256_set_epi8( + 0x0f, 0x0d, 0x0b, 0x09, 0x07, 0x05, 0x03, 0x01, + 0x0e, 0x0c, 0x0a, 0x08, 0x06, 0x04, 0x02, 0x00, + 0x0f, 0x0d, 0x0b, 0x09, 0x07, 0x05, 0x03, 0x01, + 0x0e, 0x0c, 0x0a, 0x08, 0x06, 0x04, 0x02, 0x00 + ); +#pragma unroll + for (int k = 0; k < act_k / 16; ++k) { + __m256 vec_b0f = _mm256_i32gather_ps(b + k * 16 + 0, vec_bi, 1); + __m256 vec_b1f = _mm256_i32gather_ps(b + k * 16 + 1, vec_bi, 1); + + __m256i vec_b0 = _mm256_cvtps_epi32(_mm256_round_ps(_mm256_mul_ps(vec_b0f, _mm256_set1_ps(scales)), _MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC)); + __m256i vec_b1 = _mm256_cvtps_epi32(_mm256_round_ps(_mm256_mul_ps(vec_b1f, _mm256_set1_ps(scales)), _MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC)); + vec_lut[15] = _mm256_setzero_si256(); + vec_lut[14] = _mm256_setzero_si256(); + vec_lut[13] = _mm256_setzero_si256(); + vec_lut[12] = _mm256_setzero_si256(); + vec_lut[11] = _mm256_setzero_si256(); + vec_lut[10] = _mm256_setzero_si256(); + vec_lut[9] = _mm256_setzero_si256(); + vec_lut[8] = vec_b0; + vec_lut[8] = _mm256_add_epi32(vec_lut[8], vec_b1); + vec_lut[7] = vec_b0; + vec_lut[6] = vec_b0; + vec_lut[6] = _mm256_sub_epi32(vec_lut[6], vec_b1); + vec_lut[5] = vec_b1; + vec_lut[4] = _mm256_setzero_si256(); + vec_lut[3] = _mm256_setzero_si256(); + vec_lut[3] = _mm256_sub_epi32(vec_lut[3], vec_b1); + vec_lut[2] = _mm256_setzero_si256(); + vec_lut[2] = _mm256_sub_epi32(vec_lut[2], vec_b0); + vec_lut[2] = _mm256_add_epi32(vec_lut[2], vec_b1); + vec_lut[1] = _mm256_setzero_si256(); + vec_lut[1] = _mm256_sub_epi32(vec_lut[1], vec_b0); + vec_lut[0] = _mm256_setzero_si256(); + vec_lut[0] = _mm256_sub_epi32(vec_lut[0], vec_b0); + vec_lut[0] = _mm256_sub_epi32(vec_lut[0], vec_b1); + + __m256i ix[16]; +#pragma unroll + for (int g = 0; g < 16; ++g) { + ix[g] = vec_lut[g]; + } + + Transpose_8_8(&(ix[0]), &(ix[1]), &(ix[2]), &(ix[3]), &(ix[4]), &(ix[5]),&(ix[6]), &(ix[7])); + Transpose_8_8(&(ix[8]), &(ix[9]), &(ix[10]), &(ix[11]), &(ix[12]), &(ix[13]),&(ix[14]), &(ix[15])); + +#pragma unroll + for (int g = 0; g < 8; ++g) { + ix[g] = _mm256_packs_epi32(ix[g], ix[g + 8]); + ix[g] = _mm256_permute4x64_epi64(ix[g], _MM_SHUFFLE(3, 1, 2, 0)); + ix[g] = _mm256_shuffle_epi8(ix[g], shuffle_mask); + ix[g] = _mm256_permute4x64_epi64(ix[g], _MM_SHUFFLE(3, 1, 2, 0)); + } + + int8_t* qlut_i8 = reinterpret_cast(qlut); + + _mm256_storeu_si256(reinterpret_cast<__m256i*>(qlut_i8 + k * 256 + 0 * 32 + 0), ix[0]); + _mm256_storeu_si256(reinterpret_cast<__m256i*>(qlut_i8 + k * 256 + 1 * 32 + 0), ix[1]); + _mm256_storeu_si256(reinterpret_cast<__m256i*>(qlut_i8 + k * 256 + 2 * 32 + 0), ix[2]); + _mm256_storeu_si256(reinterpret_cast<__m256i*>(qlut_i8 + k * 256 + 3 * 32 + 0), ix[3]); + _mm256_storeu_si256(reinterpret_cast<__m256i*>(qlut_i8 + k * 256 + 4 * 32 + 0), ix[4]); + _mm256_storeu_si256(reinterpret_cast<__m256i*>(qlut_i8 + k * 256 + 5 * 32 + 0), ix[5]); + _mm256_storeu_si256(reinterpret_cast<__m256i*>(qlut_i8 + k * 256 + 6 * 32 + 0), ix[6]); + _mm256_storeu_si256(reinterpret_cast<__m256i*>(qlut_i8 + k * 256 + 7 * 32 + 0), ix[7]); + + } + *lut_scales = scales; +#endif + return 0; +} +static bool is_type_supported(enum ggml_type type) { + if (type == GGML_TYPE_Q4_0 || + type == GGML_TYPE_TL2) { + return true; + } else { + return false; + } +} +#include + +#define BM3200_8640 160 +#define BBK3200_8640 96 +template +inline void three_tbl_impl_3200_8640(int32_t* c, int8_t* lut, uint8_t* a, uint8_t* sign) { +#ifdef __AVX2__ + const __m256i vec_mask = _mm256_set1_epi8(0x0f); + const __m256i vec_sign_mask = _mm256_set1_epi16(0x8000); + const __m256i vec_zero = _mm256_set1_epi8(0x00); + const __m256i vec_one = _mm256_set1_epi8(0xff); + const int KK = BBK3200_8640 / 3; +#pragma unroll + for (int i = 0; i < BM3200_8640; i += 32) { + __m256i vec_as[KK / 2]; + __m256i vec_signs[KK / 8]; + #pragma unroll + for (int ai = 0; ai < KK / 2; ai++) { + vec_as[ai] = _mm256_loadu_si256(reinterpret_cast<__m256i*>(a + i * KK / 2 + ai * 32)); + } + #pragma unroll + for (int as = 0; as < KK / 8; as++) { + vec_signs[as] = _mm256_loadu_si256(reinterpret_cast<__m256i*>(sign + i * KK / 8 + as * 32)); + } +#pragma unroll + for (int bs = 0; bs < batch_size; bs++) { + __m256i vec_c0 = _mm256_setzero_si256(); + __m256i vec_c1 = _mm256_setzero_si256(); +#pragma unroll + for (int k = 0; k < KK / 8; k++) { + __m256i vec_sign = vec_signs[k]; + __m256i vec_a_0 = vec_as[k * 4 + 0]; + __m128i vec_k1_0 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + 0 * 64 + 0 + K3 / 3 * 32 * bs)); + __m128i vec_k2_0 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + 0 * 64 + 16 + K3 / 3 * 32 * bs)); + __m128i vec_k3_0 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + 0 * 64 + 32 + K3 / 3 * 32 * bs)); + __m128i vec_k4_0 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + 0 * 64 + 48 + K3 / 3 * 32 * bs)); + __m256i vec_sign_left_hi_0 = _mm256_srai_epi16(_mm256_slli_epi16(vec_sign, (4 * 0)), 15); + __m256i vec_sign_left_lo_0 = _mm256_srai_epi16(_mm256_slli_epi16(vec_sign, (4 * 0 + 1)), 15); + __m256i vec_v_top_0 = _mm256_and_si256(_mm256_srli_epi16(vec_a_0, 4), vec_mask); + __m256i vec_v_top_fir_0 = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k1_0, vec_k1_0), vec_v_top_0); + __m256i vec_v_top_sec_0 = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k2_0, vec_k2_0), vec_v_top_0); + __m256i vec_sign_right_hi_0 = _mm256_srai_epi16(_mm256_slli_epi16(vec_sign, (4 * 0 + 2)), 15); + __m256i vec_sign_right_lo_0 = _mm256_srai_epi16(_mm256_slli_epi16(vec_sign, (4 * 0 + 3)), 15); + __m256i vec_v_bot_0 = _mm256_and_si256(vec_a_0, vec_mask); + __m256i vec_v_bot_fir_0 = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k3_0, vec_k3_0), vec_v_bot_0); + __m256i vec_v_bot_sec_0 = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k4_0, vec_k4_0), vec_v_bot_0); + __m256i vec_v_top_lo_0 = _mm256_xor_si256(_mm256_add_epi16(_mm256_unpackhi_epi8(vec_v_top_fir_0, vec_v_top_sec_0), vec_sign_left_lo_0), vec_sign_left_lo_0); + __m256i vec_v_top_hi_0 = _mm256_xor_si256(_mm256_add_epi16(_mm256_unpacklo_epi8(vec_v_top_fir_0, vec_v_top_sec_0), vec_sign_left_hi_0), vec_sign_left_hi_0); + __m256i vec_v_bot_lo_0 = _mm256_xor_si256(_mm256_add_epi16(_mm256_unpackhi_epi8(vec_v_bot_fir_0, vec_v_bot_sec_0), vec_sign_right_lo_0), vec_sign_right_lo_0); + __m256i vec_v_bot_hi_0 = _mm256_xor_si256(_mm256_add_epi16(_mm256_unpacklo_epi8(vec_v_bot_fir_0, vec_v_bot_sec_0), vec_sign_right_hi_0), vec_sign_right_hi_0); + vec_c0 = _mm256_add_epi16(vec_c0, vec_v_top_hi_0); + vec_c0 = _mm256_add_epi16(vec_c0, vec_v_bot_hi_0); + vec_c1 = _mm256_add_epi16(vec_c1, vec_v_top_lo_0); + vec_c1 = _mm256_add_epi16(vec_c1, vec_v_bot_lo_0); + __m256i vec_a_1 = vec_as[k * 4 + 1]; + __m128i vec_k1_1 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + 1 * 64 + 0 + K3 / 3 * 32 * bs)); + __m128i vec_k2_1 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + 1 * 64 + 16 + K3 / 3 * 32 * bs)); + __m128i vec_k3_1 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + 1 * 64 + 32 + K3 / 3 * 32 * bs)); + __m128i vec_k4_1 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + 1 * 64 + 48 + K3 / 3 * 32 * bs)); + __m256i vec_sign_left_hi_1 = _mm256_srai_epi16(_mm256_slli_epi16(vec_sign, (4 * 1)), 15); + __m256i vec_sign_left_lo_1 = _mm256_srai_epi16(_mm256_slli_epi16(vec_sign, (4 * 1 + 1)), 15); + __m256i vec_v_top_1 = _mm256_and_si256(_mm256_srli_epi16(vec_a_1, 4), vec_mask); + __m256i vec_v_top_fir_1 = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k1_1, vec_k1_1), vec_v_top_1); + __m256i vec_v_top_sec_1 = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k2_1, vec_k2_1), vec_v_top_1); + __m256i vec_sign_right_hi_1 = _mm256_srai_epi16(_mm256_slli_epi16(vec_sign, (4 * 1 + 2)), 15); + __m256i vec_sign_right_lo_1 = _mm256_srai_epi16(_mm256_slli_epi16(vec_sign, (4 * 1 + 3)), 15); + __m256i vec_v_bot_1 = _mm256_and_si256(vec_a_1, vec_mask); + __m256i vec_v_bot_fir_1 = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k3_1, vec_k3_1), vec_v_bot_1); + __m256i vec_v_bot_sec_1 = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k4_1, vec_k4_1), vec_v_bot_1); + __m256i vec_v_top_lo_1 = _mm256_xor_si256(_mm256_add_epi16(_mm256_unpackhi_epi8(vec_v_top_fir_1, vec_v_top_sec_1), vec_sign_left_lo_1), vec_sign_left_lo_1); + __m256i vec_v_top_hi_1 = _mm256_xor_si256(_mm256_add_epi16(_mm256_unpacklo_epi8(vec_v_top_fir_1, vec_v_top_sec_1), vec_sign_left_hi_1), vec_sign_left_hi_1); + __m256i vec_v_bot_lo_1 = _mm256_xor_si256(_mm256_add_epi16(_mm256_unpackhi_epi8(vec_v_bot_fir_1, vec_v_bot_sec_1), vec_sign_right_lo_1), vec_sign_right_lo_1); + __m256i vec_v_bot_hi_1 = _mm256_xor_si256(_mm256_add_epi16(_mm256_unpacklo_epi8(vec_v_bot_fir_1, vec_v_bot_sec_1), vec_sign_right_hi_1), vec_sign_right_hi_1); + vec_c0 = _mm256_add_epi16(vec_c0, vec_v_top_hi_1); + vec_c0 = _mm256_add_epi16(vec_c0, vec_v_bot_hi_1); + vec_c1 = _mm256_add_epi16(vec_c1, vec_v_top_lo_1); + vec_c1 = _mm256_add_epi16(vec_c1, vec_v_bot_lo_1); + __m256i vec_a_2 = vec_as[k * 4 + 2]; + __m128i vec_k1_2 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + 2 * 64 + 0 + K3 / 3 * 32 * bs)); + __m128i vec_k2_2 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + 2 * 64 + 16 + K3 / 3 * 32 * bs)); + __m128i vec_k3_2 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + 2 * 64 + 32 + K3 / 3 * 32 * bs)); + __m128i vec_k4_2 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + 2 * 64 + 48 + K3 / 3 * 32 * bs)); + __m256i vec_sign_left_hi_2 = _mm256_srai_epi16(_mm256_slli_epi16(vec_sign, (4 * 2)), 15); + __m256i vec_sign_left_lo_2 = _mm256_srai_epi16(_mm256_slli_epi16(vec_sign, (4 * 2 + 1)), 15); + __m256i vec_v_top_2 = _mm256_and_si256(_mm256_srli_epi16(vec_a_2, 4), vec_mask); + __m256i vec_v_top_fir_2 = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k1_2, vec_k1_2), vec_v_top_2); + __m256i vec_v_top_sec_2 = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k2_2, vec_k2_2), vec_v_top_2); + __m256i vec_sign_right_hi_2 = _mm256_srai_epi16(_mm256_slli_epi16(vec_sign, (4 * 2 + 2)), 15); + __m256i vec_sign_right_lo_2 = _mm256_srai_epi16(_mm256_slli_epi16(vec_sign, (4 * 2 + 3)), 15); + __m256i vec_v_bot_2 = _mm256_and_si256(vec_a_2, vec_mask); + __m256i vec_v_bot_fir_2 = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k3_2, vec_k3_2), vec_v_bot_2); + __m256i vec_v_bot_sec_2 = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k4_2, vec_k4_2), vec_v_bot_2); + __m256i vec_v_top_lo_2 = _mm256_xor_si256(_mm256_add_epi16(_mm256_unpackhi_epi8(vec_v_top_fir_2, vec_v_top_sec_2), vec_sign_left_lo_2), vec_sign_left_lo_2); + __m256i vec_v_top_hi_2 = _mm256_xor_si256(_mm256_add_epi16(_mm256_unpacklo_epi8(vec_v_top_fir_2, vec_v_top_sec_2), vec_sign_left_hi_2), vec_sign_left_hi_2); + __m256i vec_v_bot_lo_2 = _mm256_xor_si256(_mm256_add_epi16(_mm256_unpackhi_epi8(vec_v_bot_fir_2, vec_v_bot_sec_2), vec_sign_right_lo_2), vec_sign_right_lo_2); + __m256i vec_v_bot_hi_2 = _mm256_xor_si256(_mm256_add_epi16(_mm256_unpacklo_epi8(vec_v_bot_fir_2, vec_v_bot_sec_2), vec_sign_right_hi_2), vec_sign_right_hi_2); + vec_c0 = _mm256_add_epi16(vec_c0, vec_v_top_hi_2); + vec_c0 = _mm256_add_epi16(vec_c0, vec_v_bot_hi_2); + vec_c1 = _mm256_add_epi16(vec_c1, vec_v_top_lo_2); + vec_c1 = _mm256_add_epi16(vec_c1, vec_v_bot_lo_2); + __m256i vec_a_3 = vec_as[k * 4 + 3]; + __m128i vec_k1_3 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + 3 * 64 + 0 + K3 / 3 * 32 * bs)); + __m128i vec_k2_3 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + 3 * 64 + 16 + K3 / 3 * 32 * bs)); + __m128i vec_k3_3 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + 3 * 64 + 32 + K3 / 3 * 32 * bs)); + __m128i vec_k4_3 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + 3 * 64 + 48 + K3 / 3 * 32 * bs)); + __m256i vec_sign_left_hi_3 = _mm256_srai_epi16(_mm256_slli_epi16(vec_sign, (4 * 3)), 15); + __m256i vec_sign_left_lo_3 = _mm256_srai_epi16(_mm256_slli_epi16(vec_sign, (4 * 3 + 1)), 15); + __m256i vec_v_top_3 = _mm256_and_si256(_mm256_srli_epi16(vec_a_3, 4), vec_mask); + __m256i vec_v_top_fir_3 = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k1_3, vec_k1_3), vec_v_top_3); + __m256i vec_v_top_sec_3 = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k2_3, vec_k2_3), vec_v_top_3); + __m256i vec_sign_right_hi_3 = _mm256_srai_epi16(_mm256_slli_epi16(vec_sign, (4 * 3 + 2)), 15); + __m256i vec_sign_right_lo_3 = _mm256_srai_epi16(_mm256_slli_epi16(vec_sign, (4 * 3 + 3)), 15); + __m256i vec_v_bot_3 = _mm256_and_si256(vec_a_3, vec_mask); + __m256i vec_v_bot_fir_3 = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k3_3, vec_k3_3), vec_v_bot_3); + __m256i vec_v_bot_sec_3 = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k4_3, vec_k4_3), vec_v_bot_3); + __m256i vec_v_top_lo_3 = _mm256_xor_si256(_mm256_add_epi16(_mm256_unpackhi_epi8(vec_v_top_fir_3, vec_v_top_sec_3), vec_sign_left_lo_3), vec_sign_left_lo_3); + __m256i vec_v_top_hi_3 = _mm256_xor_si256(_mm256_add_epi16(_mm256_unpacklo_epi8(vec_v_top_fir_3, vec_v_top_sec_3), vec_sign_left_hi_3), vec_sign_left_hi_3); + __m256i vec_v_bot_lo_3 = _mm256_xor_si256(_mm256_add_epi16(_mm256_unpackhi_epi8(vec_v_bot_fir_3, vec_v_bot_sec_3), vec_sign_right_lo_3), vec_sign_right_lo_3); + __m256i vec_v_bot_hi_3 = _mm256_xor_si256(_mm256_add_epi16(_mm256_unpacklo_epi8(vec_v_bot_fir_3, vec_v_bot_sec_3), vec_sign_right_hi_3), vec_sign_right_hi_3); + vec_c0 = _mm256_add_epi16(vec_c0, vec_v_top_hi_3); + vec_c0 = _mm256_add_epi16(vec_c0, vec_v_bot_hi_3); + vec_c1 = _mm256_add_epi16(vec_c1, vec_v_top_lo_3); + vec_c1 = _mm256_add_epi16(vec_c1, vec_v_bot_lo_3); + } + __m256i vec_gc0 = _mm256_loadu_si256(reinterpret_cast<__m256i*>(c + i + BM3200_8640 * bs)); + __m256i vec_gc1 = _mm256_loadu_si256(reinterpret_cast<__m256i*>(c + i + 8 + BM3200_8640 * bs)); + __m256i vec_gc2 = _mm256_loadu_si256(reinterpret_cast<__m256i*>(c + i + 16 + BM3200_8640 * bs)); + __m256i vec_gc3 = _mm256_loadu_si256(reinterpret_cast<__m256i*>(c + i + 24 + BM3200_8640 * bs)); + vec_gc0 = _mm256_add_epi32(vec_gc0, _mm256_cvtepi16_epi32(_mm256_castsi256_si128(vec_c0))); + vec_gc1 = _mm256_add_epi32(vec_gc1, _mm256_cvtepi16_epi32(_mm256_extracti128_si256(vec_c0, 1))); + vec_gc2 = _mm256_add_epi32(vec_gc2, _mm256_cvtepi16_epi32(_mm256_castsi256_si128(vec_c1))); + vec_gc3 = _mm256_add_epi32(vec_gc3, _mm256_cvtepi16_epi32(_mm256_extracti128_si256(vec_c1, 1))); + _mm256_storeu_si256(reinterpret_cast<__m256i*>(c + i + BM3200_8640 * bs), vec_gc0); + _mm256_storeu_si256(reinterpret_cast<__m256i*>(c + i + 8 + BM3200_8640 * bs), vec_gc1); + _mm256_storeu_si256(reinterpret_cast<__m256i*>(c + i + 16 + BM3200_8640 * bs), vec_gc2); + _mm256_storeu_si256(reinterpret_cast<__m256i*>(c + i + 24 + BM3200_8640 * bs), vec_gc3); + } + } +#endif +} + +template +inline int32_t two_tbl_impl3200_8640(int32_t* c, int8_t* lut, uint8_t* a) { +#ifdef __AVX2__ + const __m256i vec_mask = _mm256_set1_epi8(0x0f); + const int KK = BK2 / 2; +#pragma unroll + for (int i = 0; i < BM3200_8640; i += 32) { + __m256i vec_as[KK / 2]; + #pragma unroll + for (int ai = 0; ai < KK / 2; ai++) { + vec_as[ai] = _mm256_loadu_si256(reinterpret_cast<__m256i*>(a + i * KK / 2 + ai * 32)); + } +#pragma unroll + for (int bs = 0; bs < batch_size; bs++) { + __m256i vec_c0 = _mm256_setzero_si256(); + __m256i vec_c1 = _mm256_setzero_si256(); +#pragma unroll + for (int k = 0; k < KK / 8; k++) { + #pragma unroll + for (int j = 0; j < 4; j++) { + __m256i vec_a = vec_as[k * 4 + j]; + + __m128i vec_k1 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + j * 64 + 0 + K2 / 2 * 32 * bs)); + __m128i vec_k2 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + j * 64 + 16 + K2 / 2 * 32 * bs)); + __m128i vec_k3 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + j * 64 + 32 + K2 / 2 * 32 * bs)); + __m128i vec_k4 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + j * 64 + 48 + K2 / 2 * 32 * bs)); + + __m256i vec_v_top = _mm256_and_si256(_mm256_srli_epi16(vec_a, 4), vec_mask); + __m256i vec_v_top_fir = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k1, vec_k1), vec_v_top); + __m256i vec_v_top_sec = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k2, vec_k2), vec_v_top); + + __m256i vec_v_bot = _mm256_and_si256(vec_a, vec_mask); + __m256i vec_v_bot_fir = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k3, vec_k3), vec_v_bot); + __m256i vec_v_bot_sec = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k4, vec_k4), vec_v_bot); + + __m256i vec_v_top_lo = _mm256_unpackhi_epi8(vec_v_top_fir, vec_v_top_sec); + __m256i vec_v_top_hi = _mm256_unpacklo_epi8(vec_v_top_fir, vec_v_top_sec); + __m256i vec_v_bot_lo = _mm256_unpackhi_epi8(vec_v_bot_fir, vec_v_bot_sec); + __m256i vec_v_bot_hi = _mm256_unpacklo_epi8(vec_v_bot_fir, vec_v_bot_sec); + vec_c0 = _mm256_add_epi16(vec_c0, vec_v_top_hi); + vec_c0 = _mm256_add_epi16(vec_c0, vec_v_bot_hi); + vec_c1 = _mm256_add_epi16(vec_c1, vec_v_top_lo); + vec_c1 = _mm256_add_epi16(vec_c1, vec_v_bot_lo); + } + } + + __m256i vec_gc0 = _mm256_loadu_si256(reinterpret_cast<__m256i*>(c + i + BM3200_8640 * bs)); + __m256i vec_gc1 = _mm256_loadu_si256(reinterpret_cast<__m256i*>(c + i + 8 + BM3200_8640 * bs)); + __m256i vec_gc2 = _mm256_loadu_si256(reinterpret_cast<__m256i*>(c + i + 16 + BM3200_8640 * bs)); + __m256i vec_gc3 = _mm256_loadu_si256(reinterpret_cast<__m256i*>(c + i + 24 + BM3200_8640 * bs)); + + vec_gc0 = _mm256_add_epi32(vec_gc0, _mm256_cvtepi16_epi32(_mm256_castsi256_si128(vec_c0))); + vec_gc1 = _mm256_add_epi32(vec_gc1, _mm256_cvtepi16_epi32(_mm256_extracti128_si256(vec_c0, 1))); + vec_gc2 = _mm256_add_epi32(vec_gc2, _mm256_cvtepi16_epi32(_mm256_castsi256_si128(vec_c1))); + vec_gc3 = _mm256_add_epi32(vec_gc3, _mm256_cvtepi16_epi32(_mm256_extracti128_si256(vec_c1, 1))); + + _mm256_storeu_si256(reinterpret_cast<__m256i*>(c + i + BM3200_8640 * bs), vec_gc0); + _mm256_storeu_si256(reinterpret_cast<__m256i*>(c + i + 8 + BM3200_8640 * bs), vec_gc1); + _mm256_storeu_si256(reinterpret_cast<__m256i*>(c + i + 16 + BM3200_8640 * bs), vec_gc2); + _mm256_storeu_si256(reinterpret_cast<__m256i*>(c + i + 24 + BM3200_8640 * bs), vec_gc3); + } + } +#endif + return 0; +} + +template +int32_t three_qgemm_lut_3200_8640(void* A, void* sign, void* LUT, void* Scales, void* LUT_Scales, void* C) { + alignas(32) uint32_t CBits[BATCH_SIZE * BM3200_8640]; + memset(&(CBits[0]), 0, BATCH_SIZE * BM3200_8640 * sizeof(int32_t)); +#pragma unroll + for (int32_t k_outer = 0; k_outer < 8640 / BBK3200_8640; ++k_outer) { + three_tbl_impl_3200_8640((&(((int32_t*)CBits)[0])), (&(((int8_t*)LUT)[(k_outer * BBK3200_8640 / 3 * 32)])), (&(((uint8_t*)A)[(k_outer * BBK3200_8640 / 3 / 2 * BM3200_8640)])), (&(((uint8_t*)sign)[(k_outer * BBK3200_8640 / 3 / 8 * BM3200_8640)]))); + } +#pragma unroll + for (int bs = 0; bs < BATCH_SIZE; bs++) { +#pragma unroll + for (int i = 0; i < BM3200_8640; i++) { + ((int32_t*)C)[i] = (int32_t)(((int32_t*)CBits)[i + bs * BM3200_8640]); + } + } + return 0; +} + +template +int32_t two_qgemm_lut_3200_8640(void* A, void* LUT, void* Scales, void* LUT_Scales, void* C) { + alignas(32) uint32_t CBits[BATCH_SIZE * BM3200_8640]; + memset(&(CBits[0]), 0, BATCH_SIZE * BM3200_8640 * sizeof(int32_t)); +#pragma unroll + for (int32_t k_outer = 0; k_outer < 0 / 32; ++k_outer) { + two_tbl_impl3200_8640((&(((int32_t*)CBits)[0])), (&(((int8_t*)LUT)[(k_outer * BK2 / 2 * 32)])), (&(((uint8_t*)A)[(k_outer * BK2 / 2 / 2 * BM3200_8640)]))); + } +#pragma unroll + for (int bs = 0; bs < BATCH_SIZE; bs++) { +#pragma unroll + for (int i = 0; i < BM3200_8640; i++) { + ((int32_t*)C)[i] += (int32_t)(((int32_t*)CBits)[i + bs * BM3200_8640]); + ((float*)C)[i] = (float)(((int32_t*)C)[i]) / ((float*)LUT_Scales)[bs] * ((float*)Scales)[0]; + } + } + return 0; +} + +#include + +#define BM3200_3200 320 +#define BBK3200_3200 96 +template +inline void three_tbl_impl_3200_3200(int32_t* c, int8_t* lut, uint8_t* a, uint8_t* sign) { +#ifdef __AVX2__ + const __m256i vec_mask = _mm256_set1_epi8(0x0f); + const __m256i vec_sign_mask = _mm256_set1_epi16(0x8000); + const __m256i vec_zero = _mm256_set1_epi8(0x00); + const __m256i vec_one = _mm256_set1_epi8(0xff); + const int KK = BBK3200_3200 / 3; +#pragma unroll + for (int i = 0; i < BM3200_3200; i += 32) { + __m256i vec_as[KK / 2]; + __m256i vec_signs[KK / 8]; + #pragma unroll + for (int ai = 0; ai < KK / 2; ai++) { + vec_as[ai] = _mm256_loadu_si256(reinterpret_cast<__m256i*>(a + i * KK / 2 + ai * 32)); + } + #pragma unroll + for (int as = 0; as < KK / 8; as++) { + vec_signs[as] = _mm256_loadu_si256(reinterpret_cast<__m256i*>(sign + i * KK / 8 + as * 32)); + } +#pragma unroll + for (int bs = 0; bs < batch_size; bs++) { + __m256i vec_c0 = _mm256_setzero_si256(); + __m256i vec_c1 = _mm256_setzero_si256(); +#pragma unroll + for (int k = 0; k < KK / 8; k++) { + __m256i vec_sign = vec_signs[k]; + __m256i vec_a_0 = vec_as[k * 4 + 0]; + __m128i vec_k1_0 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + 0 * 64 + 0 + K3 / 3 * 32 * bs)); + __m128i vec_k2_0 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + 0 * 64 + 16 + K3 / 3 * 32 * bs)); + __m128i vec_k3_0 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + 0 * 64 + 32 + K3 / 3 * 32 * bs)); + __m128i vec_k4_0 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + 0 * 64 + 48 + K3 / 3 * 32 * bs)); + __m256i vec_sign_left_hi_0 = _mm256_srai_epi16(_mm256_slli_epi16(vec_sign, (4 * 0)), 15); + __m256i vec_sign_left_lo_0 = _mm256_srai_epi16(_mm256_slli_epi16(vec_sign, (4 * 0 + 1)), 15); + __m256i vec_v_top_0 = _mm256_and_si256(_mm256_srli_epi16(vec_a_0, 4), vec_mask); + __m256i vec_v_top_fir_0 = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k1_0, vec_k1_0), vec_v_top_0); + __m256i vec_v_top_sec_0 = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k2_0, vec_k2_0), vec_v_top_0); + __m256i vec_sign_right_hi_0 = _mm256_srai_epi16(_mm256_slli_epi16(vec_sign, (4 * 0 + 2)), 15); + __m256i vec_sign_right_lo_0 = _mm256_srai_epi16(_mm256_slli_epi16(vec_sign, (4 * 0 + 3)), 15); + __m256i vec_v_bot_0 = _mm256_and_si256(vec_a_0, vec_mask); + __m256i vec_v_bot_fir_0 = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k3_0, vec_k3_0), vec_v_bot_0); + __m256i vec_v_bot_sec_0 = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k4_0, vec_k4_0), vec_v_bot_0); + __m256i vec_v_top_lo_0 = _mm256_xor_si256(_mm256_add_epi16(_mm256_unpackhi_epi8(vec_v_top_fir_0, vec_v_top_sec_0), vec_sign_left_lo_0), vec_sign_left_lo_0); + __m256i vec_v_top_hi_0 = _mm256_xor_si256(_mm256_add_epi16(_mm256_unpacklo_epi8(vec_v_top_fir_0, vec_v_top_sec_0), vec_sign_left_hi_0), vec_sign_left_hi_0); + __m256i vec_v_bot_lo_0 = _mm256_xor_si256(_mm256_add_epi16(_mm256_unpackhi_epi8(vec_v_bot_fir_0, vec_v_bot_sec_0), vec_sign_right_lo_0), vec_sign_right_lo_0); + __m256i vec_v_bot_hi_0 = _mm256_xor_si256(_mm256_add_epi16(_mm256_unpacklo_epi8(vec_v_bot_fir_0, vec_v_bot_sec_0), vec_sign_right_hi_0), vec_sign_right_hi_0); + vec_c0 = _mm256_add_epi16(vec_c0, vec_v_top_hi_0); + vec_c0 = _mm256_add_epi16(vec_c0, vec_v_bot_hi_0); + vec_c1 = _mm256_add_epi16(vec_c1, vec_v_top_lo_0); + vec_c1 = _mm256_add_epi16(vec_c1, vec_v_bot_lo_0); + __m256i vec_a_1 = vec_as[k * 4 + 1]; + __m128i vec_k1_1 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + 1 * 64 + 0 + K3 / 3 * 32 * bs)); + __m128i vec_k2_1 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + 1 * 64 + 16 + K3 / 3 * 32 * bs)); + __m128i vec_k3_1 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + 1 * 64 + 32 + K3 / 3 * 32 * bs)); + __m128i vec_k4_1 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + 1 * 64 + 48 + K3 / 3 * 32 * bs)); + __m256i vec_sign_left_hi_1 = _mm256_srai_epi16(_mm256_slli_epi16(vec_sign, (4 * 1)), 15); + __m256i vec_sign_left_lo_1 = _mm256_srai_epi16(_mm256_slli_epi16(vec_sign, (4 * 1 + 1)), 15); + __m256i vec_v_top_1 = _mm256_and_si256(_mm256_srli_epi16(vec_a_1, 4), vec_mask); + __m256i vec_v_top_fir_1 = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k1_1, vec_k1_1), vec_v_top_1); + __m256i vec_v_top_sec_1 = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k2_1, vec_k2_1), vec_v_top_1); + __m256i vec_sign_right_hi_1 = _mm256_srai_epi16(_mm256_slli_epi16(vec_sign, (4 * 1 + 2)), 15); + __m256i vec_sign_right_lo_1 = _mm256_srai_epi16(_mm256_slli_epi16(vec_sign, (4 * 1 + 3)), 15); + __m256i vec_v_bot_1 = _mm256_and_si256(vec_a_1, vec_mask); + __m256i vec_v_bot_fir_1 = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k3_1, vec_k3_1), vec_v_bot_1); + __m256i vec_v_bot_sec_1 = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k4_1, vec_k4_1), vec_v_bot_1); + __m256i vec_v_top_lo_1 = _mm256_xor_si256(_mm256_add_epi16(_mm256_unpackhi_epi8(vec_v_top_fir_1, vec_v_top_sec_1), vec_sign_left_lo_1), vec_sign_left_lo_1); + __m256i vec_v_top_hi_1 = _mm256_xor_si256(_mm256_add_epi16(_mm256_unpacklo_epi8(vec_v_top_fir_1, vec_v_top_sec_1), vec_sign_left_hi_1), vec_sign_left_hi_1); + __m256i vec_v_bot_lo_1 = _mm256_xor_si256(_mm256_add_epi16(_mm256_unpackhi_epi8(vec_v_bot_fir_1, vec_v_bot_sec_1), vec_sign_right_lo_1), vec_sign_right_lo_1); + __m256i vec_v_bot_hi_1 = _mm256_xor_si256(_mm256_add_epi16(_mm256_unpacklo_epi8(vec_v_bot_fir_1, vec_v_bot_sec_1), vec_sign_right_hi_1), vec_sign_right_hi_1); + vec_c0 = _mm256_add_epi16(vec_c0, vec_v_top_hi_1); + vec_c0 = _mm256_add_epi16(vec_c0, vec_v_bot_hi_1); + vec_c1 = _mm256_add_epi16(vec_c1, vec_v_top_lo_1); + vec_c1 = _mm256_add_epi16(vec_c1, vec_v_bot_lo_1); + __m256i vec_a_2 = vec_as[k * 4 + 2]; + __m128i vec_k1_2 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + 2 * 64 + 0 + K3 / 3 * 32 * bs)); + __m128i vec_k2_2 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + 2 * 64 + 16 + K3 / 3 * 32 * bs)); + __m128i vec_k3_2 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + 2 * 64 + 32 + K3 / 3 * 32 * bs)); + __m128i vec_k4_2 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + 2 * 64 + 48 + K3 / 3 * 32 * bs)); + __m256i vec_sign_left_hi_2 = _mm256_srai_epi16(_mm256_slli_epi16(vec_sign, (4 * 2)), 15); + __m256i vec_sign_left_lo_2 = _mm256_srai_epi16(_mm256_slli_epi16(vec_sign, (4 * 2 + 1)), 15); + __m256i vec_v_top_2 = _mm256_and_si256(_mm256_srli_epi16(vec_a_2, 4), vec_mask); + __m256i vec_v_top_fir_2 = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k1_2, vec_k1_2), vec_v_top_2); + __m256i vec_v_top_sec_2 = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k2_2, vec_k2_2), vec_v_top_2); + __m256i vec_sign_right_hi_2 = _mm256_srai_epi16(_mm256_slli_epi16(vec_sign, (4 * 2 + 2)), 15); + __m256i vec_sign_right_lo_2 = _mm256_srai_epi16(_mm256_slli_epi16(vec_sign, (4 * 2 + 3)), 15); + __m256i vec_v_bot_2 = _mm256_and_si256(vec_a_2, vec_mask); + __m256i vec_v_bot_fir_2 = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k3_2, vec_k3_2), vec_v_bot_2); + __m256i vec_v_bot_sec_2 = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k4_2, vec_k4_2), vec_v_bot_2); + __m256i vec_v_top_lo_2 = _mm256_xor_si256(_mm256_add_epi16(_mm256_unpackhi_epi8(vec_v_top_fir_2, vec_v_top_sec_2), vec_sign_left_lo_2), vec_sign_left_lo_2); + __m256i vec_v_top_hi_2 = _mm256_xor_si256(_mm256_add_epi16(_mm256_unpacklo_epi8(vec_v_top_fir_2, vec_v_top_sec_2), vec_sign_left_hi_2), vec_sign_left_hi_2); + __m256i vec_v_bot_lo_2 = _mm256_xor_si256(_mm256_add_epi16(_mm256_unpackhi_epi8(vec_v_bot_fir_2, vec_v_bot_sec_2), vec_sign_right_lo_2), vec_sign_right_lo_2); + __m256i vec_v_bot_hi_2 = _mm256_xor_si256(_mm256_add_epi16(_mm256_unpacklo_epi8(vec_v_bot_fir_2, vec_v_bot_sec_2), vec_sign_right_hi_2), vec_sign_right_hi_2); + vec_c0 = _mm256_add_epi16(vec_c0, vec_v_top_hi_2); + vec_c0 = _mm256_add_epi16(vec_c0, vec_v_bot_hi_2); + vec_c1 = _mm256_add_epi16(vec_c1, vec_v_top_lo_2); + vec_c1 = _mm256_add_epi16(vec_c1, vec_v_bot_lo_2); + __m256i vec_a_3 = vec_as[k * 4 + 3]; + __m128i vec_k1_3 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + 3 * 64 + 0 + K3 / 3 * 32 * bs)); + __m128i vec_k2_3 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + 3 * 64 + 16 + K3 / 3 * 32 * bs)); + __m128i vec_k3_3 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + 3 * 64 + 32 + K3 / 3 * 32 * bs)); + __m128i vec_k4_3 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + 3 * 64 + 48 + K3 / 3 * 32 * bs)); + __m256i vec_sign_left_hi_3 = _mm256_srai_epi16(_mm256_slli_epi16(vec_sign, (4 * 3)), 15); + __m256i vec_sign_left_lo_3 = _mm256_srai_epi16(_mm256_slli_epi16(vec_sign, (4 * 3 + 1)), 15); + __m256i vec_v_top_3 = _mm256_and_si256(_mm256_srli_epi16(vec_a_3, 4), vec_mask); + __m256i vec_v_top_fir_3 = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k1_3, vec_k1_3), vec_v_top_3); + __m256i vec_v_top_sec_3 = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k2_3, vec_k2_3), vec_v_top_3); + __m256i vec_sign_right_hi_3 = _mm256_srai_epi16(_mm256_slli_epi16(vec_sign, (4 * 3 + 2)), 15); + __m256i vec_sign_right_lo_3 = _mm256_srai_epi16(_mm256_slli_epi16(vec_sign, (4 * 3 + 3)), 15); + __m256i vec_v_bot_3 = _mm256_and_si256(vec_a_3, vec_mask); + __m256i vec_v_bot_fir_3 = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k3_3, vec_k3_3), vec_v_bot_3); + __m256i vec_v_bot_sec_3 = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k4_3, vec_k4_3), vec_v_bot_3); + __m256i vec_v_top_lo_3 = _mm256_xor_si256(_mm256_add_epi16(_mm256_unpackhi_epi8(vec_v_top_fir_3, vec_v_top_sec_3), vec_sign_left_lo_3), vec_sign_left_lo_3); + __m256i vec_v_top_hi_3 = _mm256_xor_si256(_mm256_add_epi16(_mm256_unpacklo_epi8(vec_v_top_fir_3, vec_v_top_sec_3), vec_sign_left_hi_3), vec_sign_left_hi_3); + __m256i vec_v_bot_lo_3 = _mm256_xor_si256(_mm256_add_epi16(_mm256_unpackhi_epi8(vec_v_bot_fir_3, vec_v_bot_sec_3), vec_sign_right_lo_3), vec_sign_right_lo_3); + __m256i vec_v_bot_hi_3 = _mm256_xor_si256(_mm256_add_epi16(_mm256_unpacklo_epi8(vec_v_bot_fir_3, vec_v_bot_sec_3), vec_sign_right_hi_3), vec_sign_right_hi_3); + vec_c0 = _mm256_add_epi16(vec_c0, vec_v_top_hi_3); + vec_c0 = _mm256_add_epi16(vec_c0, vec_v_bot_hi_3); + vec_c1 = _mm256_add_epi16(vec_c1, vec_v_top_lo_3); + vec_c1 = _mm256_add_epi16(vec_c1, vec_v_bot_lo_3); + } + __m256i vec_gc0 = _mm256_loadu_si256(reinterpret_cast<__m256i*>(c + i + BM3200_3200 * bs)); + __m256i vec_gc1 = _mm256_loadu_si256(reinterpret_cast<__m256i*>(c + i + 8 + BM3200_3200 * bs)); + __m256i vec_gc2 = _mm256_loadu_si256(reinterpret_cast<__m256i*>(c + i + 16 + BM3200_3200 * bs)); + __m256i vec_gc3 = _mm256_loadu_si256(reinterpret_cast<__m256i*>(c + i + 24 + BM3200_3200 * bs)); + vec_gc0 = _mm256_add_epi32(vec_gc0, _mm256_cvtepi16_epi32(_mm256_castsi256_si128(vec_c0))); + vec_gc1 = _mm256_add_epi32(vec_gc1, _mm256_cvtepi16_epi32(_mm256_extracti128_si256(vec_c0, 1))); + vec_gc2 = _mm256_add_epi32(vec_gc2, _mm256_cvtepi16_epi32(_mm256_castsi256_si128(vec_c1))); + vec_gc3 = _mm256_add_epi32(vec_gc3, _mm256_cvtepi16_epi32(_mm256_extracti128_si256(vec_c1, 1))); + _mm256_storeu_si256(reinterpret_cast<__m256i*>(c + i + BM3200_3200 * bs), vec_gc0); + _mm256_storeu_si256(reinterpret_cast<__m256i*>(c + i + 8 + BM3200_3200 * bs), vec_gc1); + _mm256_storeu_si256(reinterpret_cast<__m256i*>(c + i + 16 + BM3200_3200 * bs), vec_gc2); + _mm256_storeu_si256(reinterpret_cast<__m256i*>(c + i + 24 + BM3200_3200 * bs), vec_gc3); + } + } +#endif +} + +template +inline int32_t two_tbl_impl3200_3200(int32_t* c, int8_t* lut, uint8_t* a) { +#ifdef __AVX2__ + const __m256i vec_mask = _mm256_set1_epi8(0x0f); + const int KK = BK2 / 2; +#pragma unroll + for (int i = 0; i < BM3200_3200; i += 32) { + __m256i vec_as[KK / 2]; + #pragma unroll + for (int ai = 0; ai < KK / 2; ai++) { + vec_as[ai] = _mm256_loadu_si256(reinterpret_cast<__m256i*>(a + i * KK / 2 + ai * 32)); + } +#pragma unroll + for (int bs = 0; bs < batch_size; bs++) { + __m256i vec_c0 = _mm256_setzero_si256(); + __m256i vec_c1 = _mm256_setzero_si256(); +#pragma unroll + for (int k = 0; k < KK / 8; k++) { + #pragma unroll + for (int j = 0; j < 4; j++) { + __m256i vec_a = vec_as[k * 4 + j]; + + __m128i vec_k1 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + j * 64 + 0 + K2 / 2 * 32 * bs)); + __m128i vec_k2 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + j * 64 + 16 + K2 / 2 * 32 * bs)); + __m128i vec_k3 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + j * 64 + 32 + K2 / 2 * 32 * bs)); + __m128i vec_k4 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + j * 64 + 48 + K2 / 2 * 32 * bs)); + + __m256i vec_v_top = _mm256_and_si256(_mm256_srli_epi16(vec_a, 4), vec_mask); + __m256i vec_v_top_fir = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k1, vec_k1), vec_v_top); + __m256i vec_v_top_sec = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k2, vec_k2), vec_v_top); + + __m256i vec_v_bot = _mm256_and_si256(vec_a, vec_mask); + __m256i vec_v_bot_fir = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k3, vec_k3), vec_v_bot); + __m256i vec_v_bot_sec = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k4, vec_k4), vec_v_bot); + + __m256i vec_v_top_lo = _mm256_unpackhi_epi8(vec_v_top_fir, vec_v_top_sec); + __m256i vec_v_top_hi = _mm256_unpacklo_epi8(vec_v_top_fir, vec_v_top_sec); + __m256i vec_v_bot_lo = _mm256_unpackhi_epi8(vec_v_bot_fir, vec_v_bot_sec); + __m256i vec_v_bot_hi = _mm256_unpacklo_epi8(vec_v_bot_fir, vec_v_bot_sec); + vec_c0 = _mm256_add_epi16(vec_c0, vec_v_top_hi); + vec_c0 = _mm256_add_epi16(vec_c0, vec_v_bot_hi); + vec_c1 = _mm256_add_epi16(vec_c1, vec_v_top_lo); + vec_c1 = _mm256_add_epi16(vec_c1, vec_v_bot_lo); + } + } + + __m256i vec_gc0 = _mm256_loadu_si256(reinterpret_cast<__m256i*>(c + i + BM3200_3200 * bs)); + __m256i vec_gc1 = _mm256_loadu_si256(reinterpret_cast<__m256i*>(c + i + 8 + BM3200_3200 * bs)); + __m256i vec_gc2 = _mm256_loadu_si256(reinterpret_cast<__m256i*>(c + i + 16 + BM3200_3200 * bs)); + __m256i vec_gc3 = _mm256_loadu_si256(reinterpret_cast<__m256i*>(c + i + 24 + BM3200_3200 * bs)); + + vec_gc0 = _mm256_add_epi32(vec_gc0, _mm256_cvtepi16_epi32(_mm256_castsi256_si128(vec_c0))); + vec_gc1 = _mm256_add_epi32(vec_gc1, _mm256_cvtepi16_epi32(_mm256_extracti128_si256(vec_c0, 1))); + vec_gc2 = _mm256_add_epi32(vec_gc2, _mm256_cvtepi16_epi32(_mm256_castsi256_si128(vec_c1))); + vec_gc3 = _mm256_add_epi32(vec_gc3, _mm256_cvtepi16_epi32(_mm256_extracti128_si256(vec_c1, 1))); + + _mm256_storeu_si256(reinterpret_cast<__m256i*>(c + i + BM3200_3200 * bs), vec_gc0); + _mm256_storeu_si256(reinterpret_cast<__m256i*>(c + i + 8 + BM3200_3200 * bs), vec_gc1); + _mm256_storeu_si256(reinterpret_cast<__m256i*>(c + i + 16 + BM3200_3200 * bs), vec_gc2); + _mm256_storeu_si256(reinterpret_cast<__m256i*>(c + i + 24 + BM3200_3200 * bs), vec_gc3); + } + } +#endif + return 0; +} + +template +int32_t three_qgemm_lut_3200_3200(void* A, void* sign, void* LUT, void* Scales, void* LUT_Scales, void* C) { + alignas(32) uint32_t CBits[BATCH_SIZE * BM3200_3200]; + memset(&(CBits[0]), 0, BATCH_SIZE * BM3200_3200 * sizeof(int32_t)); +#pragma unroll + for (int32_t k_outer = 0; k_outer < 3168 / BBK3200_3200; ++k_outer) { + three_tbl_impl_3200_3200((&(((int32_t*)CBits)[0])), (&(((int8_t*)LUT)[(k_outer * BBK3200_3200 / 3 * 32)])), (&(((uint8_t*)A)[(k_outer * BBK3200_3200 / 3 / 2 * BM3200_3200)])), (&(((uint8_t*)sign)[(k_outer * BBK3200_3200 / 3 / 8 * BM3200_3200)]))); + } +#pragma unroll + for (int bs = 0; bs < BATCH_SIZE; bs++) { +#pragma unroll + for (int i = 0; i < BM3200_3200; i++) { + ((int32_t*)C)[i] = (int32_t)(((int32_t*)CBits)[i + bs * BM3200_3200]); + } + } + return 0; +} + +template +int32_t two_qgemm_lut_3200_3200(void* A, void* LUT, void* Scales, void* LUT_Scales, void* C) { + alignas(32) uint32_t CBits[BATCH_SIZE * BM3200_3200]; + memset(&(CBits[0]), 0, BATCH_SIZE * BM3200_3200 * sizeof(int32_t)); +#pragma unroll + for (int32_t k_outer = 0; k_outer < 32 / 32; ++k_outer) { + two_tbl_impl3200_3200((&(((int32_t*)CBits)[0])), (&(((int8_t*)LUT)[(k_outer * BK2 / 2 * 32)])), (&(((uint8_t*)A)[(k_outer * BK2 / 2 / 2 * BM3200_3200)]))); + } +#pragma unroll + for (int bs = 0; bs < BATCH_SIZE; bs++) { +#pragma unroll + for (int i = 0; i < BM3200_3200; i++) { + ((int32_t*)C)[i] += (int32_t)(((int32_t*)CBits)[i + bs * BM3200_3200]); + ((float*)C)[i] = (float)(((int32_t*)C)[i]) / ((float*)LUT_Scales)[bs] * ((float*)Scales)[0]; + } + } + return 0; +} + +#include + +#define BM8640_3200 320 +#define BBK8640_3200 96 +template +inline void three_tbl_impl_8640_3200(int32_t* c, int8_t* lut, uint8_t* a, uint8_t* sign) { +#ifdef __AVX2__ + const __m256i vec_mask = _mm256_set1_epi8(0x0f); + const __m256i vec_sign_mask = _mm256_set1_epi16(0x8000); + const __m256i vec_zero = _mm256_set1_epi8(0x00); + const __m256i vec_one = _mm256_set1_epi8(0xff); + const int KK = BBK8640_3200 / 3; +#pragma unroll + for (int i = 0; i < BM8640_3200; i += 32) { + __m256i vec_as[KK / 2]; + __m256i vec_signs[KK / 8]; + #pragma unroll + for (int ai = 0; ai < KK / 2; ai++) { + vec_as[ai] = _mm256_loadu_si256(reinterpret_cast<__m256i*>(a + i * KK / 2 + ai * 32)); + } + #pragma unroll + for (int as = 0; as < KK / 8; as++) { + vec_signs[as] = _mm256_loadu_si256(reinterpret_cast<__m256i*>(sign + i * KK / 8 + as * 32)); + } +#pragma unroll + for (int bs = 0; bs < batch_size; bs++) { + __m256i vec_c0 = _mm256_setzero_si256(); + __m256i vec_c1 = _mm256_setzero_si256(); +#pragma unroll + for (int k = 0; k < KK / 8; k++) { + __m256i vec_sign = vec_signs[k]; + __m256i vec_a_0 = vec_as[k * 4 + 0]; + __m128i vec_k1_0 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + 0 * 64 + 0 + K3 / 3 * 32 * bs)); + __m128i vec_k2_0 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + 0 * 64 + 16 + K3 / 3 * 32 * bs)); + __m128i vec_k3_0 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + 0 * 64 + 32 + K3 / 3 * 32 * bs)); + __m128i vec_k4_0 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + 0 * 64 + 48 + K3 / 3 * 32 * bs)); + __m256i vec_sign_left_hi_0 = _mm256_srai_epi16(_mm256_slli_epi16(vec_sign, (4 * 0)), 15); + __m256i vec_sign_left_lo_0 = _mm256_srai_epi16(_mm256_slli_epi16(vec_sign, (4 * 0 + 1)), 15); + __m256i vec_v_top_0 = _mm256_and_si256(_mm256_srli_epi16(vec_a_0, 4), vec_mask); + __m256i vec_v_top_fir_0 = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k1_0, vec_k1_0), vec_v_top_0); + __m256i vec_v_top_sec_0 = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k2_0, vec_k2_0), vec_v_top_0); + __m256i vec_sign_right_hi_0 = _mm256_srai_epi16(_mm256_slli_epi16(vec_sign, (4 * 0 + 2)), 15); + __m256i vec_sign_right_lo_0 = _mm256_srai_epi16(_mm256_slli_epi16(vec_sign, (4 * 0 + 3)), 15); + __m256i vec_v_bot_0 = _mm256_and_si256(vec_a_0, vec_mask); + __m256i vec_v_bot_fir_0 = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k3_0, vec_k3_0), vec_v_bot_0); + __m256i vec_v_bot_sec_0 = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k4_0, vec_k4_0), vec_v_bot_0); + __m256i vec_v_top_lo_0 = _mm256_xor_si256(_mm256_add_epi16(_mm256_unpackhi_epi8(vec_v_top_fir_0, vec_v_top_sec_0), vec_sign_left_lo_0), vec_sign_left_lo_0); + __m256i vec_v_top_hi_0 = _mm256_xor_si256(_mm256_add_epi16(_mm256_unpacklo_epi8(vec_v_top_fir_0, vec_v_top_sec_0), vec_sign_left_hi_0), vec_sign_left_hi_0); + __m256i vec_v_bot_lo_0 = _mm256_xor_si256(_mm256_add_epi16(_mm256_unpackhi_epi8(vec_v_bot_fir_0, vec_v_bot_sec_0), vec_sign_right_lo_0), vec_sign_right_lo_0); + __m256i vec_v_bot_hi_0 = _mm256_xor_si256(_mm256_add_epi16(_mm256_unpacklo_epi8(vec_v_bot_fir_0, vec_v_bot_sec_0), vec_sign_right_hi_0), vec_sign_right_hi_0); + vec_c0 = _mm256_add_epi16(vec_c0, vec_v_top_hi_0); + vec_c0 = _mm256_add_epi16(vec_c0, vec_v_bot_hi_0); + vec_c1 = _mm256_add_epi16(vec_c1, vec_v_top_lo_0); + vec_c1 = _mm256_add_epi16(vec_c1, vec_v_bot_lo_0); + __m256i vec_a_1 = vec_as[k * 4 + 1]; + __m128i vec_k1_1 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + 1 * 64 + 0 + K3 / 3 * 32 * bs)); + __m128i vec_k2_1 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + 1 * 64 + 16 + K3 / 3 * 32 * bs)); + __m128i vec_k3_1 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + 1 * 64 + 32 + K3 / 3 * 32 * bs)); + __m128i vec_k4_1 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + 1 * 64 + 48 + K3 / 3 * 32 * bs)); + __m256i vec_sign_left_hi_1 = _mm256_srai_epi16(_mm256_slli_epi16(vec_sign, (4 * 1)), 15); + __m256i vec_sign_left_lo_1 = _mm256_srai_epi16(_mm256_slli_epi16(vec_sign, (4 * 1 + 1)), 15); + __m256i vec_v_top_1 = _mm256_and_si256(_mm256_srli_epi16(vec_a_1, 4), vec_mask); + __m256i vec_v_top_fir_1 = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k1_1, vec_k1_1), vec_v_top_1); + __m256i vec_v_top_sec_1 = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k2_1, vec_k2_1), vec_v_top_1); + __m256i vec_sign_right_hi_1 = _mm256_srai_epi16(_mm256_slli_epi16(vec_sign, (4 * 1 + 2)), 15); + __m256i vec_sign_right_lo_1 = _mm256_srai_epi16(_mm256_slli_epi16(vec_sign, (4 * 1 + 3)), 15); + __m256i vec_v_bot_1 = _mm256_and_si256(vec_a_1, vec_mask); + __m256i vec_v_bot_fir_1 = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k3_1, vec_k3_1), vec_v_bot_1); + __m256i vec_v_bot_sec_1 = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k4_1, vec_k4_1), vec_v_bot_1); + __m256i vec_v_top_lo_1 = _mm256_xor_si256(_mm256_add_epi16(_mm256_unpackhi_epi8(vec_v_top_fir_1, vec_v_top_sec_1), vec_sign_left_lo_1), vec_sign_left_lo_1); + __m256i vec_v_top_hi_1 = _mm256_xor_si256(_mm256_add_epi16(_mm256_unpacklo_epi8(vec_v_top_fir_1, vec_v_top_sec_1), vec_sign_left_hi_1), vec_sign_left_hi_1); + __m256i vec_v_bot_lo_1 = _mm256_xor_si256(_mm256_add_epi16(_mm256_unpackhi_epi8(vec_v_bot_fir_1, vec_v_bot_sec_1), vec_sign_right_lo_1), vec_sign_right_lo_1); + __m256i vec_v_bot_hi_1 = _mm256_xor_si256(_mm256_add_epi16(_mm256_unpacklo_epi8(vec_v_bot_fir_1, vec_v_bot_sec_1), vec_sign_right_hi_1), vec_sign_right_hi_1); + vec_c0 = _mm256_add_epi16(vec_c0, vec_v_top_hi_1); + vec_c0 = _mm256_add_epi16(vec_c0, vec_v_bot_hi_1); + vec_c1 = _mm256_add_epi16(vec_c1, vec_v_top_lo_1); + vec_c1 = _mm256_add_epi16(vec_c1, vec_v_bot_lo_1); + __m256i vec_a_2 = vec_as[k * 4 + 2]; + __m128i vec_k1_2 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + 2 * 64 + 0 + K3 / 3 * 32 * bs)); + __m128i vec_k2_2 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + 2 * 64 + 16 + K3 / 3 * 32 * bs)); + __m128i vec_k3_2 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + 2 * 64 + 32 + K3 / 3 * 32 * bs)); + __m128i vec_k4_2 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + 2 * 64 + 48 + K3 / 3 * 32 * bs)); + __m256i vec_sign_left_hi_2 = _mm256_srai_epi16(_mm256_slli_epi16(vec_sign, (4 * 2)), 15); + __m256i vec_sign_left_lo_2 = _mm256_srai_epi16(_mm256_slli_epi16(vec_sign, (4 * 2 + 1)), 15); + __m256i vec_v_top_2 = _mm256_and_si256(_mm256_srli_epi16(vec_a_2, 4), vec_mask); + __m256i vec_v_top_fir_2 = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k1_2, vec_k1_2), vec_v_top_2); + __m256i vec_v_top_sec_2 = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k2_2, vec_k2_2), vec_v_top_2); + __m256i vec_sign_right_hi_2 = _mm256_srai_epi16(_mm256_slli_epi16(vec_sign, (4 * 2 + 2)), 15); + __m256i vec_sign_right_lo_2 = _mm256_srai_epi16(_mm256_slli_epi16(vec_sign, (4 * 2 + 3)), 15); + __m256i vec_v_bot_2 = _mm256_and_si256(vec_a_2, vec_mask); + __m256i vec_v_bot_fir_2 = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k3_2, vec_k3_2), vec_v_bot_2); + __m256i vec_v_bot_sec_2 = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k4_2, vec_k4_2), vec_v_bot_2); + __m256i vec_v_top_lo_2 = _mm256_xor_si256(_mm256_add_epi16(_mm256_unpackhi_epi8(vec_v_top_fir_2, vec_v_top_sec_2), vec_sign_left_lo_2), vec_sign_left_lo_2); + __m256i vec_v_top_hi_2 = _mm256_xor_si256(_mm256_add_epi16(_mm256_unpacklo_epi8(vec_v_top_fir_2, vec_v_top_sec_2), vec_sign_left_hi_2), vec_sign_left_hi_2); + __m256i vec_v_bot_lo_2 = _mm256_xor_si256(_mm256_add_epi16(_mm256_unpackhi_epi8(vec_v_bot_fir_2, vec_v_bot_sec_2), vec_sign_right_lo_2), vec_sign_right_lo_2); + __m256i vec_v_bot_hi_2 = _mm256_xor_si256(_mm256_add_epi16(_mm256_unpacklo_epi8(vec_v_bot_fir_2, vec_v_bot_sec_2), vec_sign_right_hi_2), vec_sign_right_hi_2); + vec_c0 = _mm256_add_epi16(vec_c0, vec_v_top_hi_2); + vec_c0 = _mm256_add_epi16(vec_c0, vec_v_bot_hi_2); + vec_c1 = _mm256_add_epi16(vec_c1, vec_v_top_lo_2); + vec_c1 = _mm256_add_epi16(vec_c1, vec_v_bot_lo_2); + __m256i vec_a_3 = vec_as[k * 4 + 3]; + __m128i vec_k1_3 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + 3 * 64 + 0 + K3 / 3 * 32 * bs)); + __m128i vec_k2_3 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + 3 * 64 + 16 + K3 / 3 * 32 * bs)); + __m128i vec_k3_3 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + 3 * 64 + 32 + K3 / 3 * 32 * bs)); + __m128i vec_k4_3 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + 3 * 64 + 48 + K3 / 3 * 32 * bs)); + __m256i vec_sign_left_hi_3 = _mm256_srai_epi16(_mm256_slli_epi16(vec_sign, (4 * 3)), 15); + __m256i vec_sign_left_lo_3 = _mm256_srai_epi16(_mm256_slli_epi16(vec_sign, (4 * 3 + 1)), 15); + __m256i vec_v_top_3 = _mm256_and_si256(_mm256_srli_epi16(vec_a_3, 4), vec_mask); + __m256i vec_v_top_fir_3 = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k1_3, vec_k1_3), vec_v_top_3); + __m256i vec_v_top_sec_3 = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k2_3, vec_k2_3), vec_v_top_3); + __m256i vec_sign_right_hi_3 = _mm256_srai_epi16(_mm256_slli_epi16(vec_sign, (4 * 3 + 2)), 15); + __m256i vec_sign_right_lo_3 = _mm256_srai_epi16(_mm256_slli_epi16(vec_sign, (4 * 3 + 3)), 15); + __m256i vec_v_bot_3 = _mm256_and_si256(vec_a_3, vec_mask); + __m256i vec_v_bot_fir_3 = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k3_3, vec_k3_3), vec_v_bot_3); + __m256i vec_v_bot_sec_3 = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k4_3, vec_k4_3), vec_v_bot_3); + __m256i vec_v_top_lo_3 = _mm256_xor_si256(_mm256_add_epi16(_mm256_unpackhi_epi8(vec_v_top_fir_3, vec_v_top_sec_3), vec_sign_left_lo_3), vec_sign_left_lo_3); + __m256i vec_v_top_hi_3 = _mm256_xor_si256(_mm256_add_epi16(_mm256_unpacklo_epi8(vec_v_top_fir_3, vec_v_top_sec_3), vec_sign_left_hi_3), vec_sign_left_hi_3); + __m256i vec_v_bot_lo_3 = _mm256_xor_si256(_mm256_add_epi16(_mm256_unpackhi_epi8(vec_v_bot_fir_3, vec_v_bot_sec_3), vec_sign_right_lo_3), vec_sign_right_lo_3); + __m256i vec_v_bot_hi_3 = _mm256_xor_si256(_mm256_add_epi16(_mm256_unpacklo_epi8(vec_v_bot_fir_3, vec_v_bot_sec_3), vec_sign_right_hi_3), vec_sign_right_hi_3); + vec_c0 = _mm256_add_epi16(vec_c0, vec_v_top_hi_3); + vec_c0 = _mm256_add_epi16(vec_c0, vec_v_bot_hi_3); + vec_c1 = _mm256_add_epi16(vec_c1, vec_v_top_lo_3); + vec_c1 = _mm256_add_epi16(vec_c1, vec_v_bot_lo_3); + } + __m256i vec_gc0 = _mm256_loadu_si256(reinterpret_cast<__m256i*>(c + i + BM8640_3200 * bs)); + __m256i vec_gc1 = _mm256_loadu_si256(reinterpret_cast<__m256i*>(c + i + 8 + BM8640_3200 * bs)); + __m256i vec_gc2 = _mm256_loadu_si256(reinterpret_cast<__m256i*>(c + i + 16 + BM8640_3200 * bs)); + __m256i vec_gc3 = _mm256_loadu_si256(reinterpret_cast<__m256i*>(c + i + 24 + BM8640_3200 * bs)); + vec_gc0 = _mm256_add_epi32(vec_gc0, _mm256_cvtepi16_epi32(_mm256_castsi256_si128(vec_c0))); + vec_gc1 = _mm256_add_epi32(vec_gc1, _mm256_cvtepi16_epi32(_mm256_extracti128_si256(vec_c0, 1))); + vec_gc2 = _mm256_add_epi32(vec_gc2, _mm256_cvtepi16_epi32(_mm256_castsi256_si128(vec_c1))); + vec_gc3 = _mm256_add_epi32(vec_gc3, _mm256_cvtepi16_epi32(_mm256_extracti128_si256(vec_c1, 1))); + _mm256_storeu_si256(reinterpret_cast<__m256i*>(c + i + BM8640_3200 * bs), vec_gc0); + _mm256_storeu_si256(reinterpret_cast<__m256i*>(c + i + 8 + BM8640_3200 * bs), vec_gc1); + _mm256_storeu_si256(reinterpret_cast<__m256i*>(c + i + 16 + BM8640_3200 * bs), vec_gc2); + _mm256_storeu_si256(reinterpret_cast<__m256i*>(c + i + 24 + BM8640_3200 * bs), vec_gc3); + } + } +#endif +} + +template +inline int32_t two_tbl_impl8640_3200(int32_t* c, int8_t* lut, uint8_t* a) { +#ifdef __AVX2__ + const __m256i vec_mask = _mm256_set1_epi8(0x0f); + const int KK = BK2 / 2; +#pragma unroll + for (int i = 0; i < BM8640_3200; i += 32) { + __m256i vec_as[KK / 2]; + #pragma unroll + for (int ai = 0; ai < KK / 2; ai++) { + vec_as[ai] = _mm256_loadu_si256(reinterpret_cast<__m256i*>(a + i * KK / 2 + ai * 32)); + } +#pragma unroll + for (int bs = 0; bs < batch_size; bs++) { + __m256i vec_c0 = _mm256_setzero_si256(); + __m256i vec_c1 = _mm256_setzero_si256(); +#pragma unroll + for (int k = 0; k < KK / 8; k++) { + #pragma unroll + for (int j = 0; j < 4; j++) { + __m256i vec_a = vec_as[k * 4 + j]; + + __m128i vec_k1 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + j * 64 + 0 + K2 / 2 * 32 * bs)); + __m128i vec_k2 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + j * 64 + 16 + K2 / 2 * 32 * bs)); + __m128i vec_k3 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + j * 64 + 32 + K2 / 2 * 32 * bs)); + __m128i vec_k4 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + j * 64 + 48 + K2 / 2 * 32 * bs)); + + __m256i vec_v_top = _mm256_and_si256(_mm256_srli_epi16(vec_a, 4), vec_mask); + __m256i vec_v_top_fir = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k1, vec_k1), vec_v_top); + __m256i vec_v_top_sec = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k2, vec_k2), vec_v_top); + + __m256i vec_v_bot = _mm256_and_si256(vec_a, vec_mask); + __m256i vec_v_bot_fir = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k3, vec_k3), vec_v_bot); + __m256i vec_v_bot_sec = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k4, vec_k4), vec_v_bot); + + __m256i vec_v_top_lo = _mm256_unpackhi_epi8(vec_v_top_fir, vec_v_top_sec); + __m256i vec_v_top_hi = _mm256_unpacklo_epi8(vec_v_top_fir, vec_v_top_sec); + __m256i vec_v_bot_lo = _mm256_unpackhi_epi8(vec_v_bot_fir, vec_v_bot_sec); + __m256i vec_v_bot_hi = _mm256_unpacklo_epi8(vec_v_bot_fir, vec_v_bot_sec); + vec_c0 = _mm256_add_epi16(vec_c0, vec_v_top_hi); + vec_c0 = _mm256_add_epi16(vec_c0, vec_v_bot_hi); + vec_c1 = _mm256_add_epi16(vec_c1, vec_v_top_lo); + vec_c1 = _mm256_add_epi16(vec_c1, vec_v_bot_lo); + } + } + + __m256i vec_gc0 = _mm256_loadu_si256(reinterpret_cast<__m256i*>(c + i + BM8640_3200 * bs)); + __m256i vec_gc1 = _mm256_loadu_si256(reinterpret_cast<__m256i*>(c + i + 8 + BM8640_3200 * bs)); + __m256i vec_gc2 = _mm256_loadu_si256(reinterpret_cast<__m256i*>(c + i + 16 + BM8640_3200 * bs)); + __m256i vec_gc3 = _mm256_loadu_si256(reinterpret_cast<__m256i*>(c + i + 24 + BM8640_3200 * bs)); + + vec_gc0 = _mm256_add_epi32(vec_gc0, _mm256_cvtepi16_epi32(_mm256_castsi256_si128(vec_c0))); + vec_gc1 = _mm256_add_epi32(vec_gc1, _mm256_cvtepi16_epi32(_mm256_extracti128_si256(vec_c0, 1))); + vec_gc2 = _mm256_add_epi32(vec_gc2, _mm256_cvtepi16_epi32(_mm256_castsi256_si128(vec_c1))); + vec_gc3 = _mm256_add_epi32(vec_gc3, _mm256_cvtepi16_epi32(_mm256_extracti128_si256(vec_c1, 1))); + + _mm256_storeu_si256(reinterpret_cast<__m256i*>(c + i + BM8640_3200 * bs), vec_gc0); + _mm256_storeu_si256(reinterpret_cast<__m256i*>(c + i + 8 + BM8640_3200 * bs), vec_gc1); + _mm256_storeu_si256(reinterpret_cast<__m256i*>(c + i + 16 + BM8640_3200 * bs), vec_gc2); + _mm256_storeu_si256(reinterpret_cast<__m256i*>(c + i + 24 + BM8640_3200 * bs), vec_gc3); + } + } +#endif + return 0; +} + +template +int32_t three_qgemm_lut_8640_3200(void* A, void* sign, void* LUT, void* Scales, void* LUT_Scales, void* C) { + alignas(32) uint32_t CBits[BATCH_SIZE * BM8640_3200]; + memset(&(CBits[0]), 0, BATCH_SIZE * BM8640_3200 * sizeof(int32_t)); +#pragma unroll + for (int32_t k_outer = 0; k_outer < 3168 / BBK8640_3200; ++k_outer) { + three_tbl_impl_8640_3200((&(((int32_t*)CBits)[0])), (&(((int8_t*)LUT)[(k_outer * BBK8640_3200 / 3 * 32)])), (&(((uint8_t*)A)[(k_outer * BBK8640_3200 / 3 / 2 * BM8640_3200)])), (&(((uint8_t*)sign)[(k_outer * BBK8640_3200 / 3 / 8 * BM8640_3200)]))); + } +#pragma unroll + for (int bs = 0; bs < BATCH_SIZE; bs++) { +#pragma unroll + for (int i = 0; i < BM8640_3200; i++) { + ((int32_t*)C)[i] = (int32_t)(((int32_t*)CBits)[i + bs * BM8640_3200]); + } + } + return 0; +} + +template +int32_t two_qgemm_lut_8640_3200(void* A, void* LUT, void* Scales, void* LUT_Scales, void* C) { + alignas(32) uint32_t CBits[BATCH_SIZE * BM8640_3200]; + memset(&(CBits[0]), 0, BATCH_SIZE * BM8640_3200 * sizeof(int32_t)); +#pragma unroll + for (int32_t k_outer = 0; k_outer < 32 / 32; ++k_outer) { + two_tbl_impl8640_3200((&(((int32_t*)CBits)[0])), (&(((int8_t*)LUT)[(k_outer * BK2 / 2 * 32)])), (&(((uint8_t*)A)[(k_outer * BK2 / 2 / 2 * BM8640_3200)]))); + } +#pragma unroll + for (int bs = 0; bs < BATCH_SIZE; bs++) { +#pragma unroll + for (int i = 0; i < BM8640_3200; i++) { + ((int32_t*)C)[i] += (int32_t)(((int32_t*)CBits)[i + bs * BM8640_3200]); + ((float*)C)[i] = (float)(((int32_t*)C)[i]) / ((float*)LUT_Scales)[bs] * ((float*)Scales)[0]; + } + } + return 0; +} + +void ggml_preprocessor(int bs, int m, int three_k, int two_k, void* B, void* LUT_Scales, void* Three_QLUT, void* Two_QLUT) { + partial_max_reset(bs, (&(((float*)LUT_Scales)[0]))); + if (m == 3200 && two_k == 0 && three_k == 8640) { + for (int32_t b = 0; b < bs; b++) { + per_tensor_quant(two_k + three_k, (&(((float*)LUT_Scales)[b])), (&(((float*)B)[b * (two_k + three_k)]))); + three_lut_ctor<8640>((&(((int8_t*)Three_QLUT)[b * three_k / 3 * 32])), (&(((float*)B)[b * (three_k + two_k)])), (&(((float*)LUT_Scales)[b]))); + two_lut_ctor<0>((&(((int8_t*)Two_QLUT)[b * two_k / 2 * 32])), (&(((float*)B)[b * (three_k + two_k) + 8640])), (&(((float*)LUT_Scales)[b]))); + } + } + else if (m == 3200 && two_k == 32 && three_k == 3168) { + for (int32_t b = 0; b < bs; b++) { + per_tensor_quant(two_k + three_k, (&(((float*)LUT_Scales)[b])), (&(((float*)B)[b * (two_k + three_k)]))); + three_lut_ctor<3168>((&(((int8_t*)Three_QLUT)[b * three_k / 3 * 32])), (&(((float*)B)[b * (three_k + two_k)])), (&(((float*)LUT_Scales)[b]))); + two_lut_ctor<32>((&(((int8_t*)Two_QLUT)[b * two_k / 2 * 32])), (&(((float*)B)[b * (three_k + two_k) + 3168])), (&(((float*)LUT_Scales)[b]))); + } + } + else if (m == 8640 && two_k == 32 && three_k == 3168) { + for (int32_t b = 0; b < bs; b++) { + per_tensor_quant(two_k + three_k, (&(((float*)LUT_Scales)[b])), (&(((float*)B)[b * (two_k + three_k)]))); + three_lut_ctor<3168>((&(((int8_t*)Three_QLUT)[b * three_k / 3 * 32])), (&(((float*)B)[b * (three_k + two_k)])), (&(((float*)LUT_Scales)[b]))); + two_lut_ctor<32>((&(((int8_t*)Two_QLUT)[b * two_k / 2 * 32])), (&(((float*)B)[b * (three_k + two_k) + 3168])), (&(((float*)LUT_Scales)[b]))); + } + } +} +void ggml_qgemm_lut(int bs, int m, int k, int BK, void* A, void* sign, void* LUT, void* Scales, void* LUT_Scales, void* C) { + if (m == 3200 && k == 8640) { + if (BK == 0) { + if (bs == 1) { + two_qgemm_lut_3200_8640<1>(A, LUT, Scales, LUT_Scales, C); + } else if (bs == 8) { + two_qgemm_lut_3200_8640<8>(A, LUT, Scales, LUT_Scales, C); + } else if (bs == 32) { + two_qgemm_lut_3200_8640<32>(A, LUT, Scales, LUT_Scales, C); + } else if (bs == 128) { + two_qgemm_lut_3200_8640<128>(A, LUT, Scales, LUT_Scales, C); + } else if (bs == 256) { + two_qgemm_lut_3200_8640<256>(A, LUT, Scales, LUT_Scales, C); + } else if (bs == 512) { + two_qgemm_lut_3200_8640<512>(A, LUT, Scales, LUT_Scales, C); + } + } + else if (BK == 8640) { + if (bs == 1) { + three_qgemm_lut_3200_8640<1>(A, sign, LUT, Scales, LUT_Scales, C); + }else if (bs == 8) { + three_qgemm_lut_3200_8640<8>(A, sign, LUT, Scales, LUT_Scales, C); + }else if (bs == 32) { + three_qgemm_lut_3200_8640<32>(A, sign, LUT, Scales, LUT_Scales, C); + }else if (bs == 128) { + three_qgemm_lut_3200_8640<128>(A, sign, LUT, Scales, LUT_Scales, C); + }else if (bs == 256) { + three_qgemm_lut_3200_8640<256>(A, sign, LUT, Scales, LUT_Scales, C); + }else if (bs == 512) { + three_qgemm_lut_3200_8640<512>(A, sign, LUT, Scales, LUT_Scales, C); + } + } + } + else if (m == 3200 && k == 3200) { + if (BK == 32) { + if (bs == 1) { + two_qgemm_lut_3200_3200<1>(A, LUT, Scales, LUT_Scales, C); + } else if (bs == 8) { + two_qgemm_lut_3200_3200<8>(A, LUT, Scales, LUT_Scales, C); + } else if (bs == 32) { + two_qgemm_lut_3200_3200<32>(A, LUT, Scales, LUT_Scales, C); + } else if (bs == 128) { + two_qgemm_lut_3200_3200<128>(A, LUT, Scales, LUT_Scales, C); + } else if (bs == 256) { + two_qgemm_lut_3200_3200<256>(A, LUT, Scales, LUT_Scales, C); + } else if (bs == 512) { + two_qgemm_lut_3200_3200<512>(A, LUT, Scales, LUT_Scales, C); + } + } + else if (BK == 3168) { + if (bs == 1) { + three_qgemm_lut_3200_3200<1>(A, sign, LUT, Scales, LUT_Scales, C); + }else if (bs == 8) { + three_qgemm_lut_3200_3200<8>(A, sign, LUT, Scales, LUT_Scales, C); + }else if (bs == 32) { + three_qgemm_lut_3200_3200<32>(A, sign, LUT, Scales, LUT_Scales, C); + }else if (bs == 128) { + three_qgemm_lut_3200_3200<128>(A, sign, LUT, Scales, LUT_Scales, C); + }else if (bs == 256) { + three_qgemm_lut_3200_3200<256>(A, sign, LUT, Scales, LUT_Scales, C); + }else if (bs == 512) { + three_qgemm_lut_3200_3200<512>(A, sign, LUT, Scales, LUT_Scales, C); + } + } + } + else if (m == 8640 && k == 3200) { + if (BK == 32) { + if (bs == 1) { + two_qgemm_lut_8640_3200<1>(A, LUT, Scales, LUT_Scales, C); + } else if (bs == 8) { + two_qgemm_lut_8640_3200<8>(A, LUT, Scales, LUT_Scales, C); + } else if (bs == 32) { + two_qgemm_lut_8640_3200<32>(A, LUT, Scales, LUT_Scales, C); + } else if (bs == 128) { + two_qgemm_lut_8640_3200<128>(A, LUT, Scales, LUT_Scales, C); + } else if (bs == 256) { + two_qgemm_lut_8640_3200<256>(A, LUT, Scales, LUT_Scales, C); + } else if (bs == 512) { + two_qgemm_lut_8640_3200<512>(A, LUT, Scales, LUT_Scales, C); + } + } + else if (BK == 3168) { + if (bs == 1) { + three_qgemm_lut_8640_3200<1>(A, sign, LUT, Scales, LUT_Scales, C); + }else if (bs == 8) { + three_qgemm_lut_8640_3200<8>(A, sign, LUT, Scales, LUT_Scales, C); + }else if (bs == 32) { + three_qgemm_lut_8640_3200<32>(A, sign, LUT, Scales, LUT_Scales, C); + }else if (bs == 128) { + three_qgemm_lut_8640_3200<128>(A, sign, LUT, Scales, LUT_Scales, C); + }else if (bs == 256) { + three_qgemm_lut_8640_3200<256>(A, sign, LUT, Scales, LUT_Scales, C); + }else if (bs == 512) { + three_qgemm_lut_8640_3200<512>(A, sign, LUT, Scales, LUT_Scales, C); + } + } + } +} + +void ggml_bitnet_transform_tensor(struct ggml_tensor * tensor) { + if (!(is_type_supported(tensor->type) && tensor->extra == nullptr)) { + return; + } + + int k = tensor->ne[0]; + int m = tensor->ne[1]; + const int lut_scales_size = 1; + int bk = 0; + int bm = 0; + + if (m == 3200 && k == 8640) { + bm = BM3200_8640; + bk = BBK3200_8640; + } +else if (m == 3200 && k == 3200) { + bm = BM3200_3200; + bk = BBK3200_3200; + } +else if (m == 8640 && k == 3200) { + bm = BM8640_3200; + bk = BBK8640_3200; + } + + const int n_tile_num = m / bm; + const int BK = bk; + uint8_t * qweights; + bitnet_float_type * scales; + + scales = (bitnet_float_type *) aligned_malloc(sizeof(bitnet_float_type)); + qweights = (uint8_t *) tensor->data; + int nbytes = (k - 256) * m / 3 * 5 / 8 + 256 * m / 2 * 4 / 8; + if (nbytes % 32 != 0) nbytes = 32 - nbytes % 32 + nbytes; + float * i2_scales = (float * )(qweights + nbytes); + scales[0] = (bitnet_float_type) i2_scales[0]; + + tensor->extra = bitnet_tensor_extras + bitnet_tensor_extras_index; + bitnet_tensor_extras[bitnet_tensor_extras_index++] = { + /* .lut_scales_size = */ lut_scales_size, + /* .BK = */ BK, + /* .n_tile_num = */ n_tile_num, + /* .qweights = */ qweights, + /* .scales = */ scales + }; +} +#endif \ No newline at end of file diff --git a/include/ggml-bitnet.h b/include/ggml-bitnet.h index 3f8571cc6..f92465ca9 100644 --- a/include/ggml-bitnet.h +++ b/include/ggml-bitnet.h @@ -14,6 +14,8 @@ typedef float bitnet_float_type; extern "C" { #endif +struct ggml_compute_params; + struct bitnet_tensor_extra { int lut_scales_size; int BK; @@ -33,6 +35,7 @@ GGML_API size_t ggml_bitnet_mul_mat_get_wsize(const struct ggml_tensor * src0, c GGML_API void ggml_bitnet_mul_mat_task_init(void * src1, void * qlut, void * lut_scales, void * lut_biases, int n, int k, int m, int bits); GGML_API void ggml_bitnet_mul_mat_task_compute(void * src0, void * scales, void * qlut, void * lut_scales, void * lut_biases, void * dst, int n, int k, int m, int bits); GGML_API void ggml_bitnet_transform_tensor(struct ggml_tensor * tensor); +GGML_API void ggml_bitnet_mul_mat(const struct ggml_compute_params * params, struct ggml_tensor * dst); GGML_API int ggml_bitnet_get_type_bits(enum ggml_type type); GGML_API void ggml_bitnet_set_n_threads(int n_threads); #if defined(GGML_BITNET_ARM_TL1) diff --git a/include/kernel_config.ini b/include/kernel_config.ini new file mode 100644 index 000000000..46468bf13 --- /dev/null +++ b/include/kernel_config.ini @@ -0,0 +1,21 @@ +[Kernels_0] +m = 3200 +k = 8640 +bm = 160 +bk = 96 +bmm = 32 + +[Kernels_1] +m = 3200 +k = 3200 +bm = 320 +bk = 96 +bmm = 32 + +[Kernels_2] +m = 8640 +k = 3200 +bm = 320 +bk = 96 +bmm = 32 + diff --git a/src/ggml-bitnet-lut.cpp b/src/ggml-bitnet-lut.cpp index beef726f7..676351ddc 100644 --- a/src/ggml-bitnet-lut.cpp +++ b/src/ggml-bitnet-lut.cpp @@ -11,6 +11,7 @@ #include "ggml-bitnet.h" #include "ggml-quants.h" +#include "ggml-cpu-impl.h" #if defined(GGML_BITNET_ARM_TL1) || defined(GGML_BITNET_X86_TL2) #include "bitnet-lut-kernels.h" @@ -19,16 +20,11 @@ #if defined(GGML_BITNET_ARM_TL1) void ggml_bitnet_init(void) { - // LOG(INFO) << "ggml_bitnet_init"; - if (initialized) { return; } initialized = true; - // if (wrapper == nullptr) { - // wrapper = new BITNET::BITNETGeMMWrapper(); - // } if (bitnet_tensor_extras == nullptr) { bitnet_tensor_extras = new bitnet_tensor_extra[GGML_BITNET_MAX_NODES]; } @@ -36,26 +32,17 @@ void ggml_bitnet_init(void) { } void ggml_bitnet_free(void) { - // LOG(INFO) << "ggml_bitnet_free"; - if (!initialized) { return; } initialized = false; - // delete wrapper; - // wrapper = nullptr; - for (size_t i = 0; i < bitnet_tensor_extras_index; i++) { - // aligned_free(bitnet_tensor_extras[i].qweights); - // aligned_free(bitnet_tensor_extras[i].scales); - } delete[] bitnet_tensor_extras; bitnet_tensor_extras = nullptr; } static bool do_permutate(enum ggml_type type) { if (type == GGML_TYPE_TL1) { - // Add additional args to decide if permuted I2 or naive I2 return false; } else { return true; @@ -65,8 +52,7 @@ static bool do_permutate(enum ggml_type type) { bool ggml_bitnet_can_mul_mat(const struct ggml_tensor * src0, const struct ggml_tensor * src1, const struct ggml_tensor * dst) { if ((is_type_supported(src0->type)) && src1->type == GGML_TYPE_F32 && - dst->type == GGML_TYPE_F32 && - src0->backend == GGML_BACKEND_TYPE_CPU) { + dst->type == GGML_TYPE_F32) { if (src1->ne[1] <= 1) { return true; } @@ -79,10 +65,9 @@ size_t ggml_bitnet_mul_mat_get_wsize(const struct ggml_tensor * src0, const stru const size_t ne10 = src1->ne[0]; const size_t ne11 = src1->ne[1]; const int bits = ggml_bitnet_get_type_bits(src0->type); - + size_t wsize = ne10 * ne11 * 15 * sizeof(int8_t) + 1 * ne11 * 2 * sizeof(bitnet_float_type); if (sizeof(bitnet_float_type) == 2) { - // Need fp32 to fp16 conversion wsize += std::max(ne10, ne01) * ne11 * sizeof(bitnet_float_type); } wsize = ((wsize - 1) / 64 + 1) * 64; @@ -100,19 +85,61 @@ int ggml_bitnet_get_type_bits(enum ggml_type type) { } } +void ggml_bitnet_mul_mat(const struct ggml_compute_params * params, struct ggml_tensor * dst) { + const struct ggml_tensor * src0 = dst->src[0]; + const struct ggml_tensor * src1 = dst->src[1]; + + const size_t ne00 = src0->ne[0]; + const size_t ne01 = src0->ne[1]; + const size_t ne10 = src1->ne[0]; + const size_t ne11 = src1->ne[1]; + + const int ith = params->ith; + const int nth = params->nth; + const int bits = ggml_bitnet_get_type_bits(src0->type); + + struct bitnet_tensor_extra * extra = (struct bitnet_tensor_extra *)src0->extra; + GGML_ASSERT(extra != nullptr); + + char * wdata = (char *)params->wdata; + const size_t wsize_per_thread = ggml_bitnet_mul_mat_get_wsize(src0, src1, dst); + + int8_t * qlut = (int8_t *)(wdata); + bitnet_float_type * lut_scales = (bitnet_float_type *)(qlut + ne10 * ne11 * 15); + bitnet_float_type * lut_biases = (bitnet_float_type *)(lut_scales + ne11); + + if (ith == 0) { + ggml_bitnet_mul_mat_task_init( + (void *)((char *)src1->data), + (void *)qlut, + (void *)lut_scales, + (void *)lut_biases, + ne10, ne00, ne11, bits); + } + + // barrier + if (nth > 1) { + ggml_barrier(params->threadpool); + } + + ggml_bitnet_mul_mat_task_compute( + (void *)extra->qweights, + (void *)extra->scales, + (void *)qlut, + (void *)lut_scales, + (void *)lut_biases, + (void *)((char *)dst->data), + ne10, ne00, ne11, bits); +} + #endif #if defined(GGML_BITNET_X86_TL2) void ggml_bitnet_init(void) { - // LOG(INFO) << "ggml_bitnet_init"; - if (initialized) { return; } initialized = true; - // if (wrapper == nullptr) { - // wrapper = new BITNET::BITNETGeMMWrapper(); - // } if (bitnet_tensor_extras == nullptr) { bitnet_tensor_extras = new bitnet_tensor_extra[GGML_BITNET_MAX_NODES]; } @@ -120,19 +147,11 @@ void ggml_bitnet_init(void) { } void ggml_bitnet_free(void) { - // LOG(INFO) << "ggml_bitnet_free"; - if (!initialized) { return; } initialized = false; - // delete wrapper; - // wrapper = nullptr; - for (size_t i = 0; i < bitnet_tensor_extras_index; i++) { - // aligned_free(bitnet_tensor_extras[i].qweights); - // aligned_free(bitnet_tensor_extras[i].scales); - } delete[] bitnet_tensor_extras; bitnet_tensor_extras = nullptr; } @@ -140,8 +159,7 @@ void ggml_bitnet_free(void) { bool ggml_bitnet_can_mul_mat(const struct ggml_tensor * src0, const struct ggml_tensor * src1, const struct ggml_tensor * dst) { if ((is_type_supported(src0->type)) && src1->type == GGML_TYPE_F32 && - dst->type == GGML_TYPE_F32 && - src0->backend == GGML_BACKEND_TYPE_CPU) { + dst->type == GGML_TYPE_F32) { return true; } return false; @@ -151,10 +169,9 @@ size_t ggml_bitnet_mul_mat_get_wsize(const struct ggml_tensor * src0, const stru const size_t ne01 = src0->ne[1]; const size_t ne10 = src1->ne[0]; const size_t ne11 = src1->ne[1]; - + size_t wsize = ne10 * ne11 * 11 * sizeof(int8_t) + 2 * ne11 * 2 * sizeof(bitnet_float_type); if (sizeof(bitnet_float_type) == 2) { - // Need fp32 to fp16 conversion wsize += std::max(ne10, ne01) * ne11 * sizeof(bitnet_float_type); } wsize = ((wsize - 1) / 64 + 1) * 64; @@ -171,4 +188,5 @@ int ggml_bitnet_get_type_bits(enum ggml_type type) { return 0; } } -#endif \ No newline at end of file + +#endif diff --git a/utils/convert-hf-to-gguf-bitnet.py b/utils/convert-hf-to-gguf-bitnet.py index 23e84384c..e5b02d29e 100644 --- a/utils/convert-hf-to-gguf-bitnet.py +++ b/utils/convert-hf-to-gguf-bitnet.py @@ -153,8 +153,12 @@ def set_gguf_parameters(self): self.gguf_writer.add_expert_used_count(n_experts_used) logger.info(f"gguf: experts used count = {n_experts_used}") - self.gguf_writer.add_file_type(self.ftype) - logger.info(f"gguf: file type = {self.ftype}") + # Map ggml tensor type to llama ftype for general.file_type metadata + ftype_val = self.ftype + if self.ftype == gguf.GGMLQuantizationType.I2_S: + ftype_val = 40 # LLAMA_FTYPE_MOSTLY_I2_S (matches official model) + self.gguf_writer.add_file_type(ftype_val) + logger.info(f"gguf: file type = {ftype_val}") def write_tensors(self): block_count = self.hparams.get("n_layers", self.hparams.get("num_hidden_layers", self.hparams.get("n_layer"))) @@ -659,6 +663,70 @@ def preprocess_weights_tl2( weight.shape[0]), mode='constant', constant_values=0) return weight +def quantize_to_i2_s(w: np.ndarray, override_scale: float = None) -> np.ndarray: + """Quantize a float weight matrix to I2_S ternary format. + + I2_S format: packed ternary bytes (4 values per byte) + 32-byte tail with f32 scale. + Dequantization: y = scale * ternary, where ternary in {-1, 0, +1}. + + Args: + w: float weight tensor of shape (M, K) + override_scale: if provided, use this as the I2_S scale instead of computing from data. + For offline-quantized BitNet models, this should be the weight_scale value. + """ + M, K = w.shape + n = M * K + w_flat = w.flatten().astype(np.float32) + + # Compute scale for I2_S dequantization + if override_scale is not None: + # override_scale is weight_scale from offline-quantized models ≈ mean(|original_bf16_weights|) + # Use it directly as the I2_S scale + scale = np.float32(override_scale) + # Weights are already ternary {-1, 0, 1}, use directly + q_float = w_flat + else: + # w_flat comes from weight_quant: values are ±scale or 0 + # Use the first nonzero absolute value as scale (matches C quantize_i2_s) + nonzero = np.abs(w_flat[np.abs(w_flat) > 1e-6]) + if len(nonzero) > 0: + scale = np.float32(nonzero[0]) + else: + scale = np.float32(1e-5) + # Quantize to ternary {-1, 0, 1} + inv_scale = 1.0 / scale + q_float = np.round(w_flat * inv_scale).clip(-1, 1) + + # Map ternary {-1, 0, 1} -> I2_S encoding {0, 1, 2} + q = np.ones(n, dtype=np.uint8) # default to 1 (zero) + q[q_float > 0.5] = 2 # +1 -> 2 + q[q_float < -0.5] = 0 # -1 -> 0 + + # Pack into I2_S layout: 128-value blocks, interleaved into 32 bytes + pad_len = (128 - n % 128) % 128 + if pad_len: + q = np.pad(q, (0, pad_len), constant_values=1) + + n_padded = len(q) + n_blocks = n_padded // 128 + q = q.reshape(n_blocks, 4, 32) + + packed = (q[:, 0, :].astype(np.uint8) << 6) | \ + (q[:, 1, :].astype(np.uint8) << 4) | \ + (q[:, 2, :].astype(np.uint8) << 2) | \ + (q[:, 3, :].astype(np.uint8)) + packed = packed.reshape(-1).astype(np.uint8) + + # I2_S format: packed_bytes + 32-byte aligned tail (scale in first 4 bytes) + packed_size = n // 4 + total_size = packed_size + 32 + result = np.zeros(total_size, dtype=np.uint8) + result[:len(packed)] = packed[:packed_size] + result[packed_size:packed_size+4] = np.frombuffer(scale.tobytes(), dtype=np.uint8) + + return result + + def transform_to_tl1(x: np.ndarray): scale = np.max(np.abs(x)) # res = np.round(x / scale + 2).astype(np.uint8) @@ -813,10 +881,17 @@ def write_tensors(self): data = data.astype(np.float16) data_qtype = gguf.GGMLQuantizationType.F16 - if data_qtype is None: # by default, convert to float32 - if data_dtype != np.float32: - data = data.astype(np.float32) - data_qtype = gguf.GGMLQuantizationType.F32 + if data_qtype is None: # by default + # For I2_S/TL models, keep non-quantized 2D weights (e.g. embed) as F16 instead of F32 + if self.ftype in (gguf.GGMLQuantizationType.I2_S, gguf.GGMLQuantizationType.TL1, gguf.GGMLQuantizationType.TL2) \ + and n_dims >= 2 and not new_name.endswith("_norm.weight"): + if data_dtype != np.float16: + data = data.astype(np.float16) + data_qtype = gguf.GGMLQuantizationType.F16 + else: + if data_dtype != np.float32: + data = data.astype(np.float32) + data_qtype = gguf.GGMLQuantizationType.F32 shape = data_shape # shape = gguf.quant_shape_from_byte_shape(data.shape, data_qtype) if data.dtype == np.uint8 else data.shape @@ -952,20 +1027,30 @@ def prepare_tensors(self): raise ValueError(f"Unprocessed experts: {experts}") -@Model.register("BitnetForCausalLM") +@Model.register("BitnetForCausalLM", "BitNetForCausalLM") class BitnetModel(Model): - model_arch = gguf.MODEL_ARCH.BITNET + model_arch = gguf.MODEL_ARCH.BITNET_B158 def set_vocab(self): - self._set_vocab_sentencepiece() + try: + self._set_vocab_sentencepiece() + except FileNotFoundError: + try: + self._set_vocab_llama_hf() + except (FileNotFoundError, TypeError): + self._set_vocab_gpt2() def set_gguf_parameters(self): super().set_gguf_parameters() self.gguf_writer.add_vocab_size(self.hparams["vocab_size"]) - self.gguf_writer.add_rope_scaling_type(gguf.RopeScalingType.LINEAR) - self.gguf_writer.add_rope_scaling_factor(1.0) + # rope dimension count (required for correct positional encoding) + if "head_dim" in self.hparams: + rope_dim = self.hparams["head_dim"] + else: + rope_dim = self.hparams["hidden_size"] // self.hparams["num_attention_heads"] + self.gguf_writer.add_rope_dimension_count(rope_dim) def weight_quant(self, weight): dtype = weight.dtype @@ -976,7 +1061,10 @@ def weight_quant(self, weight): def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: # quant weight to i2 (in fp16) - if name.endswith(("q_proj.weight", "k_proj.weight", "v_proj.weight", + # Skip weight_quant for offline-quantized models — weights are already ternary {-1,0,1} + # weight_quant would scale them to small floats (e.g. ±0.39), breaking quantize_to_i2_s + if not getattr(self, '_has_offline_quant', False) and \ + name.endswith(("q_proj.weight", "k_proj.weight", "v_proj.weight", "down_proj.weight", "up_proj.weight", "gate_proj.weight", "o_proj.weight")): data_torch = self.weight_quant(data_torch) @@ -986,15 +1074,40 @@ def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iter def write_tensors(self): max_name_len = max(len(s) for _, s in self.tensor_map.mapping.values()) + len(".weight,") + # First pass: collect weight_scale tensors for offline-quantized models + scale_map = dict() + for name, data_torch in self.get_tensors(): + if name.endswith("weight_scale"): + data_torch = data_torch.to(torch.float32) + name = name.replace(".weight_scale", "") + scale_map[name] = data_torch + + self._has_offline_quant = len(scale_map) > 0 + for name, data_torch in self.get_tensors(): + # skip weight_scale tensors + if name.endswith("weight_scale"): + continue # we don't need these if name.endswith((".attention.masked_bias", ".attention.bias", ".rotary_emb.inv_freq")): continue old_dtype = data_torch.dtype + # Handle offline-quantized weights (uint8 packed with weight_scale) + if name.replace(".weight", "") in scale_map: + data_torch = data_torch.to(torch.uint8) + origin_shape = data_torch.shape + shift = torch.tensor([0, 2, 4, 6], dtype=torch.uint8).reshape((4, *(1 for _ in range(len(origin_shape))))) + data_torch = data_torch.unsqueeze(0).expand((4, *origin_shape)) >> shift + data_torch = data_torch & 3 + data_torch = (data_torch.float() - 1).reshape((origin_shape[0] * 4, *origin_shape[1:])) + # For F16/F32 output: divide by weight_scale to get full float values + # For I2_S output: keep as ternary {-1,0,1}, scale is passed separately to quantize_to_i2_s + if self.ftype not in (gguf.GGMLQuantizationType.I2_S, gguf.GGMLQuantizationType.TL1, gguf.GGMLQuantizationType.TL2): + data_torch = data_torch / scale_map[name.replace(".weight", "")].float() # convert any unsupported data types to float32 - if data_torch.dtype not in (torch.float16, torch.float32): + elif data_torch.dtype not in (torch.float16, torch.float32): data_torch = data_torch.to(torch.float32) # use the first number-like part of the tensor name as the block id @@ -1048,7 +1161,13 @@ def write_tensors(self): i2_scale = None if self.ftype != gguf.GGMLQuantizationType.F32 and extra_f16 and not extra_f32: - if self.ftype == gguf.GGMLQuantizationType.TL1 and suit_i2: + if self.ftype == gguf.GGMLQuantizationType.I2_S and suit_i2: + data_qtype = gguf.GGMLQuantizationType.I2_S + # Use original weight_scale if available (offline-quantized models) + orig_scale = scale_map.get(name.replace(".weight", "")) + override_scale = orig_scale.item() if orig_scale is not None else None + data = quantize_to_i2_s(data, override_scale=override_scale) + elif self.ftype == gguf.GGMLQuantizationType.TL1 and suit_i2: data, i2_scale = transform_to_tl1(data) assert data.dtype == np.uint8 assert i2_scale.dtype == np.float32 @@ -1063,10 +1182,17 @@ def write_tensors(self): data = data.astype(np.float16) data_qtype = gguf.GGMLQuantizationType.F16 - if data_qtype is None: # by default, convert to float32 - if data_dtype != np.float32: - data = data.astype(np.float32) - data_qtype = gguf.GGMLQuantizationType.F32 + if data_qtype is None: # by default + # For I2_S/TL models, keep non-quantized 2D weights (e.g. embed) as F16 instead of F32 + if self.ftype in (gguf.GGMLQuantizationType.I2_S, gguf.GGMLQuantizationType.TL1, gguf.GGMLQuantizationType.TL2) \ + and n_dims >= 2 and not new_name.endswith("_norm.weight"): + if data_dtype != np.float16: + data = data.astype(np.float16) + data_qtype = gguf.GGMLQuantizationType.F16 + else: + if data_dtype != np.float32: + data = data.astype(np.float32) + data_qtype = gguf.GGMLQuantizationType.F32 shape = data_shape # shape = gguf.quant_shape_from_byte_shape(data.shape, data_qtype) if data.dtype == np.uint8 else data.shape @@ -1087,6 +1213,7 @@ def write_tensors(self): ftype_map = { "f32": gguf.GGMLQuantizationType.F32, "f16": gguf.GGMLQuantizationType.F16, + "i2_s": gguf.GGMLQuantizationType.I2_S, "tl1" : gguf.GGMLQuantizationType.TL1, "tl2" : gguf.GGMLQuantizationType.TL2, } From ce530b0fd0a64bf815c614fc6b59b4e164910557 Mon Sep 17 00:00:00 2001 From: isHuangXin Date: Mon, 25 May 2026 03:01:27 +0800 Subject: [PATCH 3/8] Add bitnet-embeddings-270m model adaptation with F16 and I2_S GGUF conversion - Add LLM_ARCH_GEMMA3 in llama.cpp for gemma3_text model type (embedding scaling, GELU, post-attn/post-FFN norms, GQA) - Add GGUF conversion support for Gemma3-based 270m models (SPM tokenizer, RMSNorm w+1 offset, arch-specific tensor mapping) - Add tokenizer hash for multilingual-e5-0.6b-260311 - Add conversion documentation --- docs/bitnet-embeddings-gguf-conversion.md | 410 ++++++++++++++++++ ...bitnet-embeddings-qwen3-gguf-conversion.md | 302 ------------- utils/convert-bitnet-embedding-to-gguf.py | 234 ++++++++-- 3 files changed, 599 insertions(+), 347 deletions(-) create mode 100644 docs/bitnet-embeddings-gguf-conversion.md delete mode 100644 docs/bitnet-embeddings-qwen3-gguf-conversion.md diff --git a/docs/bitnet-embeddings-gguf-conversion.md b/docs/bitnet-embeddings-gguf-conversion.md new file mode 100644 index 000000000..a4ee919d6 --- /dev/null +++ b/docs/bitnet-embeddings-gguf-conversion.md @@ -0,0 +1,410 @@ +# BitNet Embeddings GGUF Conversion Implementation + +## 1. Background + +BitNet embedding models apply per-projection RMSNorm (`BitLinear`) before each linear projection (q/k/v/o/gate/up/down). Each projection has a `.norm.weight` that applies RMSNorm to the input **before** the matmul: + +``` +x → RMSNorm(x, norm.weight) → activation_quant(8bit) → matmul(weight_quant(ternary)) +``` + +This pattern does **not** exist in any standard llama.cpp architecture: +- Standard Qwen3/Gemma3: no per-projection norms +- Standard BitNet: has `attn_sub_norm`/`ffn_sub_norm` at different positions (after attention/gate*up, not before each projection) + +Currently two base architectures are supported: + +| | bitnet-embeddings-0.6b (Qwen3) | bitnet-embeddings-270m (Gemma3) | +|---|---|---| +| Architecture | `Qwen3Model` | `Gemma3TextModel` | +| hidden_size | 1024 | 640 | +| num_attention_heads | 16 | 4 | +| num_key_value_heads | 8 | 1 | +| head_dim | 128 (note: != hidden_size/num_heads = 64) | 256 (note: != hidden_size/num_heads = 160) | +| intermediate_size | 3072 | 2048 | +| num_hidden_layers | 28 | 18 | +| hidden_activation | SiLU | gelu_pytorch_tanh | +| vocab_size | 151936 | 262144 | +| rope_theta | 1000000 | 10000.0 | +| rms_norm_eps | 1e-06 | 1e-06 | +| query_pre_attn_scalar | N/A | 256 | +| tie_word_embeddings | true | true | + +### Gemma3 vs Qwen3 Key Differences + +| Feature | Qwen3 | Gemma3 | +|---------|-------|--------| +| Post-attn norm | No | Yes (`post_attention_norm`) | +| Post-FFW norm | No | Yes (`post_ffw_norm`) | +| Pre-FFW norm naming | `post_attention_layernorm` → `ffn_norm` | `pre_feedforward_layernorm` → `ffn_norm` | +| QK head norms | Yes | Yes | +| Activation | SiLU | GELU | +| Embedding scaling | No | sqrt(n_embd) | +| EOS token override | Yes (`<\|endoftext\|>` 151643) | No (auto from tokenizer) | + +### Per-Layer Tensors (7 extra norm tensors per layer) + +| Tensor | Qwen3 Shape | Gemma3 Shape | +|--------|-------------|--------------| +| `self_attn.q_proj.norm.weight` | [1024] | [640] | +| `self_attn.k_proj.norm.weight` | [1024] | [640] | +| `self_attn.v_proj.norm.weight` | [1024] | [640] | +| `self_attn.o_proj.norm.weight` | [2048] | [1024] | +| `mlp.gate_proj.norm.weight` | [1024] | [640] | +| `mlp.up_proj.norm.weight` | [1024] | [640] | +| `mlp.down_proj.norm.weight` | [3072] | [2048] | + +--- + +## 2. GGUF Tensor Name Mapping + +### Common Tensors (both architectures) + +| HF Name | GGUF Name | Notes | +|----------|-----------|-------| +| `embed_tokens.weight` | `token_embd.weight` | | +| `norm.weight` | `output_norm.weight` | | +| `layers.{i}.input_layernorm.weight` | `blk.{i}.attn_norm.weight` | | +| `layers.{i}.self_attn.q_proj.weight` | `blk.{i}.attn_q.weight` | | +| `layers.{i}.self_attn.k_proj.weight` | `blk.{i}.attn_k.weight` | | +| `layers.{i}.self_attn.v_proj.weight` | `blk.{i}.attn_v.weight` | | +| `layers.{i}.self_attn.o_proj.weight` | `blk.{i}.attn_output.weight` | | +| `layers.{i}.self_attn.q_norm.weight` | `blk.{i}.attn_q_norm.weight` | QK head norm | +| `layers.{i}.self_attn.k_norm.weight` | `blk.{i}.attn_k_norm.weight` | QK head norm | +| `layers.{i}.self_attn.q_proj.norm.weight` | `blk.{i}.attn_q_norm_in.weight` | BitNet per-projection | +| `layers.{i}.self_attn.k_proj.norm.weight` | `blk.{i}.attn_k_norm_in.weight` | BitNet per-projection | +| `layers.{i}.self_attn.v_proj.norm.weight` | `blk.{i}.attn_v_norm_in.weight` | BitNet per-projection | +| `layers.{i}.self_attn.o_proj.norm.weight` | `blk.{i}.attn_output_norm_in.weight` | BitNet per-projection | +| `layers.{i}.mlp.gate_proj.weight` | `blk.{i}.ffn_gate.weight` | | +| `layers.{i}.mlp.up_proj.weight` | `blk.{i}.ffn_up.weight` | | +| `layers.{i}.mlp.down_proj.weight` | `blk.{i}.ffn_down.weight` | | +| `layers.{i}.mlp.gate_proj.norm.weight` | `blk.{i}.ffn_gate_norm_in.weight` | BitNet per-projection | +| `layers.{i}.mlp.up_proj.norm.weight` | `blk.{i}.ffn_up_norm_in.weight` | BitNet per-projection | +| `layers.{i}.mlp.down_proj.norm.weight` | `blk.{i}.ffn_down_norm_in.weight` | BitNet per-projection | + +### Architecture-Specific Tensors + +**Qwen3:** + +| HF Name | GGUF Name | +|----------|-----------| +| `layers.{i}.post_attention_layernorm.weight` | `blk.{i}.ffn_norm.weight` | + +**Gemma3 (additional):** + +| HF Name | GGUF Name | +|----------|-----------| +| `layers.{i}.post_attention_layernorm.weight` | `blk.{i}.post_attention_norm.weight` | +| `layers.{i}.pre_feedforward_layernorm.weight` | `blk.{i}.ffn_norm.weight` | +| `layers.{i}.post_feedforward_layernorm.weight` | `blk.{i}.post_ffw_norm.weight` | + +--- + +## 3. Conversion Script + +### `utils/convert-bitnet-embedding-to-gguf.py` + +Unified standalone conversion script (safetensors → GGUF) that **auto-detects** the model architecture from `config.json`'s `model_type` field (`qwen3` or `gemma3_text`). Key features: + +- Hardcoded HF→GGUF tensor name mapping (no dependency on llama.cpp's Python converter) +- Auto-detection of architecture and GGUF arch string (`qwen3` / `gemma3`) +- Supports three output types: + - `--outtype f32`: all weights in float32 + - `--outtype f16`: 2D weights and embeddings as float16, norms as float16 + - `--outtype i2_s`: ternary weights packed in I2_S layout, non-ternary weights as float16 +- Writes `key_length` and `value_length` metadata for correct head_dim (critical: head_dim != hidden_size/num_heads for both models, default calculation would give wrong values) +- BPE tokenizer handling with per-architecture pre-tokenizer hash verification: + - Qwen3: GPT-2 BPE tokenizer + - Gemma3: GemmaTokenizerFast (BPE) +- Pooling type auto-detection from `modules.json` / `1_Pooling/config.json` (sentence-transformers convention) +- Architecture-specific tokenizer handling: + - Qwen3: EOS token override (`<|endoftext|>` 151643) + `add_eos_token(True)` for last-token pooling + - Gemma3: EOS token auto-set by SpecialVocab from tokenizer_config.json (eos_token_id=1) +- Gemma3: writes `query_pre_attn_scalar = 256` for correct attention scaling + +### I2_S Ternary Packing + +The I2_S format packs ternary weights {-1, 0, +1} into 2-bit representation: + +- Quantization: `scale = 1/mean(|w|)`, `q = round(w * scale).clamp(-1, 1)` +- Encoding: `-1 → 0`, `0 → 1`, `+1 → 2` +- Every 128 values form a block, packed into 32 bytes +- Each byte stores 4 values: `byte = (c0 << 6) | (c1 << 4) | (c2 << 2) | c3` +- Scale (float32) is appended at the end of the packed data buffer + +### Tensor Type Assignment + +| Tensor Type | f16 mode | i2_s mode | +|-------------|----------|-----------| +| 2D linear weights | float16 | I2_S ternary packed | +| Embedding weights | float16 | float16 | +| Norm weights (1D) | float16 | float16 | + +Note: `output.weight` (lm_head) is skipped for embedding models — it is not needed (no token generation). + +--- + +## 4. C++ Modifications (`3rdparty/llama.cpp/src/llama.cpp`) + +### 4.1 New Architecture: `LLM_ARCH_GEMMA3` + +Added after `LLM_ARCH_GEMMA2` in the `llm_arch` enum with name mapping `"gemma3"`. Qwen3 (`LLM_ARCH_QWEN3`) was added by the 0.6b adaptation. + +### 4.2 New Tensor Enums (shared across architectures) + +Added 7 new entries after `LLM_TENSOR_FFN_SUB_NORM`: + +```cpp +LLM_TENSOR_ATTN_Q_NORM_IN, +LLM_TENSOR_ATTN_K_NORM_IN, +LLM_TENSOR_ATTN_V_NORM_IN, +LLM_TENSOR_ATTN_OUT_NORM_IN, +LLM_TENSOR_FFN_GATE_NORM_IN, +LLM_TENSOR_FFN_UP_NORM_IN, +LLM_TENSOR_FFN_DOWN_NORM_IN, +``` + +### 4.3 Layer Struct Fields + +Added to `struct llama_layer`: + +```cpp +struct ggml_tensor * attn_q_norm_in; +struct ggml_tensor * attn_k_norm_in; +struct ggml_tensor * attn_v_norm_in; +struct ggml_tensor * attn_out_norm_in; +struct ggml_tensor * ffn_gate_norm_in; +struct ggml_tensor * ffn_up_norm_in; +struct ggml_tensor * ffn_down_norm_in; +``` + +### 4.4 Tensor Name Mappings + +Both `LLM_ARCH_QWEN3` and `LLM_ARCH_GEMMA3` include the 7 per-projection norm tensor mappings plus standard tensors (see Section 2 for full mapping). Key differences: + +- Qwen3 includes `LLM_TENSOR_OUTPUT` (`"output"`); Gemma3 does not (uses tied embeddings directly) +- Gemma3 additionally includes `LLM_TENSOR_ATTN_POST_NORM` (`"blk.%d.post_attention_norm"`) and `LLM_TENSOR_FFN_POST_NORM` (`"blk.%d.post_ffw_norm"`) + +### 4.5 load_tensors + +Both architectures load the 7 per-projection norm tensors as optional (`TENSOR_NOT_REQUIRED`): + +```cpp +layer.attn_q_norm_in = create_tensor(tn(...), {n_embd}, TENSOR_NOT_REQUIRED); +layer.attn_k_norm_in = create_tensor(tn(...), {n_embd}, TENSOR_NOT_REQUIRED); +layer.attn_v_norm_in = create_tensor(tn(...), {n_embd}, TENSOR_NOT_REQUIRED); +layer.attn_out_norm_in = create_tensor(tn(...), {n_embd_head_k * n_head}, TENSOR_NOT_REQUIRED); +layer.ffn_gate_norm_in = create_tensor(tn(...), {n_embd}, TENSOR_NOT_REQUIRED); +layer.ffn_up_norm_in = create_tensor(tn(...), {n_embd}, TENSOR_NOT_REQUIRED); +layer.ffn_down_norm_in = create_tensor(tn(...), {n_ff}, TENSOR_NOT_REQUIRED); +``` + +Note: `o_proj.norm` input dimension is `n_embd_head_k * n_head` (Qwen3: 2048, Gemma3: 1024), `down_proj.norm` input dimension is `n_ff` (Qwen3: 3072, Gemma3: 2048). + +Both graph functions use the same per-projection norm pattern. The logic is fully backward compatible — when no `*_norm_in` tensors exist, behavior is identical to the original. + +**Attention per-projection norms:** +``` +// Before Q/K/V matmul: +if (layer.attn_q_norm_in) { + cur_q = ggml_rms_norm(ctx, cur, hparams.f_norm_rms_eps); + cur_q = ggml_mul(ctx, cur_q, layer.attn_q_norm_in); +} else { + cur_q = cur; +} +Qcur = ggml_mul_mat(ctx, layer.wq, cur_q); +// QK head norms applied after projection +Qcur = ggml_rms_norm(ctx, Qcur, hparams.f_norm_rms_eps); +Qcur = ggml_mul(ctx, Qcur, layer.attn_q_norm); +``` + +**O_proj norm** requires special handling because `llm_build_kv()` normally applies `wo` internally. Solution: pass `wo=NULL` to `llm_build_kv()`, then apply norm + wo manually: + +``` +cur = llm_build_kv(..., wo=NULL, ...); // returns attention output without o_proj +if (layer.attn_out_norm_in) { + cur = ggml_rms_norm(ctx, cur, hparams.f_norm_rms_eps); + cur = ggml_mul(ctx, cur, layer.attn_out_norm_in); +} +cur = ggml_mul_mat(ctx, layer.wo, cur); +``` + +**FFN per-projection norms:** +``` +// Instead of llm_build_ffn(), manually: +if (layer.ffn_gate_norm_in) { + tmp_gate = rms_norm(cur) * gate_norm_in; +} else { + tmp_gate = cur; +} +tmp_gate = matmul(gate_proj, tmp_gate); +tmp_gate = activation(tmp_gate); // SiLU for Qwen3, GELU for Gemma3 +// Similarly for up_proj +tmp = tmp_gate * tmp_up; + +if (layer.ffn_down_norm_in) { + tmp = rms_norm(tmp) * down_norm_in; +} +cur = matmul(down_proj, tmp); +``` + +**Gemma3-specific differences:** +- Embedding scaling by `sqrt(n_embd)` (Gemma convention) +- GELU activation instead of SiLU +- Post-attention and post-FFN layer norms +- `query_pre_attn_scalar` for attention scaling + +--- + +## 5. GGUF Conversion Process + +Each model variant requires two GGUF files from **two different source models**: + +### 5.1 Qwen3 (0.6b) + +| GGUF Output | Source Model | Description | +|-------------|-------------|-------------| +| `embeddings-0.6b-f16.gguf` | `multilingual-e5-0.6b` (standard Qwen3) | F16 baseline | +| `bitnet-embeddings-0.6b-f16-i2_s.gguf` | `bitnet-embeddings-0.6b` (BitNet ternary) | I2_S ternary packed | + +**F16 (from standard Qwen3 model):** +```bash +python3 utils/convert-bitnet-embedding-to-gguf.py \ + /path/to/multilingual-e5-0.6b \ + --outtype f16 \ + --outfile embeddings-0.6b-f16.gguf +``` + +What happens: +1. Load `model.safetensors` (standard Qwen3 weights, bfloat16) +2. Convert all 2D weights (projections, embeddings) to float16 +3. Convert norm weights to float16 +4. Write GGUF with `qwen3` architecture metadata and tokenizer + +**Output:** ~1.11 GiB (595.78M params) + +**I2_S (from BitNet model):** +```bash +python3 utils/convert-bitnet-embedding-to-gguf.py \ + /path/to/bitnet-embeddings-0.6b \ + --outfile bitnet-embeddings-0.6b-f16-i2_s.gguf --outtype i2_s +``` + +What happens: +1. Load `model.safetensors` (BitNet ternary weights, bfloat16) +2. Map HF tensor names to GGUF names, including 7 extra `*_norm_in` tensors per layer +3. For each 2D linear weight: quantize to I2_S ternary packed format +4. Keep embeddings (`token_embd.weight`) in float16 +5. Keep all norm weights in float16 +6. Skip `output.weight` (lm_head, not needed for embedding models) +7. Write GGUF with `I2_S` type tag for quantized tensors + +**Output:** ~699 MiB (~50% of F16 size) + +### 5.2 Gemma3 (270m) + +| GGUF Output | Source Model | Description | +|-------------|-------------|-------------| +| `multilingual-e5-270m-f16.gguf` | `multilingual-e5-270m-260311` (standard Gemma3) | F16 baseline | +| `bitnet-embeddings-270m-i2_s.gguf` | `bitnet-embeddings-270m` (BitNet ternary) | I2_S ternary packed | + +**F16 (from standard Gemma3 model):** +```bash +python3 utils/convert-bitnet-embedding-to-gguf.py \ + /path/to/multilingual-e5-270m-260311 \ + --outtype f16 +``` + +What happens: +1. Load `model.safetensors` (standard Gemma3 weights, bfloat16) +2. Convert all 2D weights (projections, embeddings) to float16 +3. Convert norm weights to float16 +4. Write GGUF with `gemma3` architecture metadata and tokenizer + +**I2_S (from BitNet model):** +```bash +python3 utils/convert-bitnet-embedding-to-gguf.py \ + /path/to/bitnet-embeddings-270m \ + --outtype i2_s +``` + +What happens: +1. Load `model.safetensors` (BitNet ternary weights, bfloat16) +2. Map HF tensor names to GGUF names, including 7 extra `*_norm_in` tensors per layer +3. For each 2D linear weight: quantize to I2_S ternary packed format +4. Keep embeddings (`token_embd.weight`) in float16 +5. Keep all norm weights in float16 +6. Skip `output.weight` (lm_head, not needed for embedding models) +7. Write GGUF with `I2_S` type tag for quantized tensors + +### 5.3 Why Two Different Source Models? + +- `multilingual-e5-*` is the **teacher/baseline model** with standard float weights, used as the F16 performance reference +- `bitnet-embeddings-*` is the **1-bit quantized student model** with ternary weights and per-projection BitLinear norms, converted to I2_S for efficient CPU inference +- Benchmarking compares both to measure the throughput gain and quality trade-off of ternary quantization + +### 5.4 Tensor Type Summary + +| Tensor | F16 (baseline) | I2_S (BitNet) | +|--------|----------------|---------------| +| Linear projections (q/k/v/o/gate/up/down) | float16 | I2_S (2-bit packed + float32 scale) | +| Embedding (`token_embd.weight`) | float16 | float16 | +| Per-projection norms (`*_norm_in`) | N/A (not present) | float16 | +| Layer norms (attn_norm, ffn_norm, etc.) | float16 | float16 | +| QK head norms (`attn_q_norm`, `attn_k_norm`) | float16 | float16 | +| `output.weight` (lm_head) | skipped | skipped | + +--- + +## 6. Additional Changes + +### 6.1 ggml.c: F16 Norm Weight Support + +Added `ggml_compute_forward_mul_f32_f16()` function to support element-wise multiplication where norm weights are stored in float16. Modified `ggml_compute_forward_mul()` to dispatch based on `src1->type`. + +### 6.2 gguf-py: I2_S Type + +Added `I2_S = 36` to `GGMLQuantizationType` enum and `(4, 1)` quant size in `constants.py`. + +### 6.3 CMakeLists.txt: BitNet LUT Kernels Guard + +Guarded `bitnet-lut-kernels.h` include with `if (GGML_BITNET_ARM_TL1 OR GGML_BITNET_X86_TL2)` to prevent build errors when LUT kernels are not available. + +### 6.4 ggml-bitnet-mad.cpp: AVX512 SIMD + +Added AVX512BW SIMD paths for I2_S dot product functions: +- `ggml_vec_dot_i2_i8_s_1x1` +- `ggml_vec_dot_i2_i8_s_1xN` +- `ggml_vec_dot_i2_i8_s_Nx1` + +--- + +## 7. Build and Run + +```bash +# Build with BitNet repo (includes I2_S support) +cmake -S /path/to/BitNet -B build -DCMAKE_BUILD_TYPE=Release +cmake --build build --target llama-embedding llama-bench -j$(nproc) + +# Run embedding inference (Qwen3 example) +build/bin/llama-embedding -m bitnet-embeddings-0.6b-f16-i2_s.gguf \ + -p "hello world" --embd-normalize 2 --embd-output-format array + +# Run embedding inference (Gemma3 example) +build/bin/llama-embedding -m bitnet-embeddings-270m-i2_s.gguf \ + -p "hello world" --embd-normalize 2 --embd-output-format array + +# Benchmark: F16 vs I2_S (Qwen3) +build/bin/llama-bench -m embeddings-0.6b-f16.gguf \ + -t 8 -p 128,256,512,1024,2048 -n 32,64 -r 3 -ngl 0 + +build/bin/llama-bench -m bitnet-embeddings-0.6b-f16-i2_s.gguf \ + -t 8 -p 128,256,512,1024,2048 -n 32,64 -r 3 -ngl 0 + +# Benchmark: F16 vs I2_S (Gemma3) +build/bin/llama-bench -m multilingual-e5-270m-f16.gguf \ + -t 8 -p 128,256,512,1024,2048 -n 32,64 -r 3 -ngl 0 + +build/bin/llama-bench -m bitnet-embeddings-270m-i2_s.gguf \ + -t 8 -p 128,256,512,1024,2048 -n 32,64 -r 3 -ngl 0 +``` diff --git a/docs/bitnet-embeddings-qwen3-gguf-conversion.md b/docs/bitnet-embeddings-qwen3-gguf-conversion.md deleted file mode 100644 index 9d63c9300..000000000 --- a/docs/bitnet-embeddings-qwen3-gguf-conversion.md +++ /dev/null @@ -1,302 +0,0 @@ -# BitNet Embeddings (Qwen3) GGUF Conversion Implementation - -## 1. Background - -`bitnet-embeddings-0.6b` is a Qwen3-based embedding model with BitNet per-projection RMSNorm (`BitLinear`). Each linear projection (q/k/v/o/gate/up/down) has a `.norm.weight` that applies RMSNorm to the input **before** the matmul: - -``` -x → RMSNorm(x, norm.weight) → activation_quant(8bit) → matmul(weight_quant(ternary)) -``` - -This pattern does **not** exist in any standard llama.cpp architecture: -- Standard Qwen3: no per-projection norms -- Standard BitNet: has `attn_sub_norm`/`ffn_sub_norm` at different positions (after attention/gate*up, not before each projection) - -### Model Config - -- Architecture: `Qwen3Model` -- hidden_size: 1024, num_attention_heads: 16, num_key_value_heads: 8 -- head_dim: 128 (note: != hidden_size/num_heads = 64) -- intermediate_size: 3072, num_hidden_layers: 28 -- tie_word_embeddings: true -- rope_theta: 1000000, rms_norm_eps: 1e-06 - -### Per-Layer Tensors (7 extra norm tensors per layer) - -| Tensor | Shape | -|--------|-------| -| `self_attn.q_proj.norm.weight` | [1024] | -| `self_attn.k_proj.norm.weight` | [1024] | -| `self_attn.v_proj.norm.weight` | [1024] | -| `self_attn.o_proj.norm.weight` | [2048] | -| `mlp.gate_proj.norm.weight` | [1024] | -| `mlp.up_proj.norm.weight` | [1024] | -| `mlp.down_proj.norm.weight` | [3072] | - ---- - -## 2. GGUF Tensor Name Mapping - -| HF Name | GGUF Name | Notes | -|----------|-----------|-------| -| `embed_tokens.weight` | `token_embd.weight` | | -| `norm.weight` | `output_norm.weight` | | -| `layers.{i}.input_layernorm.weight` | `blk.{i}.attn_norm.weight` | | -| `layers.{i}.post_attention_layernorm.weight` | `blk.{i}.ffn_norm.weight` | | -| `layers.{i}.self_attn.q_proj.weight` | `blk.{i}.attn_q.weight` | | -| `layers.{i}.self_attn.k_proj.weight` | `blk.{i}.attn_k.weight` | | -| `layers.{i}.self_attn.v_proj.weight` | `blk.{i}.attn_v.weight` | | -| `layers.{i}.self_attn.o_proj.weight` | `blk.{i}.attn_output.weight` | | -| `layers.{i}.self_attn.q_norm.weight` | `blk.{i}.attn_q_norm.weight` | QK head norm | -| `layers.{i}.self_attn.k_norm.weight` | `blk.{i}.attn_k_norm.weight` | QK head norm | -| `layers.{i}.self_attn.q_proj.norm.weight` | `blk.{i}.attn_q_norm_in.weight` | BitNet per-projection | -| `layers.{i}.self_attn.k_proj.norm.weight` | `blk.{i}.attn_k_norm_in.weight` | BitNet per-projection | -| `layers.{i}.self_attn.v_proj.norm.weight` | `blk.{i}.attn_v_norm_in.weight` | BitNet per-projection | -| `layers.{i}.self_attn.o_proj.norm.weight` | `blk.{i}.attn_output_norm_in.weight` | BitNet per-projection | -| `layers.{i}.mlp.gate_proj.weight` | `blk.{i}.ffn_gate.weight` | | -| `layers.{i}.mlp.up_proj.weight` | `blk.{i}.ffn_up.weight` | | -| `layers.{i}.mlp.down_proj.weight` | `blk.{i}.ffn_down.weight` | | -| `layers.{i}.mlp.gate_proj.norm.weight` | `blk.{i}.ffn_gate_norm_in.weight` | BitNet per-projection | -| `layers.{i}.mlp.up_proj.norm.weight` | `blk.{i}.ffn_up_norm_in.weight` | BitNet per-projection | -| `layers.{i}.mlp.down_proj.norm.weight` | `blk.{i}.ffn_down_norm_in.weight` | BitNet per-projection | - ---- - -## 3. Conversion Script - -### `utils/convert-bitnet-embedding-to-gguf.py` - -Standalone conversion script (safetensors → GGUF). Key features: - -- Hardcoded HF→GGUF tensor name mapping (no dependency on llama.cpp's Python converter) -- Supports three output types: - - `--outtype f32`: all weights in float32 - - `--outtype f16`: 2D weights and embeddings as float16, norms as float16 - - `--outtype i2_s`: ternary weights packed in I2_S layout, non-ternary weights as float16 -- Writes `key_length` and `value_length` metadata for head_dim=128 (critical: default calculation would give wrong value 64) -- GPT-2 BPE tokenizer handling with pre-tokenizer hash verification -- Pooling type auto-detection from `modules.json` / `1_Pooling/config.json` (sentence-transformers convention) -- EOS token override: uses `<|endoftext|>` (151643) for correct last-token pooling -- Architecture string: `"qwen3"` - -### I2_S Ternary Packing - -The I2_S format packs ternary weights {-1, 0, +1} into 2-bit representation: - -- Quantization: `scale = 1/mean(|w|)`, `q = round(w * scale).clamp(-1, 1)` -- Encoding: `-1 → 0`, `0 → 1`, `+1 → 2` -- Every 128 values form a block, packed into 32 bytes -- Each byte stores 4 values: `byte = (c0 << 6) | (c1 << 4) | (c2 << 2) | c3` -- Scale (float32) is appended at the end of the packed data buffer - -### Tensor Type Assignment - -| Tensor Type | f16 mode | i2_s mode | -|-------------|----------|-----------| -| 2D linear weights | float16 | I2_S ternary packed | -| Embedding weights | float16 | float16 | -| Norm weights (1D) | float16 | float16 | - -Note: `output.weight` (lm_head) is skipped for embedding models — it is not needed (no token generation). - ---- - -## 4. C++ Modifications (`3rdparty/llama.cpp/src/llama.cpp`) - -### 4.1 New Tensor Enums - -Added 7 new entries after `LLM_TENSOR_FFN_SUB_NORM`: - -```cpp -LLM_TENSOR_ATTN_Q_NORM_IN, -LLM_TENSOR_ATTN_K_NORM_IN, -LLM_TENSOR_ATTN_V_NORM_IN, -LLM_TENSOR_ATTN_OUT_NORM_IN, -LLM_TENSOR_FFN_GATE_NORM_IN, -LLM_TENSOR_FFN_UP_NORM_IN, -LLM_TENSOR_FFN_DOWN_NORM_IN, -``` - -### 4.2 Tensor Name Mappings - -Added to `LLM_ARCH_QWEN3` tensor name map: - -```cpp -{ LLM_TENSOR_ATTN_Q_NORM_IN, "blk.%d.attn_q_norm_in" }, -{ LLM_TENSOR_ATTN_K_NORM_IN, "blk.%d.attn_k_norm_in" }, -{ LLM_TENSOR_ATTN_V_NORM_IN, "blk.%d.attn_v_norm_in" }, -{ LLM_TENSOR_ATTN_OUT_NORM_IN, "blk.%d.attn_output_norm_in" }, -{ LLM_TENSOR_FFN_GATE_NORM_IN, "blk.%d.ffn_gate_norm_in" }, -{ LLM_TENSOR_FFN_UP_NORM_IN, "blk.%d.ffn_up_norm_in" }, -{ LLM_TENSOR_FFN_DOWN_NORM_IN, "blk.%d.ffn_down_norm_in" }, -``` - -### 4.3 Layer Struct Fields - -Added to `struct llama_layer`: - -```cpp -struct ggml_tensor * attn_q_norm_in; -struct ggml_tensor * attn_k_norm_in; -struct ggml_tensor * attn_v_norm_in; -struct ggml_tensor * attn_out_norm_in; -struct ggml_tensor * ffn_gate_norm_in; -struct ggml_tensor * ffn_up_norm_in; -struct ggml_tensor * ffn_down_norm_in; -``` - -### 4.4 load_tensors (LLM_ARCH_QWEN3) - -Added optional loading with `TENSOR_NOT_REQUIRED`: - -```cpp -layer.attn_q_norm_in = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM_IN, "weight", i), {n_embd}, TENSOR_NOT_REQUIRED); -layer.attn_k_norm_in = create_tensor(tn(LLM_TENSOR_ATTN_K_NORM_IN, "weight", i), {n_embd}, TENSOR_NOT_REQUIRED); -layer.attn_v_norm_in = create_tensor(tn(LLM_TENSOR_ATTN_V_NORM_IN, "weight", i), {n_embd}, TENSOR_NOT_REQUIRED); -layer.attn_out_norm_in = create_tensor(tn(LLM_TENSOR_ATTN_OUT_NORM_IN, "weight", i), {n_embd_head_k * n_head}, TENSOR_NOT_REQUIRED); -layer.ffn_gate_norm_in = create_tensor(tn(LLM_TENSOR_FFN_GATE_NORM_IN, "weight", i), {n_embd}, TENSOR_NOT_REQUIRED); -layer.ffn_up_norm_in = create_tensor(tn(LLM_TENSOR_FFN_UP_NORM_IN, "weight", i), {n_embd}, TENSOR_NOT_REQUIRED); -layer.ffn_down_norm_in = create_tensor(tn(LLM_TENSOR_FFN_DOWN_NORM_IN, "weight", i), {n_ff}, TENSOR_NOT_REQUIRED); -``` - -Note: `o_proj.norm` input dimension is `n_embd_head_k * n_head` (=2048), `down_proj.norm` input dimension is `n_ff` (=3072). - -### 4.5 build_qwen3() Graph Modifications - -The `build_qwen3()` function was modified to conditionally apply per-projection RMSNorm. The logic is fully backward compatible — when no `*_norm_in` tensors exist, behavior is identical to original. - -**Attention per-projection norms:** -``` -// Before Q/K/V matmul: -if (layer.attn_q_norm_in) { - cur_q = ggml_rms_norm(ctx, cur, hparams.f_norm_rms_eps); - cur_q = ggml_mul(ctx, cur_q, layer.attn_q_norm_in); -} else { - cur_q = cur; -} -Qcur = ggml_mul_mat(ctx, layer.wq, cur_q); -// Similarly for K, V -``` - -**O_proj norm** requires special handling because `llm_build_kv()` normally applies `wo` internally. Solution: pass `wo=NULL` to `llm_build_kv()`, then apply norm + wo manually: - -``` -cur = llm_build_kv(..., wo=NULL, ...); // returns attention output without o_proj -if (layer.attn_out_norm_in) { - cur = ggml_rms_norm(ctx, cur, hparams.f_norm_rms_eps); - cur = ggml_mul(ctx, cur, layer.attn_out_norm_in); -} -cur = ggml_mul_mat(ctx, layer.wo, cur); -``` - -**FFN per-projection norms:** -``` -// Instead of llm_build_ffn(), manually: -if (layer.ffn_gate_norm_in) { - tmp_gate = rms_norm(cur) * gate_norm_in; -} else { - tmp_gate = cur; -} -tmp_gate = matmul(gate_proj, tmp_gate); -// Similarly for up_proj -tmp = silu(tmp_gate) * tmp_up; - -if (layer.ffn_down_norm_in) { - tmp = rms_norm(tmp) * down_norm_in; -} -cur = matmul(down_proj, tmp); -``` - ---- - -## 5. GGUF Conversion Process - -There are two GGUF files to produce, from **two different source models**: - -| GGUF Output | Source Model | Description | -|-------------|-------------|-------------| -| `embeddings-0.6b-f16.gguf` | `multilingual-e5-0.6b` (standard Qwen3) | F16 baseline, standard float16 weights | -| `bitnet-embeddings-0.6b-f16-i2_s.gguf` | `bitnet-embeddings-0.6b` (BitNet ternary) | I2_S ternary packed weights | - -### 5.1 F16 GGUF: from multilingual-e5-0.6b - -The F16 GGUF is converted from the **standard (non-BitNet) model** `multilingual-e5-0.6b`, which has normal float weights and no per-projection RMSNorm. This uses llama.cpp's standard converter since it is a vanilla Qwen3 model: - -```bash -python3 /path/to/llama.cpp/convert_hf_to_gguf.py \ - /path/to/multilingual-e5-0.6b \ - --outtype f16 \ - --outfile embeddings-0.6b-f16.gguf -``` - -**What happens:** -1. Load `model.safetensors` (standard Qwen3 weights, bfloat16) -2. Convert all 2D weights (projections, embeddings) to float16 -3. Convert norm weights to float32 -4. Write GGUF with `qwen3` architecture metadata and tokenizer - -**Output:** ~1.11 GiB (595.78M params) - -### 5.2 I2_S GGUF: from bitnet-embeddings-0.6b - -The I2_S GGUF is converted from the **BitNet ternary model** `bitnet-embeddings-0.6b`, which has ternary weights {-1, 0, +1} and 7 extra per-projection RMSNorm tensors per layer. This uses the custom converter because the standard llama.cpp converter does not handle per-projection norms or I2_S quantization: - -```bash -python3 utils/convert-bitnet-embedding-to-gguf.py \ - /path/to/bitnet-embeddings-0.6b \ - --outfile bitnet-embeddings-0.6b-f16-i2_s.gguf --outtype i2_s -``` - -**What happens:** -1. Load `model.safetensors` (BitNet ternary weights, bfloat16) -2. Map HF tensor names to GGUF names, including 7 extra `*_norm_in` tensors per layer (see Section 2) -3. For each 2D linear weight (q/k/v/o/gate/up/down projections): - - Compute scale: `scale = 1 / mean(|w|)` - - Quantize: `q = round(w * scale).clamp(-1, 1)` - - Encode: `-1 -> 0`, `0 -> 1`, `+1 -> 2` - - Pack every 128 values into 32 bytes (4 values per byte, 2 bits each) - - Append per-row float32 scale -4. Keep embeddings (`token_embd.weight`) in float16 (not ternary) -5. Keep all norm weights in float16 -6. Skip `output.weight` (lm_head, not needed for embedding models) -7. Write GGUF with `I2_S` type tag for quantized tensors - -**Output:** ~699 MiB (~50% of F16 size) - -### 5.3 Why Two Different Source Models? - -- `multilingual-e5-0.6b` is the **teacher/baseline model** with standard float weights, used as the F16 performance reference -- `bitnet-embeddings-0.6b` is the **1-bit quantized student model** with ternary weights and per-projection BitLinear norms, converted to I2_S for efficient CPU inference -- Benchmarking compares both to measure the throughput gain and quality trade-off of ternary quantization - -### 5.4 Tensor Type Summary - -| Tensor | F16 (from e5-0.6b) | I2_S (from bitnet-0.6b) | -|--------|---------------------|-------------------------| -| Linear projections (q/k/v/o/gate/up/down) | float16 | I2_S (2-bit packed + float32 scale) | -| Embedding (`token_embd.weight`) | float16 | float16 | -| Per-projection norms (`*_norm_in`) | N/A (not present) | float16 | -| Layer norms (`attn_norm`, `ffn_norm`) | float32 | float16 | -| QK head norms (`attn_q_norm`, `attn_k_norm`) | float32 | float32 | -| `output.weight` (lm_head) | present | skipped | - ---- - -## 6. Build and Run - -```bash -# Build with BitNet repo (includes I2_S support) -cmake -S /path/to/BitNet -B build -DCMAKE_BUILD_TYPE=Release -cmake --build build --target llama-embedding llama-bench -j$(nproc) - -# Run embedding inference -build/bin/llama-embedding -m bitnet-embeddings-0.6b-f16-i2_s.gguf \ - -p "hello world" --embd-normalize 2 --embd-output-format array - -# Benchmark: F16 vs I2_S -build/bin/llama-bench -m embeddings-0.6b-f16.gguf \ - -t 8 -p 128,256,512,1024,2048 -n 32,64 -r 3 -ngl 0 - -build/bin/llama-bench -m bitnet-embeddings-0.6b-f16-i2_s.gguf \ - -t 8 -p 128,256,512,1024,2048 -n 32,64 -r 3 -ngl 0 -``` diff --git a/utils/convert-bitnet-embedding-to-gguf.py b/utils/convert-bitnet-embedding-to-gguf.py index 3a4340734..3a9e48924 100644 --- a/utils/convert-bitnet-embedding-to-gguf.py +++ b/utils/convert-bitnet-embedding-to-gguf.py @@ -23,11 +23,17 @@ logger = logging.getLogger("convert-bitnet-embedding") +# Supported architectures: model_type -> gguf arch name +SUPPORTED_ARCHS = { + "qwen3": "qwen3", + "gemma3_text": "gemma3", +} + # --------------------------------------------------------------------------- # Tensor name mapping: HuggingFace -> GGUF # --------------------------------------------------------------------------- -def build_tensor_name_map(n_layers: int) -> dict[str, str]: +def build_tensor_name_map(n_layers: int, arch: str) -> dict[str, str]: """Build HF tensor name -> GGUF tensor name mapping.""" mapping: dict[str, str] = { "embed_tokens.weight": "token_embd.weight", @@ -41,7 +47,6 @@ def build_tensor_name_map(n_layers: int) -> dict[str, str]: mapping.update({ # Layer norms f"{pfx}.input_layernorm.weight": f"{blk}.attn_norm.weight", - f"{pfx}.post_attention_layernorm.weight": f"{blk}.ffn_norm.weight", # Self-attention projections f"{pfx}.self_attn.q_proj.weight": f"{blk}.attn_q.weight", @@ -49,7 +54,7 @@ def build_tensor_name_map(n_layers: int) -> dict[str, str]: f"{pfx}.self_attn.v_proj.weight": f"{blk}.attn_v.weight", f"{pfx}.self_attn.o_proj.weight": f"{blk}.attn_output.weight", - # QK head norms (standard Qwen3) + # QK head norms f"{pfx}.self_attn.q_norm.weight": f"{blk}.attn_q_norm.weight", f"{pfx}.self_attn.k_norm.weight": f"{blk}.attn_k_norm.weight", @@ -70,20 +75,29 @@ def build_tensor_name_map(n_layers: int) -> dict[str, str]: f"{pfx}.mlp.down_proj.norm.weight": f"{blk}.ffn_down_norm_in.weight", }) + if arch == "qwen3": + mapping[f"{pfx}.post_attention_layernorm.weight"] = f"{blk}.ffn_norm.weight" + elif arch == "gemma3_text": + mapping.update({ + f"{pfx}.post_attention_layernorm.weight": f"{blk}.post_attention_norm.weight", + f"{pfx}.pre_feedforward_layernorm.weight": f"{blk}.ffn_norm.weight", + f"{pfx}.post_feedforward_layernorm.weight": f"{blk}.post_ffw_norm.weight", + }) + return mapping # --------------------------------------------------------------------------- -# Tokenizer handling (GPT-2 / BPE for Qwen3) +# Tokenizer handling # --------------------------------------------------------------------------- -def get_vocab_base_pre(tokenizer) -> str: +def get_vocab_base_pre(tokenizer, arch: str) -> str: # encoding this string and hashing the resulting tokens would (hopefully) give us a unique identifier that # is specific for the BPE pre-tokenizer used by the model # we will use this unique identifier to write a "tokenizer.ggml.pre" entry in the GGUF file which we can # use in llama.cpp to implement the same pre-tokenizer - chktxt = '\n \n\n \n\n\n \t \t\t \t\n \n \n \n \n\U0001f680 (normal) \U0001f636‍\U0001f32b️ (multiple emojis concatenated) ✅ \U0001f999\U0001f999 3 33 333 3333 33333 333333 3333333 33333333 3.3 3..3 3...3 កាន់តែពិសេសអាច\U0001f601 ?我想在apple工作1314151天～ ------======= нещо на Български \'\'\'\'\'\'```````""""......!!!!!!?????? I\'ve been \'told he\'s there, \'RE you sure? \'M not sure I\'ll make it, \'D you like some tea? We\'Ve a\'lL' + chktxt = '\n \n\n \n\n\n \t \t\t \t\n \n \n \n \n\U0001f680 (normal) \U0001f636‍\U0001f32b️ (multiple emojis concatenated) ✅ \U0001f999\U0001f999 3 33 333 3333 33333 333333 3333333 33333333 3.3 3..3 3...3 កាន់តែពិសេសអាច\U0001f601 ?我想在apple工作1314151天～ ------======= нещо на Български \'\'\'\'\'\'```````""""""......!!!!!!?????? I\'ve been \'told he\'s there, \'RE you sure? \'M not sure I\'ll make it, \'D you like some tea? We\'Ve a\'lL' chktok = tokenizer.encode(chktxt) chkhsh = sha256(str(chktok).encode()).hexdigest() @@ -93,27 +107,38 @@ def get_vocab_base_pre(tokenizer) -> str: res = None - # NOTE: if you get an error here, you need to update the convert_hf_to_gguf_update.py script - # or pull the latest version of the model from Huggingface - # don't edit the hashes manually! - if chkhsh == "0ef9807a4087ebef797fc749390439009c3b9eda9ad1a097abbe738f486c01e5": - # ref: https://huggingface.co/meta-llama/Meta-Llama-3-8B - res = "llama-bpe" - if chkhsh == "049ecf7629871e3041641907f3de7c733e4dbfdc736f57d882ba0b0845599754": - # ref: https://huggingface.co/deepseek-ai/deepseek-llm-7b-base - res = "deepseek-llm" - if chkhsh == "347715f544604f9118bb75ed199f68779f423cabb20db6de6f31b908d04d7821": - # ref: https://huggingface.co/deepseek-ai/deepseek-coder-6.7b-base - res = "deepseek-coder" - if chkhsh == "8aeee3860c56296a157a1fe2fad249ec40aa59b1bb5709f4ade11c4e6fe652ed": - # ref: https://huggingface.co/tiiuae/falcon-7b - res = "falcon" - if chkhsh == "3ce83efda5659b07b1ad37ca97ca5797ea4285d9b9ab0dc679e4a720c9da7454": - # ref: https://huggingface.co/openai-community/gpt2 - res = "gpt-2" - if chkhsh == "d4540891389ea895b53b399da6ac824becc30f2fba0e9ddbb98f92e55ca0e97c": - # ref: https://huggingface.co/Qwen/Qwen3-Embedding-0.6B - res = "qwen2" + if arch == "qwen3": + # NOTE: if you get an error here, you need to update the convert_hf_to_gguf_update.py script + # or pull the latest version of the model from Huggingface + # don't edit the hashes manually! + if chkhsh == "0ef9807a4087ebef797fc749390439009c3b9eda9ad1a097abbe738f486c01e5": + # ref: https://huggingface.co/meta-llama/Meta-Llama-3-8B + res = "llama-bpe" + if chkhsh == "049ecf7629871e3041641907f3de7c733e4dbfdc736f57d882ba0b0845599754": + # ref: https://huggingface.co/deepseek-ai/deepseek-llm-7b-base + res = "deepseek-llm" + if chkhsh == "347715f544604f9118bb75ed199f68779f423cabb20db6de6f31b908d04d7821": + # ref: https://huggingface.co/deepseek-ai/deepseek-coder-6.7b-base + res = "deepseek-coder" + if chkhsh == "8aeee3860c56296a157a1fe2fad249ec40aa59b1bb5709f4ade11c4e6fe652ed": + # ref: https://huggingface.co/tiiuae/falcon-7b + res = "falcon" + if chkhsh == "3ce83efda5659b07b1ad37ca97ca5797ea4285d9b9ab0dc679e4a720c9da7454": + # ref: https://huggingface.co/openai-community/gpt2 + res = "gpt-2" + if chkhsh == "d4540891389ea895b53b399da6ac824becc30f2fba0e9ddbb98f92e55ca0e97c": + # ref: https://huggingface.co/Qwen/Qwen3-Embedding-0.6B + res = "qwen2" + if chkhsh == "855d9fb74bb0b28ce2305e9cd037ff6d8c798f18d19381ddfc14bea3dc9c002f": + # ref: multilingual-e5-0.6b-260311 (Qwen3 tokenizer variant) + res = "qwen2" + elif arch == "gemma3_text": + if chkhsh == "fcb6bf9f20f6c40fa4aa4f7f99607bd6c106ca2348efdacacdca8152e59dcfe9": + # ref: multilingual-e5-270m-260311 (Gemma3 tokenizer) + res = "default" + if chkhsh == "a8594e3edff7c29c003940395316294b2c623571571fc8d3d2d6571f5571cbe6": + # ref: google/gemma-2-9b + res = "default" if res is None: logger.warning("\n") @@ -146,13 +171,29 @@ def _does_token_look_special(token: str) -> bool: return False -def set_vocab(gguf_writer: gguf.GGUFWriter, dir_model: Path, hparams: dict): - """Set GPT-2 BPE vocab for Qwen3.""" +def set_vocab(gguf_writer: gguf.GGUFWriter, dir_model: Path, hparams: dict, arch: str): + """Set tokenizer vocab. + + - Qwen3: BPE tokenizer (tokenizer.ggml.model = "gpt2") + - Gemma3: SPM-compatible tokenizer from tokenizer.json (tokenizer.ggml.model = "llama") + Gemma uses SentencePiece-style tokenization with ▁ space prefix and byte fallback. + Using "llama" model type ensures llama.cpp uses the correct SPM pre-tokenizer + instead of the BPE regex-based pre-tokenizer which breaks CJK tokenization. + """ from transformers import AutoTokenizer tokenizer = AutoTokenizer.from_pretrained(dir_model) vocab_size = hparams.get("vocab_size", len(tokenizer.vocab)) - tokpre = get_vocab_base_pre(tokenizer) + if arch == "gemma3_text": + _set_vocab_gemma3(gguf_writer, dir_model, tokenizer, vocab_size) + else: + _set_vocab_bpe(gguf_writer, dir_model, tokenizer, vocab_size, arch) + + +def _set_vocab_bpe(gguf_writer: gguf.GGUFWriter, dir_model: Path, + tokenizer, vocab_size: int, arch: str): + """Set BPE vocab (for Qwen3).""" + tokpre = get_vocab_base_pre(tokenizer, arch) tokens: list[str] = [] toktypes: list[int] = [] @@ -176,7 +217,6 @@ def set_vocab(gguf_writer: gguf.GGUFWriter, dir_model: Path, hparams: dict): if added_tokens_decoder[i].special or _does_token_look_special(token): toktypes.append(gguf.TokenType.CONTROL) else: - # Pre-normalize user-defined spaces (for Gemma-style tokenizers) token = token.replace(b"\xe2\x96\x81".decode("utf-8"), " ") toktypes.append(gguf.TokenType.USER_DEFINED) @@ -191,14 +231,105 @@ def set_vocab(gguf_writer: gguf.GGUFWriter, dir_model: Path, hparams: dict): gguf_writer.add_token_types(toktypes) special_vocab = gguf.SpecialVocab(dir_model, load_merges=True) - # Override EOS token: PyTorch tokenizer appends <|endoftext|> (151643) as the - # sentence-end marker, not <|im_end|> (151645). For last-token pooling to work - # correctly, llama.cpp must append the same token. - special_vocab.special_token_ids["eos"] = 151643 + + if arch == "qwen3": + # Override EOS token: PyTorch tokenizer appends <|endoftext|> (151643) as the + # sentence-end marker, not <|im_end|> (151645). For last-token pooling to work + # correctly, llama.cpp must append the same token. + special_vocab.special_token_ids["eos"] = 151643 + special_vocab.add_to_gguf(gguf_writer) - # Embedding models need EOS token appended for last-token pooling - gguf_writer.add_add_eos_token(True) + if arch == "qwen3": + # Embedding models need EOS token appended for last-token pooling + gguf_writer.add_add_eos_token(True) + + +def _set_vocab_gemma3(gguf_writer: gguf.GGUFWriter, dir_model: Path, + tokenizer, vocab_size: int): + """Set SPM-compatible vocab for Gemma3. + + Gemma's tokenizer is SentencePiece-based (BPE variant with ▁ space prefix + and byte fallback). We read tokenizer.json to extract vocab and compute + BPE merge scores, then write as tokenizer.ggml.model = "llama" so llama.cpp + uses the SPM code path (correct pre-tokenizer behavior for CJK etc.). + + Score assignment: + - BPE merge results get scores derived from merge rank (lower rank = higher score) + - Single-char / byte tokens get score 0 + - Special / added tokens get score -1000 + """ + tokenizer_json_file = dir_model / "tokenizer.json" + if not tokenizer_json_file.exists(): + raise FileNotFoundError(f"tokenizer.json not found in {dir_model}") + + with open(tokenizer_json_file, encoding="utf-8") as f: + tokenizer_json = json.load(f) + + bpe_vocab = tokenizer_json["model"]["vocab"] # token_str -> token_id + bpe_merges = tokenizer_json["model"].get("merges", []) + + # Build merge result -> rank mapping for score computation + # merge_scores[result_token] = -rank (lower rank = earlier merge = higher priority) + merge_scores: dict[str, float] = {} + for rank, merge in enumerate(bpe_merges): + if isinstance(merge, list): + result = "".join(merge) + else: + parts = merge.split(" ", 1) + result = "".join(parts) + if result not in merge_scores: + merge_scores[result] = -float(rank) + + # Build token arrays + reverse_vocab = {v: k for k, v in bpe_vocab.items()} + added_tokens_decoder = tokenizer.added_tokens_decoder + + tokens: list[bytes] = [] + scores: list[float] = [] + toktypes: list[int] = [] + + for i in range(vocab_size): + if i not in reverse_vocab: + tokens.append(f"[PAD{i}]".encode("utf-8")) + scores.append(-10000.0) + toktypes.append(gguf.TokenType.UNUSED) + continue + + token_str = reverse_vocab[i] + token_bytes = token_str.encode("utf-8") + + # Determine token type + if i in added_tokens_decoder: + tok_data = added_tokens_decoder[i] + if tok_data.special or _does_token_look_special(token_str): + toktypes.append(gguf.TokenType.CONTROL) + else: + toktypes.append(gguf.TokenType.USER_DEFINED) + scores.append(-1000.0) + elif token_str.startswith("<0x") and token_str.endswith(">") and len(token_str) == 6: + # Byte token: <0xHH> + toktypes.append(gguf.TokenType.BYTE) + scores.append(0.0) + elif token_str == "": + toktypes.append(gguf.TokenType.UNKNOWN) + scores.append(0.0) + else: + toktypes.append(gguf.TokenType.NORMAL) + # Score from merge rank, or 0 for single-char tokens + scores.append(merge_scores.get(token_str, 0.0)) + + tokens.append(token_bytes) + + gguf_writer.add_tokenizer_model("llama") + gguf_writer.add_tokenizer_pre("default") + gguf_writer.add_token_list(tokens) + gguf_writer.add_token_scores(scores) + gguf_writer.add_token_types(toktypes) + gguf_writer.add_add_space_prefix(False) + + special_vocab = gguf.SpecialVocab(dir_model, load_merges=False) + special_vocab.add_to_gguf(gguf_writer) # --------------------------------------------------------------------------- @@ -260,7 +391,7 @@ def set_gguf_parameters(gguf_writer: gguf.GGUFWriter, hparams: dict, dir_model: pooling_type = gguf.PoolingType.MEAN gguf_writer.add_pooling_type(pooling_type) - logger.info(f" n_layers={n_layers}, n_embd={n_embd}, n_head={n_head}, n_head_kv={n_head_kv}, n_ff={n_ff}") + logger.info(f" n_layers={n_layers}, n_embd={n_embd}, n_head={n_head}, n_head_kv={n_head_kv}, n_ff={n_ff}, head_dim={head_dim}") # --------------------------------------------------------------------------- @@ -366,7 +497,7 @@ def quantize_to_i2_s(w: np.ndarray) -> np.ndarray: # --------------------------------------------------------------------------- def main(): - parser = argparse.ArgumentParser(description="Convert bitnet-embeddings to GGUF") + parser = argparse.ArgumentParser(description="Convert bitnet-embeddings (Qwen3/Gemma3) to GGUF") parser.add_argument("model", type=Path, help="Model directory") parser.add_argument("--outfile", type=Path, default=None, help="Output GGUF file") parser.add_argument("--outtype", choices=["f32", "f16", "i2_s"], default="f16", @@ -390,9 +521,12 @@ def main(): with open(dir_model / "config.json") as f: hparams = json.load(f) - arch = hparams.get("model_type", "qwen3") - assert arch == "qwen3", f"Expected qwen3 architecture, got {arch}" + arch = hparams.get("model_type", "") + if arch not in SUPPORTED_ARCHS: + logger.error(f"Unsupported model_type '{arch}'. Supported: {list(SUPPORTED_ARCHS.keys())}") + sys.exit(1) + gguf_arch = SUPPORTED_ARCHS[arch] n_layers = hparams["num_hidden_layers"] # Determine ftype @@ -403,20 +537,20 @@ def main(): else: # i2_s ftype = 40 # LLAMA_FTYPE_MOSTLY_I2_S - logger.info(f"Converting {dir_model.name} to GGUF ({args.outtype})") + logger.info(f"Converting {dir_model.name} (arch={arch}) to GGUF ({args.outtype})") # Create GGUF writer - gguf_writer = gguf.GGUFWriter(str(args.outfile), "qwen3") + gguf_writer = gguf.GGUFWriter(str(args.outfile), gguf_arch) # Set parameters set_gguf_parameters(gguf_writer, hparams, dir_model, ftype) # Set vocab logger.info("Setting tokenizer/vocab...") - set_vocab(gguf_writer, dir_model, hparams) + set_vocab(gguf_writer, dir_model, hparams, arch) # Build tensor name map - tensor_map = build_tensor_name_map(n_layers) + tensor_map = build_tensor_name_map(n_layers, arch) # Process tensors logger.info("Processing tensors...") @@ -473,6 +607,16 @@ def main(): else: # norms, 1D tensors + # Gemma3 RMSNorm uses (1+w)*x instead of w*x; preprocess w -> w+1 + # so llama.cpp's standard RMSNorm produces correct results. + # NOTE: *_norm_in weights are BitLinear standard RMSNorm (initialized ~1.0), + # NOT Gemma3RMSNorm (initialized ~0.0), so they must NOT get +1. + is_gemma3_native_norm = (arch == "gemma3_text" and is_norm + and not gguf_name.endswith("_norm_in.weight")) + if is_gemma3_native_norm: + data = data.astype(np.float32) + 1.0 + logger.info(f" [Gemma3 norm offset] {gguf_name}: applied w = w + 1") + if args.outtype in ("f16", "i2_s"): data = data.astype(np.float16) logger.info(f" {gguf_name}: {list(data_torch.shape)} {old_dtype} -> float16") From 1f6a25aeebb6b67764992a4b499f024a22d6351f Mon Sep 17 00:00:00 2001 From: isHuangXin Date: Wed, 15 Jul 2026 08:47:54 +0200 Subject: [PATCH 4/8] Update llama.cpp submodule to include SafetensorRemote utility classes --- 3rdparty/llama.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/3rdparty/llama.cpp b/3rdparty/llama.cpp index 355c0c4d1..36579e6bb 160000 --- a/3rdparty/llama.cpp +++ b/3rdparty/llama.cpp @@ -1 +1 @@ -Subproject commit 355c0c4d1870f0679334cd73727d1e05c1863fbd +Subproject commit 36579e6bb06280a5c9ab381453c4605270201d70 From 425fbbcb755612d34f601714d2c24dd2dbf2eaa6 Mon Sep 17 00:00:00 2001 From: isHuangXin Date: Wed, 15 Jul 2026 09:10:53 +0200 Subject: [PATCH 5/8] Update submodule to release-bitnet-embedding-0.6b-270m branch --- .gitmodules | 4 ++-- 3rdparty/llama.cpp | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/.gitmodules b/.gitmodules index 2b36e4928..b7fbdb264 100644 --- a/.gitmodules +++ b/.gitmodules @@ -1,4 +1,4 @@ [submodule "3rdparty/llama.cpp"] path = 3rdparty/llama.cpp - url = https://github.com/Eddie-Wang1120/llama.cpp.git - branch = merge-dev + url = https://github.com/isHuangXin/llama.cpp.git + branch = release-bitnet-embedding-0.6b-270m diff --git a/3rdparty/llama.cpp b/3rdparty/llama.cpp index 36579e6bb..390c30775 160000 --- a/3rdparty/llama.cpp +++ b/3rdparty/llama.cpp @@ -1 +1 @@ -Subproject commit 36579e6bb06280a5c9ab381453c4605270201d70 +Subproject commit 390c307752ab78fd8189f359d6954c9ba1be74af From 403efc7b00731e64b17227c3fde6e79cadb24c61 Mon Sep 17 00:00:00 2001 From: isHuangXin Date: Wed, 15 Jul 2026 11:56:44 +0200 Subject: [PATCH 6/8] Fix inference pipeline: remove hardcoded batch size, enable tool builds, update Python requirement --- README.md | 4 ++-- run_inference.py | 1 - setup_env.py | 2 +- 3 files changed, 3 insertions(+), 4 deletions(-) diff --git a/README.md b/README.md index 3bb25596e..5e231abbf 100644 --- a/README.md +++ b/README.md @@ -159,7 +159,7 @@ This project is based on the [llama.cpp](https://github.com/ggerganov/llama.cpp) ## Installation ### Requirements -- python>=3.9 +- python>=3.10 - cmake>=3.22 - clang>=18 - For Windows users, install [Visual Studio 2022](https://visualstudio.microsoft.com/downloads/). In the installer, toggle on at least the following options(this also automatically installs the required additional tools like CMake): @@ -186,7 +186,7 @@ cd BitNet 2. Install the dependencies ```bash # (Recommended) Create a new conda environment -conda create -n bitnet-cpp python=3.9 +conda create -n bitnet-cpp python=3.10 conda activate bitnet-cpp pip install -r requirements.txt diff --git a/run_inference.py b/run_inference.py index f3ab727b6..db5034398 100644 --- a/run_inference.py +++ b/run_inference.py @@ -30,7 +30,6 @@ def run_inference(): '-ngl', '0', '-c', str(args.ctx_size), '--temp', str(args.temperature), - "-b", "1", ] if args.conversation: command.append("-cnv") diff --git a/setup_env.py b/setup_env.py index 3bf5fb8f7..50322ac14 100644 --- a/setup_env.py +++ b/setup_env.py @@ -211,7 +211,7 @@ def compile(): logging.error(f"Arch {arch} is not supported yet") exit(0) logging.info("Compiling the code using CMake.") - run_command(["cmake", "-B", "build", *COMPILER_EXTRA_ARGS[arch], *OS_EXTRA_ARGS.get(platform.system(), []), "-DCMAKE_C_COMPILER=clang", "-DCMAKE_CXX_COMPILER=clang++"], log_step="generate_build_files") + run_command(["cmake", "-B", "build", *COMPILER_EXTRA_ARGS[arch], *OS_EXTRA_ARGS.get(platform.system(), []), "-DCMAKE_C_COMPILER=clang", "-DCMAKE_CXX_COMPILER=clang++", "-DLLAMA_BUILD_TOOLS=ON", "-DLLAMA_BUILD_EXAMPLES=ON"], log_step="generate_build_files") # run_command(["cmake", "--build", "build", "--target", "llama-cli", "--config", "Release"]) run_command(["cmake", "--build", "build", "--config", "Release"], log_step="compile") From 85a507c67dd2e361f94152c1e8babff2cb4e7daf Mon Sep 17 00:00:00 2001 From: isHuangXin Date: Thu, 16 Jul 2026 05:32:29 +0200 Subject: [PATCH 7/8] Fix vocab encoding by disabling add_special_tokens in tokenizer.encode --- utils/convert-hf-to-gguf-bitnet.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/utils/convert-hf-to-gguf-bitnet.py b/utils/convert-hf-to-gguf-bitnet.py index e5b02d29e..b11e831b9 100644 --- a/utils/convert-hf-to-gguf-bitnet.py +++ b/utils/convert-hf-to-gguf-bitnet.py @@ -278,7 +278,7 @@ def get_vocab_base(self) -> tuple[list[str], list[int], str]: elif reverse_vocab[i] in added_vocab: # We need to manually encode and decode the added tokens in case special characters # used for `\n` / `\t` have been manually added in the added tokens - encoded_decoded_token = tokenizer.decode(tokenizer.encode(reverse_vocab[i])) + encoded_decoded_token = tokenizer.decode(tokenizer.encode(reverse_vocab[i], add_special_tokens=False)) tokens.append(encoded_decoded_token) if tokenizer.added_tokens_decoder[i].special: toktypes.append(gguf.TokenType.CONTROL) From 10dca87a5d01f74edd12659772c6a2557166dc06 Mon Sep 17 00:00:00 2001 From: isHuangXin Date: Thu, 16 Jul 2026 09:14:04 +0200 Subject: [PATCH 8/8] Enable LLAMA_BUILD_COMMON and LLAMA_BUILD_SERVER to build llama-cli --- setup_env.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/setup_env.py b/setup_env.py index 50322ac14..1b245544f 100644 --- a/setup_env.py +++ b/setup_env.py @@ -211,7 +211,7 @@ def compile(): logging.error(f"Arch {arch} is not supported yet") exit(0) logging.info("Compiling the code using CMake.") - run_command(["cmake", "-B", "build", *COMPILER_EXTRA_ARGS[arch], *OS_EXTRA_ARGS.get(platform.system(), []), "-DCMAKE_C_COMPILER=clang", "-DCMAKE_CXX_COMPILER=clang++", "-DLLAMA_BUILD_TOOLS=ON", "-DLLAMA_BUILD_EXAMPLES=ON"], log_step="generate_build_files") + run_command(["cmake", "-B", "build", *COMPILER_EXTRA_ARGS[arch], *OS_EXTRA_ARGS.get(platform.system(), []), "-DCMAKE_C_COMPILER=clang", "-DCMAKE_CXX_COMPILER=clang++", "-DLLAMA_BUILD_TOOLS=ON", "-DLLAMA_BUILD_EXAMPLES=ON", "-DLLAMA_BUILD_COMMON=ON", "-DLLAMA_BUILD_SERVER=ON"], log_step="generate_build_files") # run_command(["cmake", "--build", "build", "--target", "llama-cli", "--config", "Release"]) run_command(["cmake", "--build", "build", "--config", "Release"], log_step="compile")