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 1f86f058d..390c30775 160000 --- a/3rdparty/llama.cpp +++ b/3rdparty/llama.cpp @@ -1 +1 @@ -Subproject commit 1f86f058de0c3f4098dedae2ae8653c335c868a1 +Subproject commit 390c307752ab78fd8189f359d6954c9ba1be74af 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/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/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/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..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++"], 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") diff --git a/src/ggml-bitnet-lut.cpp b/src/ggml-bitnet-lut.cpp index 59422d548..676351ddc 100644 --- a/src/ggml-bitnet-lut.cpp +++ b/src/ggml-bitnet-lut.cpp @@ -5,23 +5,26 @@ #include #include +#ifdef __x86_64__ +#include +#endif + #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" +#endif #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]; } @@ -29,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; @@ -58,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; } @@ -72,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; @@ -93,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]; } @@ -113,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; } @@ -133,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; @@ -144,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; @@ -164,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/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..3a9e48924 --- /dev/null +++ b/utils/convert-bitnet-embedding-to-gguf.py @@ -0,0 +1,646 @@ +#!/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") + +# 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, arch: str) -> 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", + + # 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 + 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", + }) + + 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 +# --------------------------------------------------------------------------- + +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' + + chktok = tokenizer.encode(chktxt) + chkhsh = sha256(str(chktok).encode()).hexdigest() + + logger.debug(f"chktok: {chktok}") + logger.debug(f"chkhsh: {chkhsh}") + + res = None + + 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") + 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, 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)) + + 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] = [] + + 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: + 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) + + 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) + + 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) + + +# --------------------------------------------------------------------------- +# 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}, head_dim={head_dim}") + + +# --------------------------------------------------------------------------- +# 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 (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", + 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", "") + 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 + 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} (arch={arch}) to GGUF ({args.outtype})") + + # Create GGUF writer + 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, arch) + + # Build tensor name map + tensor_map = build_tensor_name_map(n_layers, arch) + + # 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 + # 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") + 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() diff --git a/utils/convert-hf-to-gguf-bitnet.py b/utils/convert-hf-to-gguf-bitnet.py index 23e84384c..b11e831b9 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"))) @@ -274,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) @@ -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, }