diff --git a/.github/workflows/unsloth-sd-prebuilt.yml b/.github/workflows/unsloth-sd-prebuilt.yml index 4997219464..146ac6dc2e 100644 --- a/.github/workflows/unsloth-sd-prebuilt.yml +++ b/.github/workflows/unsloth-sd-prebuilt.yml @@ -389,6 +389,16 @@ jobs: if: ${{ matrix.cpu_floor != '' }} run: bash tooling/scripts/unsloth/cpu_variants_gate.sh check "dist/sd-${{ needs.resolve.outputs.tag }}-bin-${{ matrix.label }}.zip" "${{ matrix.cpu_floor }}" + # The runner has the newest ISA the bundle uses, so it cannot see a variant library that crashes + # older CPUs while ggml dlopens every variant to score it. Replay startup on emulated CPUs: Intel SDE + # Nehalem .. Sapphire Rapids on x86_64, qemu Cortex-A53 .. max on aarch64. + - name: CPU dispatch gate (emulated CPUs) + if: ${{ matrix.cpu_floor != '' && runner.os == 'Linux' }} + run: | + set -euo pipefail + if [ "$(uname -m)" = "aarch64" ]; then sudo apt-get install -y -qq qemu-user; fi + bash tooling/scripts/unsloth/cpu_dispatch_gate.sh "dist/sd-${{ needs.resolve.outputs.tag }}-bin-${{ matrix.label }}.zip" + - name: Upload bundle uses: actions/upload-artifact@ea165f8d65b6e75b540449e92b4886f43607fa02 # v4.6.2 with: diff --git a/CMakeLists.txt b/CMakeLists.txt index d7bfda5999..d8694e253b 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -328,6 +328,13 @@ if (NOT SD_USE_SYSTEM_GGML) add_definitions(-DGGML_MAX_NAME=160) endif() +# x86 and aarch64 CPU: build llamafile's tinyBLAS GEMM (as llama.cpp does). ggml-patches/0017 picks the weight types +# it takes per CPU from measurements (Q8_0 on x86; F16 / BF16 where tinyBLAS was faster, see ggml-cpu.c); +# GGML_CPU_SGEMM_TYPES overrides the set. +if (NOT DEFINED GGML_LLAMAFILE AND CMAKE_SYSTEM_PROCESSOR MATCHES "^(x86_64|AMD64|amd64|x64|aarch64|arm64|ARM64)$") + set(GGML_LLAMAFILE_DEFAULT ON) +endif() + # deps include(cmake/ggml.cmake) diff --git a/README.md b/README.md index b4ab2f5ab2..1337c5b92f 100644 --- a/README.md +++ b/README.md @@ -48,6 +48,7 @@ API and command-line option may change frequently.*** - [FLUX.2-dev/FLUX.2-klein](./docs/flux2.md) - [FLUX.3 Image Klein / Klein Base](./docs/flux3.md) - [MOSS-Image-Base / MOSS-Image-SFT](./docs/moss_image.md) + - [Fibo 2.0 Dev](./docs/fibo2.md) - [Lens](./docs/lens.md) - [Chroma](./docs/chroma.md) - [Chroma1-Radiance](./docs/chroma_radiance.md) @@ -79,6 +80,7 @@ API and command-line option may change frequently.*** - [Boogu Image Edit](./docs/boogu_image.md) - [Mage-Flow-Edit](./docs/mage_flow.md#image-editing) - [LLaDA-Image Edit](./docs/llada_image.md#image-editing) + - [Fibo 2.0 Dev / Edit Turbo](./docs/fibo2.md#editing-and-turbo) - Video Models - [Wan2.1/Wan2.2](./docs/wan.md) - [MiniMax-H3](./docs/minimax_h3.md) diff --git a/assets/fibo2/dev_edit.png b/assets/fibo2/dev_edit.png new file mode 100644 index 0000000000..fd6e06be13 Binary files /dev/null and b/assets/fibo2/dev_edit.png differ diff --git a/assets/fibo2/t2i.png b/assets/fibo2/t2i.png new file mode 100644 index 0000000000..688ffb36ad Binary files /dev/null and b/assets/fibo2/t2i.png differ diff --git a/docs/adetailer.md b/docs/adetailer.md index 2bc64df2c0..f42831823c 100644 --- a/docs/adetailer.md +++ b/docs/adetailer.md @@ -1,9 +1,10 @@ # ADetailer `sd-cli` can run a YOLOv8 object detector on an existing or newly generated -image and perform a cropped inpaint pass for every detected object. The first -implementation supports YOLOv8 detection checkpoints. YOLOv8 segmentation and -MediaPipe models are not supported yet. +image, and `sd-server` can run it after image generation. Both perform a cropped +inpaint pass for every detected object. The implementation supports YOLOv8 +detection checkpoints. YOLOv8 segmentation and MediaPipe models are not +supported yet. ## Convert a detector @@ -68,6 +69,36 @@ ADetailer can also run automatically after normal image generation: --extra-ad-args "confidence=0.3,denoising_strength=0.4,inpaint_width=512,inpaint_height=512" ``` +The same options can be supplied when starting `sd-server`. When `--ad-model` is set, its detector runs after OpenAI-compatible and native asynchronous image generation. SDAPI requests enable ADetailer through `alwayson_scripts`, as described below. If no object is detected, the original image is returned unchanged. RGBA output keeps its original alpha channel, including when detected regions are repaired. Detector initialization or detail-pass failures are reported as request failures. + +Server image-generation requests limit base, hires, and detail-pass dimensions to 4096 pixels per side. ADetailer step overrides are limited to 1–100; `steps=0` in `extra_ad_args` keeps inheriting the main generation steps. These limits also apply to inherited dimensions and `extra_ad_args` overrides. Requests exceeding the limits return a parameter error before generation. CLI limits are unchanged. + +### SDAPI requests + +`/sdapi/v1/txt2img` and `/sdapi/v1/img2img` accept the WebUI ADetailer argument format: + +```json +{ + "prompt": "portrait photo", + "alwayson_scripts": { + "ADetailer": { + "args": [true, false, { + "ad_model": "face_yolov8n.pt", + "ad_prompt": "[PROMPT], detailed face", + "ad_use_steps": true, + "ad_steps": 20 + }] + } + } +} +``` + +The detector must be configured at server startup with `--ad-model`. `ad_model` must match that file's basename; the original `.pt` basename is also accepted for a converted detector. Requests cannot load arbitrary detector paths. Omitting the script, setting its first argument to `false`, selecting `"None"`, or setting `ad_tab_enable` to `false` disables detailing. Dictionary-only argument arrays are also accepted. Multiple enabled dictionaries run in order using the configured detector. + +Supported settings include prompts, confidence, mask filtering by area or confidence, mask count and area ratios, dilation/erosion, offsets, mask merging/inversion, blur, denoising strength, masked or full-image inpainting, padding, and optional overrides for inpaint dimensions, steps, CFG, sampler/scheduler, and CLIP skip. The corresponding `ad_use_*` flags control whether overrides apply. Otherwise dimensions, steps, CFG, sampler, scheduler, and CLIP skip inherit the main request. Denoising strength defaults to `0.4`. Pixel mask settings are limited to the larger output dimension, with a minimum limit of 32 and a maximum of 16384 pixels. + +Model/VAE switching, non-default noise multipliers, face restoration, ControlNet, and skipping the initial img2img pass are not supported; active requests for these features return an error. Unknown settings and invalid active values also return errors. During detailing, `/sdapi/v1/progress` remains active with an `adetailer` job; its percentage and preview stay fixed until detailing finishes. + An empty ADetailer prompt inherits the main prompt. `[PROMPT]` inserts the main prompt, `[SEP]` assigns different prompts to consecutive masks, and `[SKIP]` skips the corresponding mask. @@ -82,6 +113,7 @@ All settings other than the detector path and prompts are passed through | `nms` | `0.45` | NMS IoU threshold | | `max_detections` | `100` | Maximum detections retained after NMS | | `mask_k_largest` | `0` | Keep only the largest K masks; zero keeps all | +| `mask_filter_method` | `area` | Rank retained masks by `area` or `confidence` | | `mask_min_ratio` | `0` | Minimum bbox area relative to the image | | `mask_max_ratio` | `1` | Maximum bbox area relative to the image | | `dilate_erode` | `4` | Positive values dilate; negative values erode | @@ -90,6 +122,7 @@ All settings other than the detector path and prompts are passed through | `merge_masks`, `invert_mask` | `false` | Boolean alternatives to `mask_mode` | | `mask_blur` | `4` | Final composite feather radius | | `inpaint_padding` | `32` | Padding around the detected region | +| `inpaint_only_masked` | `true` | Crop to the detected region; `false` uses the full image with the detection mask | | `inpaint_width`, `inpaint_height` | mode-specific | `512x512` after generation; input/output size in `adetailer` mode | | `denoising_strength` | mode-specific | `0.4` after generation; inherits `--strength` in `adetailer` mode | | `steps` | `0` | Detail steps; zero inherits the main generation | diff --git a/docs/backend.md b/docs/backend.md index 057d3f20ef..45175135e7 100644 --- a/docs/backend.md +++ b/docs/backend.md @@ -241,9 +241,13 @@ sd-cli -m model.safetensors -p "a cat" --conv-direct off,diffusion=on CUDA 2D requires a compiled Turing-or-newer MMA architecture; HIP 2D requires the supported CDNA MFMA or RDNA WMMA architectures described below for flash attention. MUSA 2D has a SIMT implementation. CUDA, HIP, and MUSA 3D also have - a SIMT implementation, so they do not require MMA. Other backends default to - `off`. Results and decision logs are cached per compute backend and dimension - in the context, as with flash attention. Each dimension is selected independently, + a SIMT implementation, so they do not require MMA. The CPU backend selects + `on` for 2D and 3D when ggml carries the row-run direct kernels + (`scripts/unsloth/ggml-patches/0014`, applied by the prebuilt workflow) and + `off` otherwise: the direct kernels use a fixed patch buffer instead of the + full im2col matrix and match im2col bit for bit with F16 weights. Other + backends default to `off`. Results and decision logs are cached per compute + backend and dimension in the context, as with flash attention. Each dimension is selected independently, and all compute devices in a module must qualify. Parameter-storage devices do not affect this choice. diff --git a/docs/edit.md b/docs/edit.md index 62ecff2af5..9f33862d2a 100644 --- a/docs/edit.md +++ b/docs/edit.md @@ -19,6 +19,7 @@ Depending on the architecture, different models handle reference images differen | **Krea2 (Community Edit LoRAs)** | `krea2_ostris_edit` | | [**Mage-Flow-Edit**](./mage_flow.md#image-editing) | `mage_flow` | | [**LLaDA-Image**](./llada_image.md#image-editing) | `llada_image` | +| [**Fibo 2.0 Dev / Edit Turbo**](./fibo2.md#examples) | `fibo2` | | **Anima (Community Edit LoRAs)** | `cosmos_reference` | Stable-diffusion.spp also supports basic Unet-based editing models like instruct-pix2pix or CosXL-Edit. This document is not about those. @@ -50,6 +51,7 @@ The `--ref-image-args` argument accepts a comma-separated list of key-value pair | `flux_kontext` | FLUX.1 Kontext | | `longcat` | LongCat Image Edit | | `flux2` | FLUX.2 models | +| `fibo2` | Fibo 2.0 Dev / Edit Turbo | | `qwen` | Qwen Image Edit | | `qwen_layered` | Qwen Image Layered | | `z_image_omni` | Boogu, Z-Image Omni | @@ -73,7 +75,12 @@ If presets are insufficient, you can manually configure the following parameters | `ref_index_mode` | string | Behavior of the RoPE index. | `fixed`, `increase`, `decrease` | | `force_ref_timestep_zero` | bool | Forces timestep=0 for reference tokens. | `true`, `false` (Krea2 only) | | `resize_before_vae` | bool | Whether reference images are resized before VAE encoding. | `true`, `false` | -| `vae_input_max_pixels` | int | Maximum pixel area for VAE reference inputs. | Integer | +| `resize_vae_to_target` | bool | Use output dimensions for every reference, overriding area-based sizing. | `true`, `false` | +| `resize_single_ref_to_target` | bool | Use output dimensions when exactly one reference is supplied. | `true`, `false` | +| `vae_size_buckets` | bool | Choose aligned dimensions within ±10% of the requested area, minimizing aspect-ratio cropping. Ignored when using output dimensions. | `true`, `false` | +| `vae_center_crop` | bool | Scale to cover the selected dimensions with Lanczos, then center-crop. Otherwise resize directly. | `true`, `false` | +| `vae_input_max_pixels` | int | VAE area budget, also capped by output area in ordinary resizing. With size buckets, this is the target area and output area does not cap it. | Integer | +| `vlm_use_vae_input` | bool | Reuse the resized/cropped VAE input for the VLM, before the VLM's own resize. Has no effect when `resize_before_vae=false`. | `true`, `false` | | `vlm_resize_mode` | string | How to resize VLM reference inputs. | `longest_side`, `area`, `none` | | `vlm_max_size` | int | Maximum VLM input size; interpreted according to `vlm_resize_mode`. | Integer | | `vlm_min_size` | int | Minimum VLM input size; interpreted according to `vlm_resize_mode`. | Integer | @@ -88,6 +95,7 @@ For a technical overview of how each preset is configured, see the table below. | `flux_kontext` | No | `fixed` | `none` | | | `longcat` | Yes | `fixed` | `area` | | | `flux2` | No | `increase` | `none` | | +| `fibo2` | Yes | `increase` | `longest` | Shared reference crops; clean reference timesteps; see [Fibo 2.0](./fibo2.md) | | `qwen` | Yes | `increase` | `area` | | | `qwen_layered` | Yes | `decrease` | `area` | | | `mage_flow` | Yes | `increase` | `longest` | `vlm_max_size = 384`, VAE input resized to target | @@ -99,3 +107,5 @@ For a technical overview of how each preset is configured, see the table below. **Additional Default Notes:** - **VLM Input Sizes:** For most presets, `vlm_max_size` and `vlm_min_size` are set to `-1`, meaning the values are model-dependent and handled automatically. In `area` mode they represent pixel area; in `longest_side` mode they represent a side length in pixels. - **VAE Input Size:** `vae_input_max_pixels` defaults to $1024 \times 1024$ pixels (`1048576`). +- **VAE Geometry:** `resize_vae_to_target` defaults to `true` for Mage-Flow and LLaDA-Image. Fibo enables `resize_single_ref_to_target`, `vae_size_buckets`, and `vae_center_crop`. These options are otherwise `false`. Bucket dimensions use the model's alignment (normally 16, Qwen Image 32), with a minimum side of four alignment units; if no bucket fits, output dimensions are used. +- **Shared Encoder Input:** `vlm_use_vae_input` defaults to `true` for Fibo and Qwen Image 2.1, otherwise `false`. All these options can be overridden individually after selecting a preset. diff --git a/docs/fibo2.md b/docs/fibo2.md new file mode 100644 index 0000000000..064d800892 --- /dev/null +++ b/docs/fibo2.md @@ -0,0 +1,53 @@ +# How to Use + +Fibo 2.0 uses a Fibo diffusion transformer, the standard FLUX.2 VAE, and Qwen3-VL 4B as the LLM text encoder. The Qwen tokenizer is built in. + +## Download weights + +- Download Fibo 2.0 Dev + - safetensors: https://huggingface.co/Comfy-Org/Fibo-2.0/tree/main/diffusion_models +- Download Fibo 2.0 Edit Turbo + - safetensors: https://huggingface.co/Comfy-Org/Fibo-2.0/tree/main/diffusion_models +- Download vae + - safetensors: https://huggingface.co/black-forest-labs/FLUX.2-dev/tree/main +- Download Qwen3-VL 4B + - safetensors: https://huggingface.co/Comfy-Org/Fibo-2.0/tree/main/text_encoders + - gguf: https://huggingface.co/Qwen/Qwen3-VL-4B-Instruct-GGUF/tree/main + +The original Diffusers weights from [briaai/Fibo-2.0-dev](https://huggingface.co/briaai/Fibo-2.0-dev) and [briaai/Fibo-2.0-edit-turbo](https://huggingface.co/briaai/Fibo-2.0-edit-turbo) are also supported. Pass the transformer's `diffusion_pytorch_model.safetensors.index.json` to `--diffusion-model` and keep its shards alongside it. + +INT8 convrot weights load directly; see [INT8 convrot](./int8_convrot.md) for backend requirements. + +## Examples + +### Fibo 2.0 Dev + +Dev supports both text-to-image generation and reference-image editing. + +``` +.\bin\Release\sd-cli.exe --diffusion-model ..\models\diffusion_models\fibo2_dev_int8_convrot.safetensors --llm ..\models\text_encoders\Qwen3-VL-4B-Instruct-Q4_K_M.gguf --llm_vision ..\models\text_encoders\Qwen3-VL-4B-Instruct-mmproj-BF16.gguf --vae ..\models\vae\flux2_ae.safetensors -p "a lovely cat holding a sign says 'fibo2.cpp'" --cfg-scale 5.0 --sampling-method euler -v --offload-to-cpu +``` + + + +``` +.\bin\Release\sd-cli.exe --diffusion-model ..\models\diffusion_models\fibo2_dev_int8_convrot.safetensors --llm ..\models\text_encoders\Qwen3-VL-4B-Instruct-Q4_K_M.gguf --llm_vision ..\models\text_encoders\Qwen3-VL-4B-Instruct-mmproj-BF16.gguf --vae ..\models\vae\flux2_ae.safetensors -r ..\assets\flux\flux1-dev-q8_0.png -p "change 'flux.cpp' to 'fibo2.cpp'" --cfg-scale 5.0 --sampling-method euler -v --offload-to-cpu +``` + + + +### Fibo 2.0 Edit Turbo + +Edit Turbo is distilled for reference-image editing in 4 steps with CFG 1. + +``` +.\bin\Release\sd-cli.exe --diffusion-model ..\models\diffusion_models\fibo2_edit_turbo_int8_convrot.safetensors --llm ..\models\text_encoders\Qwen3-VL-4B-Instruct-Q4_K_M.gguf --llm_vision ..\models\text_encoders\Qwen3-VL-4B-Instruct-mmproj-BF16.gguf --vae ..\models\vae\flux2_ae.safetensors -r ..\assets\flux\flux1-dev-q8_0.png -p "change 'flux.cpp' to 'fibo2.cpp'" --cfg-scale 1 --steps 4 --sampling-method euler -v --offload-to-cpu +``` + +The BF16 files accept the same options. Use dimensions divisible by 16. Repeat `-r` to use multiple reference images; instructions can refer to ``, ``, etc. See [Image Editing](./edit.md) for reference-image preprocessing options. Separate GGUF text encoders need their vision projector passed with `--llm_vision` for editing. + +The automatic `fibo2` reference preset combines `resize_single_ref_to_target=true`, `vae_size_buckets=true`, `vae_center_crop=true`, and `vlm_use_vae_input=true`. A single reference is scaled to cover the output dimensions and center-cropped. Multiple references use dimensions divisible by 16, chosen within ±10% of one megapixel to minimize cropping for each image's aspect ratio; `vae_input_max_pixels` overrides that target area. The VAE and vision encoder receive the same crops, with smaller copies used by the vision encoder. These are generic reference-image options and can be overridden individually: for example, `--ref-image-args vae_center_crop=false` uses direct resizing, while `resize_before_vae=false` skips this automatic crop and resize step entirely. + +The vision encoder defaults to `vlm_resize_mode=longest_side,vlm_max_size=768`. Override these with `--ref-image-args`, for example `vlm_max_size=512`, or use `vlm_resize_mode=none` to skip size-based scaling. Vision inputs are still aligned to the encoder's patch grid. + +Plain prompts and structured JSON captions are accepted. This integration uses the caption-only text-rendering path (`render_text_images=False`); it does not generate auxiliary glyph images or expand prompts with Bria VGL. Native Fibo mask preprocessing is not connected to `--mask`. diff --git a/docs/performance.md b/docs/performance.md index 18491039eb..2d8440241c 100644 --- a/docs/performance.md +++ b/docs/performance.md @@ -32,8 +32,10 @@ segment count change. path. Performance depends on GPU architecture, backend, driver, model, and tensor shapes, so enabling direct convolution can either improve or reduce speed. -Automatic selection uses small benchmarks on Vulkan and implicit GEMM checks -on other backends. It does not guarantee the fastest end-to-end generation. +Automatic selection uses small benchmarks on Vulkan, implicit GEMM checks on +CUDA/HIP/MUSA, and the direct kernels on CPU when ggml carries the row-run +patch (see [backend selection](backend.md)). It does not guarantee the fastest +end-to-end generation. Compare `on` and `off` with the same model and generation settings over repeated runs, keeping flash attention unchanged and separating startup/probe time from steady-state generation time. See [backend selection](backend.md) diff --git a/docs/rpc.md b/docs/rpc.md index ea9edf1c31..cb28e256f0 100644 --- a/docs/rpc.md +++ b/docs/rpc.md @@ -131,7 +131,7 @@ If you want to allow connections from other machines on the network: > **Drivers & Hardware:** Ensure the Server machine has the necessary drivers installed and functional (e.g., Nvidia Drivers for CUDA, Vulkan SDK, or Metal). If no devices are found, the server will simply fallback to CPU usage. - +``` + +The options can be given in either order; no model is required. The client exits +with an error if RPC support is unavailable or a server cannot provide devices. -### Step B: Run with RPC device +### Step C: Run with RPC device If everything is working correctly, you can now run the client while offloading some or all of the work to the RPC server. @@ -217,4 +217,4 @@ The client will map these servers to sequential device IDs (e.g., RPC0 from the ## 6. Performance Considerations -RPC performance is heavily dependent on network bandwidth, as large weights and activations must be transferred back and forth over the network, especially for large models, or when using high resolutions. For best results, ensure your network connection is stable and has sufficient bandwidth (>1Gbps recommended). This shoumd not be a concern if you are running the server and client on the same machine, as the data transfer will happen over the loopback interface. \ No newline at end of file +RPC performance is heavily dependent on network bandwidth, as large weights and activations must be transferred back and forth over the network, especially for large models, or when using high resolutions. For best results, ensure your network connection is stable and has sufficient bandwidth (>1Gbps recommended). This shoumd not be a concern if you are running the server and client on the same machine, as the data transfer will happen over the loopback interface. diff --git a/docs/taesd.md b/docs/taesd.md index 954086d180..8331f441bf 100644 --- a/docs/taesd.md +++ b/docs/taesd.md @@ -49,3 +49,5 @@ curl -L -o taeqi2_1.safetensors https://huggingface.co/madebyollin/taeqi2_1/reso ``` Replace `--vae PATH` with `--taesd taeqi2_1.safetensors` in the [Qwen Image 2.1 examples](qwen_image_2.1.md) to use it for encoding and decoding. For previews only, keep `--vae PATH` and add `--taesd taeqi2_1.safetensors --taesd-preview-only --preview tae`. + +`sd-server` accepts the same `--taesd` and `--taesd-preview-only` options; the preview method is selected per request instead of via `--preview`. diff --git a/examples/cli/main.cpp b/examples/cli/main.cpp index ca8ce4e79e..ee1e56f220 100644 --- a/examples/cli/main.cpp +++ b/examples/cli/main.cpp @@ -47,7 +47,6 @@ struct SDCliParams { int preview_interval = 1; std::string preview_path = "preview.png"; int preview_fps = 16; - bool taesd_preview = false; bool preview_noisy = false; bool color = false; bool metadata_raw = false; @@ -114,10 +113,6 @@ struct SDCliParams { "--color", "colors the logging tags according to level", true, &color}, - {"", - "--taesd-preview-only", - std::string("prevents usage of taesd for decoding the final image. (for use with --preview ") + previews_str[PREVIEW_TAE] + ")", - true, &taesd_preview}, {"", "--preview-noisy", "enables previewing noisy inputs of the models rather than the denoised outputs", @@ -268,7 +263,6 @@ struct SDCliParams { << " preview_interval: " << preview_interval << ",\n" << " preview_path: \"" << preview_path << "\",\n" << " preview_fps: " << preview_fps << ",\n" - << " taesd_preview: " << (taesd_preview ? "true" : "false") << ",\n" << " preview_noisy: " << (preview_noisy ? "true" : "false") << ",\n" << " imatrix_out: \"" << imatrix_out << "\",\n" << " metadata_raw: " << (metadata_raw ? "true" : "false") << ",\n" @@ -301,6 +295,10 @@ void parse_args(int argc, const char** argv, SDCliParams& cli_params, SDContextP log_level = cli_params.log_level; log_color = cli_params.color; + if (ctx_params.list_devices) { + exit(ctx_params.print_devices() ? 0 : 1); + } + bool valid = cli_params.resolve_and_validate(); if (valid && cli_params.mode != METADATA) { valid = ctx_params.resolve_and_validate(cli_params.mode) && @@ -881,7 +879,7 @@ int main(int argc, const char* argv[]) { } } - sd_ctx_params_t sd_ctx_params = ctx_params.to_sd_ctx_params_t(cli_params.taesd_preview); + sd_ctx_params_t sd_ctx_params = ctx_params.to_sd_ctx_params_t(); SDImageVec results; int num_results = 0; diff --git a/examples/common/common.cpp b/examples/common/common.cpp index b9fb780fdb..dab999a9f4 100644 --- a/examples/common/common.cpp +++ b/examples/common/common.cpp @@ -596,6 +596,11 @@ ArgOptions SDContextParams::get_options() { }; options.bool_options = { + {"", + "--list-devices", + "list available ggml backend devices, including --rpc-servers (one 'namedescription' per line), and exit; " + "the names are the device names accepted by --backend and --params-backend", + true, &list_devices}, {"", "--disable-prefetch", "disable asynchronous next-segment weight prefetch (defaults to false)", @@ -636,6 +641,10 @@ ArgOptions SDContextParams::get_options() { "--sage-attn", "use native CUDA SageAttention in the diffusion model, with flash/default attention fallback", true, &sage_attn}, + {"", + "--taesd-preview-only", + "prevents usage of taesd for decoding the final image; the VAE keeps encoding / decoding and taesd is only used for step previews", + true, &taesd_preview_only}, }; auto on_type_arg = [&](int argc, const char** argv, int index) { @@ -819,23 +828,22 @@ ArgOptions SDContextParams::get_options() { "but it usually offers faster inference speed and, in some cases, lower memory usage. " "The at_runtime mode, on the other hand, is exactly the opposite.", on_lora_apply_mode_arg}, - {"", - "--list-devices", - "list available ggml backend devices (one 'namedescription' per line) and exit; " - "the names are the device names accepted by --backend and --params-backend", - [](int /*argc*/, const char** /*argv*/, int /*index*/) { - size_t device_list_size = sd_list_devices(nullptr, 0); - std::vector devices(device_list_size + 1); - sd_list_devices(devices.data(), devices.size()); - fputs(devices.data(), stdout); - std::exit(0); - return 0; - }}, }; return options; } +bool SDContextParams::print_devices() const { + if (!sd_add_rpc_servers(rpc_servers.c_str())) { + return false; + } + size_t device_list_size = sd_list_devices(nullptr, 0); + std::vector devices(device_list_size + 1); + sd_list_devices(devices.data(), devices.size()); + fputs(devices.data(), stdout); + return true; +} + void SDContextParams::build_embedding_map() { static const std::vector valid_ext = {".gguf", ".safetensors", ".pt", ".ckpt"}; @@ -1010,6 +1018,7 @@ std::string SDContextParams::to_string() const { << " attn_scale: " << attn_scale << ",\n" << " diffusion_conv_direct: " << (diffusion_conv_direct ? "true" : "false") << ",\n" << " vae_conv_direct: " << (vae_conv_direct ? "true" : "false") << ",\n" + << " taesd_preview_only: " << (taesd_preview_only ? "true" : "false") << ",\n" << " prediction: " << sd_prediction_name(prediction) << ",\n" << " lora_apply_mode: " << sd_lora_apply_mode_name(lora_apply_mode) << ",\n" << " force_sdxl_vae_conv_scale: " << (force_sdxl_vae_conv_scale ? "true" : "false") << "\n" @@ -1017,7 +1026,7 @@ std::string SDContextParams::to_string() const { return oss.str(); } -sd_ctx_params_t SDContextParams::to_sd_ctx_params_t(bool taesd_preview) { +sd_ctx_params_t SDContextParams::to_sd_ctx_params_t() { prepare_backend_assignments(); embedding_vec.clear(); embedding_vec.reserve(embedding_map.size()); @@ -1067,7 +1076,7 @@ sd_ctx_params_t SDContextParams::to_sd_ctx_params_t(bool taesd_preview) { sd_ctx_params.sage_attn = sage_attn; sd_ctx_params.linear_scale = linear_scale; sd_ctx_params.attn_scale = attn_scale; - sd_ctx_params.tae_preview_only = taesd_preview; + sd_ctx_params.tae_preview_only = taesd_preview_only; sd_ctx_params.diffusion_conv_direct = diffusion_conv_direct; sd_ctx_params.vae_conv_direct = vae_conv_direct; sd_ctx_params.force_sdxl_vae_conv_scale = force_sdxl_vae_conv_scale; @@ -2219,7 +2228,7 @@ bool SDGenerationParams::from_json_str( if (target_custom_sigmas != nullptr && sample_json.contains("custom_sigmas") && sample_json["custom_sigmas"].is_array() && - !sample_json["custom_sigmas"].empty()) { // empty [] = "not specified", keep server/CLI default + !sample_json["custom_sigmas"].empty()) { // empty [] = "not specified", keep server/CLI default *target_custom_sigmas = sample_json["custom_sigmas"].get>(); } if (sample_json.contains("sample_method") && sample_json["sample_method"].is_string()) { diff --git a/examples/common/common.h b/examples/common/common.h index 91e75c6851..7b9f691af6 100644 --- a/examples/common/common.h +++ b/examples/common/common.h @@ -168,6 +168,7 @@ struct SDContextParams { std::string model_args; bool auto_fit = true; std::string rpc_servers; + bool list_devices = false; std::string effective_backend; std::string effective_params_backend; bool enable_mmap = false; @@ -179,6 +180,7 @@ struct SDContextParams { bool sage_attn = false; bool diffusion_conv_direct = false; bool vae_conv_direct = false; + bool taesd_preview_only = false; prediction_t prediction = PREDICTION_COUNT; lora_apply_mode_t lora_apply_mode = LORA_APPLY_AUTO; @@ -189,13 +191,14 @@ struct SDContextParams { float flow_shift = INFINITY; ArgOptions get_options(); + bool print_devices() const; void build_embedding_map(); void prepare_backend_assignments(); bool resolve(SDMode mode); bool validate(SDMode mode); bool resolve_and_validate(SDMode mode); std::string to_string() const; - sd_ctx_params_t to_sd_ctx_params_t(bool taesd_preview); + sd_ctx_params_t to_sd_ctx_params_t(); }; struct SDGenerationParams { diff --git a/examples/server/api.md b/examples/server/api.md index bb42523661..fc64fb0296 100644 --- a/examples/server/api.md +++ b/examples/server/api.md @@ -280,7 +280,8 @@ TAE slot. | `model` | `checkpoints` | | `diffusion_model`, `high_noise_diffusion_model`, `uncond_diffusion_model` | `diffusion_models` | | `clip_l`, `clip_g`, `t5xxl`, `llm`, `audio_encoder`, `embeddings_connectors` | `text_encoders` | -| `clip_vision`, `llm_vision` | `clip_vision` | +| `clip_vision` | `clip_vision` | +| `llm_vision` | `text_encoders` or `clip_vision` | | `vae`, `audio_vae` | `vae` | | `tae` (legacy alias: `taesd`) | `vae` or `vae_approx` | | `tokenizer` | `tokenizers` | @@ -302,11 +303,15 @@ new model architectures. Optional `context` settings apply to each model load and override the matching startup setting. Supported fields are `vae_format` (`auto`, `flux`, `sd3`, -`flux2`, or `wan`) and `model_args` (a string in the same format as -`--model-args`; replaces the startup string). Other startup settings remain -unchanged. Switching to a recipe that omits these fields restores the startup -values. For example, PiD uses `{"vae_format": "flux"}` and Qwen Image Edit -2511 uses `{"model_args": "qwen_image_zero_cond_t=true"}`. +`flux2`, or `wan`), `model_args` (a string in the same format as +`--model-args`; replaces the startup string), and `backend` / `params_backend` +(strings in the same format as `--backend` / `--params-backend`; replace the +startup assignments for that load, so a recipe can place components on +specific devices, e.g. keep a large text encoder on CPU while another recipe +uses the GPU). Other startup settings remain unchanged. Switching to a +recipe that omits these fields restores the startup values. For example, PiD +uses `{"vae_format": "flux"}` and Qwen Image Edit 2511 uses +`{"model_args": "qwen_image_zero_cond_t=true"}`. Supported defaults: @@ -1032,7 +1037,11 @@ Notes: #### `POST /sdcpp/v1/jobs/{id}/cancel` -Attempts to cancel an accepted job. +Cancels a queued job immediately, or requests interruption of a generating image or +video job. A generating job returns `200 OK` with `cancel_requested: true` and +remains `generating` until execution stops at a cancellation checkpoint. Continue +polling until its status becomes `cancelled`. Cancelled jobs have no result. +Repeated cancellation requests are safe; terminal jobs are returned unchanged. Typical status codes: diff --git a/examples/server/async_jobs.cpp b/examples/server/async_jobs.cpp index 7f3c8a1bfc..3e40e506fa 100644 --- a/examples/server/async_jobs.cpp +++ b/examples/server/async_jobs.cpp @@ -3,13 +3,46 @@ #include "async_jobs.h" #include +#include #include #include +#include #include "common/log.h" #include "common/media_io.h" #include "common/resource_owners.hpp" +class JobCancellationMonitor { + std::mutex mutex_; + std::condition_variable cv_; + bool done_ = false; + std::thread thread_; + +public: + // The caller holds sd_ctx_mutex until this monitor has joined, keeping the + // context alive and preventing cancellation from reaching another request. + JobCancellationMonitor(sd_ctx_t* ctx, const AsyncGenerationJob& job) + : thread_([this, ctx, &job]() { + std::unique_lock lock(mutex_); + while (!done_) { + if (job.cancel_requested.load()) { + // Each generation pass resets the library flag, including ADetailer. + sd_cancel_generation(ctx, SD_CANCEL_ALL); + } + cv_.wait_for(lock, std::chrono::milliseconds(10), [this]() { return done_; }); + } + }) {} + + ~JobCancellationMonitor() { + { + std::lock_guard lock(mutex_); + done_ = true; + } + cv_.notify_one(); + thread_.join(); + } +}; + static void preview_callback(int step, int frame_count, sd_image_t* frames, bool is_noisy, void* data) { (void)is_noisy; if (frame_count <= 0 || frames == nullptr || frames[0].data == nullptr) { @@ -142,6 +175,7 @@ bool cancel_queued_job(AsyncJobManager& manager, AsyncGenerationJob& job) { } manager.queue.erase(new_end, manager.queue.end()); + job.cancel_requested.store(true); job.status = AsyncJobStatus::Cancelled; job.completed_at = unix_timestamp_now_ms(); job.result_images_b64.clear(); @@ -156,13 +190,14 @@ bool cancel_queued_job(AsyncJobManager& manager, AsyncGenerationJob& job) { json make_async_job_json(const AsyncJobManager& manager, const AsyncGenerationJob& job) { json result; - result["id"] = job.id; - result["kind"] = async_job_kind_name(job.kind); - result["status"] = async_job_status_name(job.status); - result["created"] = job.created_at; - result["started"] = job.started_at == 0 ? json(nullptr) : json(job.started_at); - result["completed"] = job.completed_at == 0 ? json(nullptr) : json(job.completed_at); - result["queue_position"] = 0; + result["id"] = job.id; + result["kind"] = async_job_kind_name(job.kind); + result["status"] = async_job_status_name(job.status); + result["cancel_requested"] = job.cancel_requested.load(); + result["created"] = job.created_at; + result["started"] = job.started_at == 0 ? json(nullptr) : json(job.started_at); + result["completed"] = job.completed_at == 0 ? json(nullptr) : json(job.completed_at); + result["queue_position"] = 0; if (job.status == AsyncJobStatus::Queued) { size_t position = 1; @@ -235,6 +270,10 @@ bool execute_img_gen_job(ServerRuntime& runtime, { std::lock_guard lock(*runtime.sd_ctx_mutex); + if (job.cancel_requested.load()) { + return false; + } + JobCancellationMonitor cancellation(runtime.sd_ctx, job); set_preview_callback(job); sd_image_t* raw_results = nullptr; int num_results = 0; @@ -244,10 +283,17 @@ bool execute_img_gen_job(ServerRuntime& runtime, } clear_preview_callback(); results.adopt(raw_results, num_results); + if (job.cancel_requested.load()) { + return false; + } + if (!results.empty() && !apply_server_adetailer(runtime, job.img_gen, params, results)) { + error_message = "ADetailer failed"; + return false; + } } const int num_results = results.count(); - if (num_results <= 0) { + if (job.cancel_requested.load() || num_results <= 0) { error_message = "generate_image returned no results"; return false; } @@ -310,6 +356,10 @@ bool execute_vid_gen_job(ServerRuntime& runtime, { std::lock_guard lock(*runtime.sd_ctx_mutex); + if (job.cancel_requested.load()) { + return false; + } + JobCancellationMonitor cancellation(runtime.sd_ctx, job); set_preview_callback(job); sd_image_t* raw_results = nullptr; if (!generate_video(runtime.sd_ctx, ¶ms, &raw_results, &num_results, &generated_audio, &output_fps)) { @@ -320,7 +370,7 @@ bool execute_vid_gen_job(ServerRuntime& runtime, } num_results = results.count(); - if (num_results <= 0) { + if (job.cancel_requested.load() || num_results <= 0) { free_sd_audio(generated_audio); error_message = "generate_video returned no results"; return false; @@ -406,7 +456,11 @@ void async_job_worker(ServerRuntime& runtime) { } job->completed_at = unix_timestamp_now_ms(); - if (ok) { + if (job->cancel_requested.load()) { + job->status = AsyncJobStatus::Cancelled; + job->error_code = "cancelled"; + job->error_message = "job cancelled by client"; + } else if (ok) { job->status = AsyncJobStatus::Completed; job->result_images_b64 = std::move(output_images); job->result_media_b64 = std::move(output_media_b64); diff --git a/examples/server/async_jobs.h b/examples/server/async_jobs.h index 049a9d73b3..aa54370b1d 100644 --- a/examples/server/async_jobs.h +++ b/examples/server/async_jobs.h @@ -1,5 +1,6 @@ #pragma once +#include #include #include #include @@ -32,9 +33,10 @@ struct AsyncGenerationJob { std::string id; AsyncJobKind kind = AsyncJobKind::ImgGen; AsyncJobStatus status = AsyncJobStatus::Queued; - int64_t created_at = unix_timestamp_now_ms(); - int64_t started_at = 0; - int64_t completed_at = 0; + std::atomic cancel_requested{false}; + int64_t created_at = unix_timestamp_now_ms(); + int64_t started_at = 0; + int64_t completed_at = 0; ImgGenJobRequest img_gen; VidGenJobRequest vid_gen; std::vector result_images_b64; diff --git a/examples/server/frontend b/examples/server/frontend index e36e34a2e2..34e5ad6edc 160000 --- a/examples/server/frontend +++ b/examples/server/frontend @@ -1 +1 @@ -Subproject commit e36e34a2e2204e3df654b16cadc112d8ce8f3174 +Subproject commit 34e5ad6edcaf827b7fe1e9ffa55ee86cd2fc43d2 diff --git a/examples/server/main.cpp b/examples/server/main.cpp index 170877686b..a10cb66f1c 100644 --- a/examples/server/main.cpp +++ b/examples/server/main.cpp @@ -48,6 +48,10 @@ static void parse_args(int argc, log_level = svr_params.log_level; log_color = svr_params.color; + if (ctx_params.list_devices) { + exit(ctx_params.print_devices() ? 0 : 1); + } + const bool random_seed_requested = default_gen_params.seed < 0; if (!svr_params.resolve_and_validate() || @@ -89,7 +93,7 @@ int main(int argc, const char** argv) { LOG_DEBUG("%s", ctx_params.to_string().c_str()); LOG_DEBUG("%s", default_gen_params.to_string().c_str()); - sd_ctx_params_t sd_ctx_params = ctx_params.to_sd_ctx_params_t(false); + sd_ctx_params_t sd_ctx_params = ctx_params.to_sd_ctx_params_t(); ServerModels models; if (!ctx_params.model_path.empty() || !ctx_params.diffusion_model_path.empty()) { const auto started = std::chrono::steady_clock::now(); diff --git a/examples/server/model_catalog.cpp b/examples/server/model_catalog.cpp index a3c59206d2..8efe624b80 100644 --- a/examples/server/model_catalog.cpp +++ b/examples/server/model_catalog.cpp @@ -32,7 +32,7 @@ static const ModelSlot k_model_slots[] = { {"clip_vision", "CLIP Vision", {"clip_vision"}, &SDContextParams::clip_vision_path}, {"t5xxl", "T5", {"text_encoders"}, &SDContextParams::t5xxl_path}, {"llm", "LLM Text Encoder", {"text_encoders"}, &SDContextParams::llm_path}, - {"llm_vision", "LLM Vision Encoder", {"clip_vision"}, &SDContextParams::llm_vision_path}, + {"llm_vision", "LLM Vision Encoder", {"text_encoders", "clip_vision"}, &SDContextParams::llm_vision_path}, {"tokenizer", "Tokenizer", {"tokenizers"}, &SDContextParams::tokenizer}, {"vae", "VAE", {"vae"}, &SDContextParams::vae_path}, {"audio_vae", "Audio VAE", {"vae"}, &SDContextParams::audio_vae_path}, @@ -273,7 +273,7 @@ static void apply_recipe_defaults(const json& defaults, SDGenerationParams& para } static void apply_recipe_context(const json& settings, SDContextParams& params) { - check_keys(settings, {"vae_format", "model_args"}); + check_keys(settings, {"vae_format", "model_args", "backend", "params_backend"}); if (settings.contains("vae_format")) { const auto format = settings.at("vae_format").get(); if (format != "auto" && format != "flux" && format != "sd3" && format != "flux2" && format != "wan") @@ -283,6 +283,12 @@ static void apply_recipe_context(const json& settings, SDContextParams& params) if (settings.contains("model_args")) { params.model_args = settings.at("model_args").get(); } + if (settings.contains("backend")) { + params.backend = settings.at("backend").get(); + } + if (settings.contains("params_backend")) { + params.params_backend = settings.at("params_backend").get(); + } } static const std::map k_recipe_features = { diff --git a/examples/server/recipes/fibo2-dev.json b/examples/server/recipes/fibo2-dev.json new file mode 100644 index 0000000000..aa7e814602 --- /dev/null +++ b/examples/server/recipes/fibo2-dev.json @@ -0,0 +1,39 @@ +{ + "schema_version": 1, + "id": "fibo2-dev", + "name": "Fibo 2.0 Dev", + "description": "Text-to-image generation and native reference-image editing. Select a separate vision projector for editing when it is not bundled in the text encoder.", + "model_url": "https://huggingface.co/briaai/Fibo-2.0-dev", + "slots": [ + { + "id": "diffusion_model", + "required": true + }, + { + "id": "llm", + "required": true, + "label": "Qwen 3 VL 4B Text Encoder" + }, + { + "id": "llm_vision", + "required": false, + "label": "Vision projector (required for separate GGUF text encoders)" + }, + { + "id": "vae", + "required": true, + "label": "FLUX.2 VAE" + } + ], + "defaults": { + "sampling_method": "euler", + "cfg_scale": 5, + "steps": 20, + "width": 1024, + "height": 1024 + }, + "features": [ + "flow_shift", + "ref_images" + ] +} diff --git a/examples/server/recipes/fibo2-edit-turbo.json b/examples/server/recipes/fibo2-edit-turbo.json new file mode 100644 index 0000000000..b84e1cce44 --- /dev/null +++ b/examples/server/recipes/fibo2-edit-turbo.json @@ -0,0 +1,39 @@ +{ + "schema_version": 1, + "id": "fibo2-edit-turbo", + "name": "Fibo 2.0 Edit Turbo", + "description": "Dedicated reference-image editing checkpoint; supply reference images. Select a separate vision projector when it is not bundled in the text encoder.", + "model_url": "https://huggingface.co/briaai/Fibo-2.0-edit-turbo", + "slots": [ + { + "id": "diffusion_model", + "required": true + }, + { + "id": "llm", + "required": true, + "label": "Qwen 3 VL 4B Text Encoder" + }, + { + "id": "llm_vision", + "required": false, + "label": "Vision projector (required for separate GGUF text encoders)" + }, + { + "id": "vae", + "required": true, + "label": "FLUX.2 VAE" + } + ], + "defaults": { + "sampling_method": "euler", + "cfg_scale": 1, + "steps": 4, + "width": 1024, + "height": 1024 + }, + "features": [ + "flow_shift", + "ref_images" + ] +} diff --git a/examples/server/routes_models.cpp b/examples/server/routes_models.cpp index ea14eb821a..9e4e3963c1 100644 --- a/examples/server/routes_models.cpp +++ b/examples/server/routes_models.cpp @@ -94,7 +94,7 @@ static void change_model(ServerRuntime& runtime, const httplib::Request& req, ht *runtime.ctx_params = unload ? models.base_context : std::move(context); *runtime.default_gen_params = unload ? models.base_defaults : std::move(defaults); if (!unload) { - auto params = runtime.ctx_params->to_sd_ctx_params_t(false); + auto params = runtime.ctx_params->to_sd_ctx_params_t(); models.context.reset(new_sd_ctx(¶ms)); if (!models.context) { throw std::runtime_error("Model initialization failed; check the server log for component compatibility and memory errors"); diff --git a/examples/server/routes_openai.cpp b/examples/server/routes_openai.cpp index 29dfe313a7..27a7e18f90 100644 --- a/examples/server/routes_openai.cpp +++ b/examples/server/routes_openai.cpp @@ -75,7 +75,7 @@ static bool build_openai_generation_request(const httplib::Request& req, error_message = "invalid params"; return false; } - return true; + return validate_server_img_gen_request(runtime, request, error_message); } static bool build_openai_edit_request(const httplib::Request& req, @@ -158,22 +158,22 @@ static bool build_openai_edit_request(const httplib::Request& req, request.gen_params.height = height; request.gen_params.batch_count = n; - std::string sd_cpp_extra_args_str = extract_and_remove_sd_cpp_extra_args(request.gen_params.prompt); - const json extra_args = sd_cpp_extra_args_str.empty() - ? json::object() - : json::parse(sd_cpp_extra_args_str, nullptr, false); - const bool has_extra_init_image_override = extra_args.is_object() && extra_args.contains("init_image"); - const bool has_extra_ref_images_override = extra_args.is_object() && extra_args.contains("ref_images"); - const bool has_extra_strength_override = extra_args.is_object() && - extra_args.contains("strength") && - extra_args["strength"].is_number(); - const bool has_extra_mask_image = extra_args.is_object() && - extra_args.contains("mask_image") && - !extra_args["mask_image"].is_null(); + std::string sd_cpp_extra_args_str = extract_and_remove_sd_cpp_extra_args(request.gen_params.prompt); + const json extra_args = sd_cpp_extra_args_str.empty() + ? json::object() + : json::parse(sd_cpp_extra_args_str, nullptr, false); + const bool has_extra_init_image_override = extra_args.is_object() && extra_args.contains("init_image"); + const bool has_extra_ref_images_override = extra_args.is_object() && extra_args.contains("ref_images"); + const bool has_extra_strength_override = extra_args.is_object() && + extra_args.contains("strength") && + extra_args["strength"].is_number(); + const bool has_extra_mask_image = extra_args.is_object() && + extra_args.contains("mask_image") && + !extra_args["mask_image"].is_null(); const bool supports_ref_image_conditioning = sd_ctx_supports_ref_image_conditioning(runtime.sd_ctx); const bool may_need_uploaded_init_image = !has_extra_init_image_override && - (!supports_ref_image_conditioning || !mask_bytes.empty() || - has_extra_mask_image || has_extra_strength_override || has_extra_ref_images_override); + (!supports_ref_image_conditioning || !mask_bytes.empty() || + has_extra_mask_image || has_extra_strength_override || has_extra_ref_images_override); SDImageOwner uploaded_image_for_init; bool uploaded_image_was_first_reference = false; for (auto& bytes : images_bytes) { @@ -252,7 +252,7 @@ static bool build_openai_edit_request(const httplib::Request& req, return false; } - return true; + return validate_server_img_gen_request(runtime, request, error_message); } static bool execute_sync_img_gen_request(ServerRuntime& runtime, @@ -270,6 +270,10 @@ static bool execute_sync_img_gen_request(ServerRuntime& runtime, num_results = 0; } results.adopt(raw_results, num_results); + if (!results.empty() && !apply_server_adetailer(runtime, request, img_gen_params, results)) { + error_message = "ADetailer failed"; + return false; + } } if (results.empty()) { diff --git a/examples/server/routes_sdapi.cpp b/examples/server/routes_sdapi.cpp index 4b8e7a5b1e..014b0b5f89 100644 --- a/examples/server/routes_sdapi.cpp +++ b/examples/server/routes_sdapi.cpp @@ -3,9 +3,12 @@ #include #include #include +#include #include +#include #include #include +#include #include "common/common.h" #include "common/media_io.h" @@ -81,6 +84,202 @@ static enum sample_method_t get_sdapi_sample_method(std::string name) { return it != hardcoded.end() ? it->second : SAMPLE_METHOD_COUNT; } +static bool parse_sdapi_adetailer(const json& body, + ServerRuntime& runtime, + bool img2img, + ImgGenJobRequest& request, + std::string& error_message) { + request.adetailer_enabled = false; + request.adetailer_passes.clear(); + if (!body.contains("alwayson_scripts")) { + return true; + } + try { + const auto& scripts = body.at("alwayson_scripts"); + if (!scripts.is_object()) { + throw std::invalid_argument("alwayson_scripts must be an object"); + } + if (!scripts.contains("ADetailer")) { + return true; + } + const auto& args = scripts.at("ADetailer").at("args"); + if (!args.is_array()) { + throw std::invalid_argument("ADetailer.args must be an array"); + } + if (args.empty() || (args[0].is_boolean() && !args[0].get())) { + return true; + } + static const std::unordered_set keys = { + "ad_model", "ad_model_classes", "ad_tab_enable", "ad_prompt", "ad_negative_prompt", + "ad_confidence", "ad_mask_filter_method", "ad_mask_k", "ad_mask_min_ratio", "ad_mask_max_ratio", + "ad_dilate_erode", "ad_x_offset", "ad_y_offset", "ad_mask_merge_invert", "ad_mask_blur", + "ad_denoising_strength", "ad_inpaint_only_masked", "ad_inpaint_only_masked_padding", + "ad_use_inpaint_width_height", "ad_inpaint_width", "ad_inpaint_height", "ad_use_steps", "ad_steps", + "ad_use_cfg_scale", "ad_cfg_scale", "ad_use_sampler", "ad_sampler", "ad_scheduler", + "ad_use_checkpoint", "ad_checkpoint", "ad_use_vae", "ad_vae", "ad_use_noise_multiplier", + "ad_noise_multiplier", "ad_use_clip_skip", "ad_clip_skip", "ad_restore_face", "ad_controlnet_model", + "ad_controlnet_module", "ad_controlnet_weight", "ad_controlnet_guidance_start", "ad_controlnet_guidance_end", "is_api"}; + for (const auto& config : args) { + if (!config.is_object()) { + continue; + } + auto boolean = [&](const char* key, bool fallback) { + if (!config.contains(key)) + return fallback; + if (!config.at(key).is_boolean()) + throw std::invalid_argument(std::string(key) + " must be boolean"); + return config.at(key).get(); + }; + auto text = [&](const char* key, const char* fallback) { + if (!config.contains(key)) + return std::string(fallback); + if (!config.at(key).is_string()) + throw std::invalid_argument(std::string(key) + " must be a string"); + return config.at(key).get(); + }; + auto number = [&](const char* key, double fallback, double minimum, double maximum) { + double value = fallback; + if (config.contains(key)) { + if (!config.at(key).is_number()) + throw std::invalid_argument(std::string(key) + " must be numeric"); + value = config.at(key).get(); + } + if (!std::isfinite(value) || value < minimum || value > maximum) + throw std::invalid_argument(std::string(key) + " is out of range"); + return value; + }; + auto integer = [&](const char* key, int fallback, int minimum, int maximum = std::numeric_limits::max()) { + if (config.contains(key) && !config.at(key).is_number_integer()) + throw std::invalid_argument(std::string(key) + " must be an integer"); + return static_cast(number(key, fallback, minimum, maximum)); + }; + const std::string model = text("ad_model", "None"); + if (model == "None" || !boolean("ad_tab_enable", true)) { + continue; + } + for (const auto& item : config.items()) { + if (keys.count(item.key()) == 0) + throw std::invalid_argument("unsupported ADetailer setting: " + item.key()); + } + const auto& configured = runtime.default_gen_params->ad_model_path; + if (configured.empty()) + throw std::invalid_argument("ADetailer requires a server-configured --ad-model"); + const auto path = fs::u8path(configured); + if (model.find_first_of("/\\") != std::string::npos || + (lower_ascii(model) != lower_ascii(path.filename().u8string()) && lower_ascii(model) != lower_ascii(path.stem().u8string() + ".pt"))) { + throw std::invalid_argument("ADetailer model must match the server-configured detector"); + } + if (!text("ad_model_classes", "").empty()) + throw std::invalid_argument("ad_model_classes is not supported"); + if (boolean("ad_use_checkpoint", false)) { + const auto checkpoint = config.contains("ad_checkpoint") && config.at("ad_checkpoint").is_null() ? "" : text("ad_checkpoint", "Use same checkpoint"); + if (!checkpoint.empty() && checkpoint != "None" && checkpoint != "Use same checkpoint") + throw std::invalid_argument("ADetailer checkpoint switching is not supported"); + } + if (boolean("ad_use_vae", false)) { + const auto vae = config.contains("ad_vae") && config.at("ad_vae").is_null() ? "" : text("ad_vae", "Use same VAE"); + if (!vae.empty() && vae != "None" && vae != "Use same VAE") + throw std::invalid_argument("ADetailer VAE switching is not supported"); + } + if (boolean("ad_use_noise_multiplier", false) && number("ad_noise_multiplier", 1, 0.5, 1.5) != 1) + throw std::invalid_argument("ADetailer noise multiplier is not supported"); + if (boolean("ad_restore_face", false)) + throw std::invalid_argument("ADetailer face restoration is not supported"); + if (text("ad_controlnet_model", "None") != "None") + throw std::invalid_argument("ADetailer ControlNet is not supported"); + + ServerADetailerPass pass; + pass.prompt = text("ad_prompt", ""); + pass.negative_prompt = text("ad_negative_prompt", ""); + pass.width = request.gen_params.get_resolved_width(); + pass.height = request.gen_params.get_resolved_height(); + pass.strength = static_cast(number("ad_denoising_strength", 0.4, 0, 1)); + auto append = [&](const char* key, const std::string& value) { + if (!pass.extra_args.empty()) + pass.extra_args += ","; + pass.extra_args += std::string(key) + "=" + value; + }; + auto append_number = [&](const char* target, const char* key, double fallback, double minimum, double maximum) { + append(target, json(number(key, fallback, minimum, maximum)).dump()); + }; + auto append_integer = [&](const char* target, const char* key, int fallback, int minimum, int maximum = std::numeric_limits::max()) { + append(target, std::to_string(integer(key, fallback, minimum, maximum))); + }; + append_number("confidence", "ad_confidence", 0.3, 0, 1); + append_integer("mask_k_largest", "ad_mask_k", 0, 0); + const auto filter = text("ad_mask_filter_method", "Area"); + if (filter != "Area" && filter != "Confidence") + throw std::invalid_argument("invalid ad_mask_filter_method"); + append("mask_filter_method", lower_ascii(filter)); + const auto min_ratio = number("ad_mask_min_ratio", 0, 0, 1); + const auto max_ratio = number("ad_mask_max_ratio", 1, 0, 1); + if (min_ratio > max_ratio) + throw std::invalid_argument("ad_mask_min_ratio exceeds ad_mask_max_ratio"); + append("mask_min_ratio", json(min_ratio).dump()); + append("mask_max_ratio", json(max_ratio).dump()); + // Bound pixel operations before converting them to the native mask parameters. + const int mask_limit = std::clamp(std::max(pass.width, pass.height), 32, 16384); + append_integer("dilate_erode", "ad_dilate_erode", 4, -mask_limit, mask_limit); + append_integer("x_offset", "ad_x_offset", 0, -mask_limit, mask_limit); + append_integer("y_offset", "ad_y_offset", 0, -mask_limit, mask_limit); + append_integer("mask_blur", "ad_mask_blur", 4, 0, mask_limit); + append_integer("inpaint_padding", "ad_inpaint_only_masked_padding", 32, 0, mask_limit); + append("inpaint_only_masked", boolean("ad_inpaint_only_masked", true) ? "true" : "false"); + const auto merge = text("ad_mask_merge_invert", "None"); + if (merge != "None" && merge != "Merge" && merge != "Merge and Invert") + throw std::invalid_argument("invalid ad_mask_merge_invert"); + append("mask_mode", merge == "None" ? "none" : merge == "Merge" ? "merge" + : "merge_invert"); + if (boolean("ad_use_inpaint_width_height", false)) { + pass.width = integer("ad_inpaint_width", 512, 1, k_max_generation_dimension); + pass.height = integer("ad_inpaint_height", 512, 1, k_max_generation_dimension); + } + if (boolean("ad_use_steps", false)) + append_integer("steps", "ad_steps", 28, 1, k_max_adetailer_steps); + if (boolean("ad_use_cfg_scale", false)) + append_number("cfg_scale", "ad_cfg_scale", 7, 0, std::numeric_limits::max()); + if (boolean("ad_use_clip_skip", false)) { + pass.clip_skip = integer("ad_clip_skip", 1, 1); + if (pass.clip_skip > 12) + throw std::invalid_argument("ad_clip_skip is out of range"); + } + if (boolean("ad_use_sampler", false)) { + std::string sampler = text("ad_sampler", "DPM++ 2M Karras"); + if (sampler != "Use same sampler") { + const std::string suffix = " karras"; + const bool legacy_karras = sampler.size() > suffix.size() && lower_ascii(sampler.substr(sampler.size() - suffix.size())) == suffix; + if (legacy_karras) + sampler.resize(sampler.size() - suffix.size()); + const auto method = get_sdapi_sample_method(sampler); + if (method == SAMPLE_METHOD_COUNT) + throw std::invalid_argument("unsupported ADetailer sampler"); + append("sample_method", sd_sample_method_name(method)); + if (legacy_karras && text("ad_scheduler", "Use same scheduler") == "Use same scheduler") + append("scheduler", sd_scheduler_name(KARRAS_SCHEDULER)); + } + const auto scheduler = text("ad_scheduler", "Use same scheduler"); + if (scheduler == "Automatic") { + pass.automatic_scheduler = true; + } else if (scheduler != "Use same scheduler") { + const auto value = str_to_scheduler(scheduler.c_str()); + if (value == SCHEDULER_COUNT) + throw std::invalid_argument("unsupported ADetailer scheduler"); + append("scheduler", sd_scheduler_name(value)); + } + } + request.adetailer_passes.push_back(std::move(pass)); + } + if (!request.adetailer_passes.empty() && img2img && args.size() > 1 && args[1].is_boolean() && args[1].get()) { + throw std::invalid_argument("ADetailer skip img2img is not supported"); + } + request.adetailer_enabled = !request.adetailer_passes.empty(); + return true; + } catch (const std::exception& error) { + error_message = std::string("invalid ADetailer args: ") + error.what(); + return false; + } +} + static bool build_sdapi_img_gen_request(const json& j, ServerRuntime& runtime, bool img2img, @@ -259,7 +458,8 @@ static bool build_sdapi_img_gen_request(const json& j, return false; } - return true; + return parse_sdapi_adetailer(j, runtime, img2img, request, error_message) && + validate_server_img_gen_request(runtime, request, error_message); } static nlohmann::json prepare_info_field(const SDContextParams& ctx_params, @@ -349,7 +549,7 @@ static void preview_callback(int step, int frame_count, sd_image_t* frames, bool } struct SDAPIPreviewGuard { - SDAPIPreviewGuard(const std::string& id_task, const sd_img_gen_params_t& params, const std::string& job_type) { + SDAPIPreviewGuard(const std::string& id_task, const sd_img_gen_params_t& params, const std::string& job_type, bool detailing = false) { std::lock_guard lock(g_preview.mutex); g_preview.id_task = id_task; g_preview.total = 0; @@ -358,11 +558,20 @@ struct SDAPIPreviewGuard { g_preview.step = 0; g_preview.active = true; // Each batch item has a base pass and, when enabled, a hires pass. - g_preview.job_count = static_cast(std::max(1, params.batch_count)) * (params.hires.enabled ? 2 : 1); + g_preview.job_count = static_cast(std::max(1, params.batch_count)) * (params.hires.enabled ? 2 : 1) + (detailing ? 1 : 0); g_preview.job_no = 0; sd_set_preview_callback(preview_callback, PREVIEW_PROJ, 1, true, false, nullptr); } + void begin_detailing() { + sd_set_preview_callback(nullptr, PREVIEW_NONE, 1, false, false, nullptr); + std::lock_guard lock(g_preview.mutex); + g_preview.job_type = "adetailer"; + g_preview.job_no = static_cast(g_preview.job_count - 1); + g_preview.step = 0; + g_preview.total = 0; + } + ~SDAPIPreviewGuard() { sd_set_preview_callback(nullptr, PREVIEW_NONE, 1, false, false, nullptr); std::lock_guard lock(g_preview.mutex); @@ -410,13 +619,21 @@ void register_sdapi_endpoints(httplib::Server& svr, ServerRuntime& rt) { { std::lock_guard lock(*runtime->sd_ctx_mutex); - SDAPIPreviewGuard preview_guard(id_task, img_gen_params, img2img ? "img2img" : "txt2img"); + const bool detailing = request.adetailer_enabled && !runtime->default_gen_params->ad_model_path.empty(); + SDAPIPreviewGuard preview_guard(id_task, img_gen_params, img2img ? "img2img" : "txt2img", detailing); sd_image_t* raw_results = nullptr; if (!generate_image(runtime->sd_ctx, &img_gen_params, &raw_results, &num_results)) { raw_results = nullptr; num_results = 0; } results.adopt(raw_results, num_results); + if (detailing && !results.empty()) + preview_guard.begin_detailing(); + if (!results.empty() && !apply_server_adetailer(*runtime, request, img_gen_params, results)) { + res.status = 500; + res.set_content(R"({"error":"ADetailer failed"})", "application/json"); + return; + } } if (results.empty()) { diff --git a/examples/server/routes_sdcpp.cpp b/examples/server/routes_sdcpp.cpp index 7e93e5edbf..7839d5c53b 100644 --- a/examples/server/routes_sdcpp.cpp +++ b/examples/server/routes_sdcpp.cpp @@ -206,7 +206,7 @@ static json make_img_gen_features_json() { {"hires", true}, {"cache", true}, {"cancel_queued", true}, - {"cancel_generating", false}, + {"cancel_generating", true}, {"preview", true}, }; } @@ -221,7 +221,7 @@ static json make_vid_gen_features_json() { {"vae_tiling", true}, {"cache", true}, {"cancel_queued", true}, - {"cancel_generating", false}, + {"cancel_generating", true}, {"preview", true}, }; } @@ -381,7 +381,7 @@ static json make_capabilities_json(ServerRuntime& runtime) { json top_level_output_formats = json::array(); json top_level_features = { {"cancel_queued", true}, - {"cancel_generating", false}, + {"cancel_generating", true}, }; std::string current_mode = ""; if (supports_img) { @@ -410,9 +410,9 @@ static json make_capabilities_json(ServerRuntime& runtime) { result["defaults_by_mode"] = defaults_by_mode; result["limits"] = { {"min_width", 64}, - {"max_width", 4096}, + {"max_width", k_max_generation_dimension}, {"min_height", 64}, - {"max_height", 4096}, + {"max_height", k_max_generation_dimension}, {"max_batch_count", 8}, {"max_queue_size", manager.max_pending_jobs}, {"max_upscale_width", k_max_upscale_dimension}, @@ -457,7 +457,7 @@ static bool parse_img_gen_request(const json& body, error_message = "invalid generation parameters"; return false; } - return true; + return validate_server_img_gen_request(runtime, request, error_message); } static bool parse_vid_gen_request(const json& body, @@ -840,8 +840,9 @@ void register_sdcpp_api_endpoints(httplib::Server& svr, ServerRuntime& rt) { } if (job.status == AsyncJobStatus::Generating) { - res.status = 409; - res.set_content(R"({"error":"job is currently generating and cannot be interrupted yet"})", "application/json"); + job.cancel_requested.store(true); + res.status = 200; + res.set_content(make_async_job_json(manager, job).dump(), "application/json"); return; } diff --git a/examples/server/runtime.cpp b/examples/server/runtime.cpp index 34140ea353..5b33b338a5 100644 --- a/examples/server/runtime.cpp +++ b/examples/server/runtime.cpp @@ -3,6 +3,7 @@ #include #include #include +#include #include #include #include @@ -18,6 +19,157 @@ namespace fs = std::filesystem; +bool validate_server_img_gen_request(const ServerRuntime& runtime, + const ImgGenJobRequest& request, + std::string& error_message) { + auto dimensions_valid = [&](double width, double height) { + if (!std::isfinite(width) || !std::isfinite(height) || + width < 1 || height < 1 || width > k_max_generation_dimension || height > k_max_generation_dimension) { + error_message = "generation width and height must be between 1 and " + std::to_string(k_max_generation_dimension); + return false; + } + return true; + }; + const auto& params = request.gen_params; + const int width = params.get_resolved_width(); + const int height = params.get_resolved_height(); + if (!dimensions_valid(width, height)) { + return false; + } + if (params.hires_enabled) { + const double hires_width = params.hires_width > 0 ? params.hires_width : params.hires_height > 0 ? params.hires_height + : std::round(width * params.hires_scale); + const double hires_height = params.hires_height > 0 ? params.hires_height : params.hires_width > 0 ? params.hires_width + : std::round(height * params.hires_scale); + if (!dimensions_valid(hires_width, hires_height)) { + return false; + } + } + if (!request.adetailer_enabled || runtime.default_gen_params->ad_model_path.empty()) { + return true; + } + auto pass_valid = [&](const ServerADetailerPass& pass) { + if (!dimensions_valid(pass.width, pass.height)) { + return false; + } + // Match the native key normalization, but leave full option parsing to ADetailer. + auto trim = [](const std::string& text) { + const auto first = text.find_first_not_of(" \t\r\n"); + return first == std::string::npos ? std::string() : text.substr(first, text.find_last_not_of(" \t\r\n") - first + 1); + }; + std::stringstream stream(pass.extra_args); + std::string item; + while (std::getline(stream, item, ',')) { + const auto equals = item.find('='); + if (equals == std::string::npos) { + continue; + } + auto key = trim(item.substr(0, equals)); + std::transform(key.begin(), key.end(), key.begin(), [](unsigned char c) { return static_cast(std::tolower(c)); }); + if (key != "inpaint_width" && key != "inpaint_height" && key != "steps") { + continue; + } + const auto value = trim(item.substr(equals + 1)); + const int minimum = key == "steps" ? 0 : 1; + const int maximum = key == "steps" ? k_max_adetailer_steps : k_max_generation_dimension; + try { + size_t used = 0; + const int parsed = std::stoi(value, &used); + if (used == value.size() && parsed >= minimum && parsed <= maximum) { + continue; + } + } catch (const std::exception&) { + } + error_message = "ADetailer " + key + " must be between " + std::to_string(minimum) + " and " + std::to_string(maximum); + return false; + } + return true; + }; + if (request.adetailer_passes.empty()) { + ServerADetailerPass pass; + pass.extra_args = params.extra_ad_args; + return pass_valid(pass); + } + return std::all_of(request.adetailer_passes.begin(), request.adetailer_passes.end(), pass_valid); +} + +bool apply_server_adetailer(ServerRuntime& runtime, + const ImgGenJobRequest& request, + const sd_img_gen_params_t& params, + SDImageVec& results) { + const auto& gen_params = request.gen_params; + const auto& ad_model_path = runtime.default_gen_params->ad_model_path; + if (!request.adetailer_enabled || ad_model_path.empty()) { + return true; + } + std::string error_message; + if (!validate_server_img_gen_request(runtime, request, error_message)) { + LOG_ERROR("%s", error_message.c_str()); + return false; + } + + std::vector passes = request.adetailer_passes; + if (passes.empty()) { + ServerADetailerPass pass; + pass.prompt = gen_params.ad_prompt; + pass.negative_prompt = gen_params.ad_negative_prompt; + pass.extra_args = gen_params.extra_ad_args; + passes.push_back(std::move(pass)); + } + + ADetailerCtxPtr ad_ctx(new_adetailer_ctx(ad_model_path.c_str(), + runtime.ctx_params->n_threads, + runtime.ctx_params->effective_backend.c_str(), + runtime.ctx_params->effective_params_backend.c_str())); + if (ad_ctx == nullptr) { + LOG_ERROR("new_adetailer_ctx failed"); + return false; + } + + for (const auto& pass : passes) { + sd_adetailer_params_t ad_params{}; + ad_params.prompt = pass.prompt.empty() ? nullptr : pass.prompt.c_str(); + ad_params.negative_prompt = pass.negative_prompt.empty() ? nullptr : pass.negative_prompt.c_str(); + ad_params.extra_ad_args = pass.extra_args.c_str(); + for (int i = 0; i < results.count(); ++i) { + if (results[i].data == nullptr) { + continue; + } + sd_img_gen_params_t ad_generation_params = params; + ad_generation_params.seed = params.seed + i; + ad_generation_params.width = pass.width; + ad_generation_params.height = pass.height; + ad_generation_params.strength = pass.strength; + if (pass.clip_skip > 0) { + ad_generation_params.clip_skip = pass.clip_skip; + } + if (pass.automatic_scheduler) { + ad_generation_params.sample_params.scheduler = SCHEDULER_COUNT; + } + + sd_image_t* detailed_images = nullptr; + int detailed_count = 0; + if (!adetail_image(ad_ctx.get(), + runtime.sd_ctx, + results[i], + &ad_params, + &ad_generation_params, + &detailed_images, + &detailed_count) || + detailed_count <= 0 || detailed_images == nullptr || detailed_images[0].data == nullptr) { + free_sd_images(detailed_images, detailed_count); + LOG_ERROR("ADetailer failed for image %d", i + 1); + return false; + } + free(results[i].data); + results[i] = detailed_images[0]; + detailed_images[0] = {0, 0, 0, nullptr}; + free_sd_images(detailed_images, detailed_count); + } + } + return true; +} + static std::string lower_ascii(std::string value) { std::transform(value.begin(), value.end(), value.begin(), [](unsigned char c) { return static_cast(std::tolower(c)); diff --git a/examples/server/runtime.h b/examples/server/runtime.h index d9272f6c24..8986f38270 100644 --- a/examples/server/runtime.h +++ b/examples/server/runtime.h @@ -67,18 +67,38 @@ struct ServerRuntime { ServerModels* models; }; +struct ServerADetailerPass { + std::string prompt; + std::string negative_prompt; + std::string extra_args; + int width = 512; + int height = 512; + float strength = 0.4f; + int clip_skip = -1; + bool automatic_scheduler = false; +}; + struct ImgGenJobRequest { SDGenerationParams gen_params; std::string output_format = "png"; int output_compression = 100; std::string preview_mode = "none"; int preview_interval = 1; + bool adetailer_enabled = true; + std::vector adetailer_passes; sd_img_gen_params_t to_sd_img_gen_params_t() { return gen_params.to_sd_img_gen_params_t(); } }; +constexpr int k_max_generation_dimension = 4096; +constexpr int k_max_adetailer_steps = 100; + +bool validate_server_img_gen_request(const ServerRuntime& runtime, + const ImgGenJobRequest& request, + std::string& error_message); + struct VidGenJobRequest { SDGenerationParams gen_params; std::string output_format = "webm"; @@ -111,3 +131,7 @@ void refresh_lora_cache(ServerRuntime& rt); std::string get_lora_full_path(ServerRuntime& rt, const std::string& path); void refresh_upscaler_cache(ServerRuntime& rt); int64_t unix_timestamp_now_ms(); +bool apply_server_adetailer(ServerRuntime& runtime, + const ImgGenJobRequest& request, + const sd_img_gen_params_t& params, + SDImageVec& results); diff --git a/ggml b/ggml index dc998771e7..c1f80e507d 160000 --- a/ggml +++ b/ggml @@ -1 +1 @@ -Subproject commit dc998771e700d49bbdf9a39cede22bb5b3311767 +Subproject commit c1f80e507db1d0ce49f3bf7f81bc79246d7a2611 diff --git a/include/stable-diffusion.h b/include/stable-diffusion.h index 404d8cbddc..85fe51d463 100644 --- a/include/stable-diffusion.h +++ b/include/stable-diffusion.h @@ -640,6 +640,11 @@ SD_API const char* sd_version(void); // terminator. Passing nullptr or buffer_size 0 only queries the required size. SD_API size_t sd_list_devices(char* buffer, size_t buffer_size); +// Register comma-separated RPC endpoints (host:port) before listing devices. +// Repeated endpoints are safe; nullptr or an empty string is a no-op. +// Returns false if RPC is unavailable or an endpoint has no usable devices. +SD_API bool sd_add_rpc_servers(const char* servers); + // for C API, caller needs to call free_sd_images to free the memory after use // This helps avoid CRT problems on Windows when memory is allocated in the library but freed in the caller, which may use a different CRT. SD_API void free_sd_images(sd_image_t* result_images, int num_images); diff --git a/scripts/unsloth/cpu_dispatch_gate.sh b/scripts/unsloth/cpu_dispatch_gate.sh new file mode 100755 index 0000000000..5ebd2ff731 --- /dev/null +++ b/scripts/unsloth/cpu_dispatch_gate.sh @@ -0,0 +1,179 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: MIT +# Copyright 2026-present the Unsloth AI Inc. team. +# +# cpu_dispatch_gate.sh +# +# Runs the bundle's sd-cli on emulated CPUs that are older or newer than the build runner: Intel SDE on +# Linux x86_64 (Nehalem .. Sapphire Rapids), qemu-aarch64 on Linux aarch64 (Cortex-A53 .. max). For each CPU +# it fails when +# - the process dies or SDE reports an instruction the emulated CPU does not have. ggml dlopens EVERY CPU +# variant library while it scores them, so a static initializer that the compiler vectorised with a +# variant's AVX512 / SVE flags crashes every host without that ISA, before any variant is picked; +# - no CPU variant is loaded, or the loaded variant needs a feature the emulated CPU lacks; +# - on a 128-bit SVE vector length (CPU names ending in -sve128) a variant other than armv8.2_2 is loaded: +# ggml-patches/0020 skips the SVE variants there unless GGML_CPU_ARM_SVE128=1 (names ending in +keep +# set it and must then load an SVE variant); +# - a variant that has the iqk GEMM kernels built does not enable them by default (ARM: every batch size; +# x86 AVX2 and up: 2-12 src1 rows, ggml-patches/0018). Variants without them (pre-AVX2 x86, Armv8.0) +# report built types 0x0. +# The runner's own CPU cannot catch either: it has the newest ISA the bundle uses. +# +# Env: SDE_DIR (an unpacked Intel SDE; downloaded and checked against SDE_SHA256 when unset), +# QEMU_AARCH64 (default qemu-aarch64), CPU_DISPATCH_GATE_CPUS (space separated subset of the CPUs below). + +set -euo pipefail + +zip="${1:?usage: cpu_dispatch_gate.sh }" + +SDE_URL="https://downloadmirror.intel.com/859732/sde-external-9.58.0-2025-06-16-lin.tar.xz" +SDE_SHA256="f849acecad4c9b108259c643b2688fd65c35723cd23368abe5dd64b917cc18c0" + +# ggml CPU variant -> features it is built with (ggml/src/CMakeLists.txt ggml_add_cpu_backend_variant). +# AVX512 stands for F + CD + VL + DQ + BW, the set ggml's AVX512 variants require. +variant_needs() { + case "$1" in + x64) echo "" ;; + sse42) echo "SSE42" ;; + sandybridge) echo "SSE42 AVX" ;; + ivybridge) echo "SSE42 AVX F16C" ;; + piledriver) echo "SSE42 AVX F16C FMA" ;; + haswell) echo "SSE42 AVX F16C FMA AVX2 BMI2" ;; + skylakex) echo "SSE42 AVX F16C FMA AVX2 BMI2 AVX512" ;; + cannonlake) echo "SSE42 AVX F16C FMA AVX2 BMI2 AVX512 AVX512_VBMI" ;; + cascadelake) echo "SSE42 AVX F16C FMA AVX2 BMI2 AVX512 AVX512_VNNI" ;; + icelake) echo "SSE42 AVX F16C FMA AVX2 BMI2 AVX512 AVX512_VBMI AVX512_VNNI" ;; + cooperlake) echo "SSE42 AVX F16C FMA AVX2 BMI2 AVX512 AVX512_VNNI AVX512_BF16" ;; + zen4) echo "SSE42 AVX F16C FMA AVX2 BMI2 AVX512 AVX512_VBMI AVX512_VNNI AVX512_BF16" ;; + alderlake) echo "SSE42 AVX F16C FMA AVX2 BMI2 AVX_VNNI" ;; + sapphirerapids) echo "SSE42 AVX F16C FMA AVX2 BMI2 AVX512 AVX512_VBMI AVX512_VNNI AVX512_BF16 AMX_TILE AMX_INT8" ;; + armv8.0_1) echo "" ;; + armv8.2_1) echo "DOTPROD" ;; + armv8.2_2) echo "DOTPROD FP16_VA" ;; + armv8.2_3) echo "DOTPROD FP16_VA SVE" ;; + armv8.6_1) echo "DOTPROD FP16_VA SVE I8MM" ;; + armv8.6_2) echo "DOTPROD FP16_VA SVE I8MM SVE2" ;; + armv9.2_1) echo "DOTPROD FP16_VA SVE I8MM SME" ;; + armv9.2_2) echo "DOTPROD FP16_VA SVE I8MM SVE2 SME" ;; + *) echo "UNKNOWN_VARIANT" ;; + esac +} + +# Emulated CPU -> features it has. x86: SDE chip flags; aarch64: qemu -cpu models. +X86_HSW="SSE42 AVX F16C FMA AVX2 BMI2" +X86_SKX="$X86_HSW AVX512" +cpu_has() { + case "${1%-sve128}" in + nhm|wsm) echo "SSE42" ;; + snb) echo "SSE42 AVX" ;; + ivb) echo "SSE42 AVX F16C" ;; + hsw|bdw|skl) echo "$X86_HSW" ;; + adl) echo "$X86_HSW AVX_VNNI" ;; + skx) echo "$X86_SKX" ;; + clx) echo "$X86_SKX AVX512_VNNI" ;; + icl|icx) echo "$X86_SKX AVX512_VBMI AVX512_VNNI" ;; + cpx) echo "$X86_SKX AVX512_VNNI AVX512_BF16" ;; + spr) echo "$X86_SKX AVX512_VBMI AVX512_VNNI AVX512_BF16 AMX_TILE AMX_INT8 AVX_VNNI" ;; + cortex-a53) echo "" ;; + cortex-a76|neoverse-n1) echo "DOTPROD FP16_VA" ;; + neoverse-v1) echo "DOTPROD FP16_VA SVE I8MM" ;; + max) echo "DOTPROD FP16_VA SVE I8MM SVE2 SME" ;; + *) echo "UNKNOWN_CPU" ;; + esac +} + +work="" +cleanup() { if [ -n "$work" ]; then rm -rf -- "$work"; fi; } +trap cleanup EXIT +work="$(mktemp -d "${RUNNER_TEMP:-/tmp}/cpu-dispatch-XXXXXX")" +unzip -q "$zip" -d "$work/bundle" +cli="$(find "$work/bundle" -type f -name sd-cli | head -n1)" +[ -n "$cli" ] || { echo "ERROR: no sd-cli in $zip" >&2; exit 1; } +mkdir -p "$work/cwd" + +arch="$(uname -m)" +case "$arch" in + x86_64) + cpus="${CPU_DISPATCH_GATE_CPUS:-nhm wsm snb ivb hsw bdw adl skx clx icl icx spr}" + if [ -z "${SDE_DIR:-}" ]; then + curl -fsSL --retry 3 -o "$work/sde.tar.xz" "$SDE_URL" + echo "$SDE_SHA256 $work/sde.tar.xz" | sha256sum --check --strict + mkdir -p "$work/sde" + tar -xJf "$work/sde.tar.xz" -C "$work/sde" --strip-components=1 + SDE_DIR="$work/sde" + fi + emu() { emu_cmd=("$SDE_DIR/sde64" "-$1" --); } + ;; + aarch64) + cpus="${CPU_DISPATCH_GATE_CPUS:-cortex-a53 cortex-a76 neoverse-v1 max neoverse-v1-sve128 max-sve128 max-sve128+keep}" + qemu="${QEMU_AARCH64:-qemu-aarch64}" + command -v "$qemu" > /dev/null || { echo "ERROR: $qemu not installed (the workflow installs qemu-user)" >&2; exit 1; } + # -sve128: the same model with a 16-byte default SVE vector length (Neoverse N2 / V2 class hardware) + emu() { + case "$1" in + *-sve128) emu_cmd=("$qemu" -cpu "${1%-sve128},sve-default-vector-length=16") ;; + *) emu_cmd=("$qemu" -cpu "$1") ;; + esac + } + ;; + *) + echo "cpu_dispatch_gate: nothing to emulate on $arch" + exit 0 + ;; +esac + +fail=0 +printf '%-18s %-16s %-6s %s\n' cpu variant rc note +for entry in $cpus; do + cpu="${entry%+keep}" + keep="" + [ "$cpu" != "$entry" ] && keep=1 + has=" $(cpu_has "$cpu") " + # Nonexistent model: sd-cli loads (and scores) every CPU variant, initialises the chosen one, then stops. + # Empty cwd defeats ggml's cwd search. GGML_CPU_IQK_DEBUG makes the iqk GEMM report what it was built with. + emu "$cpu" + set +e + out="$(cd "$work/cwd" && env -u LD_LIBRARY_PATH GGML_CPU_IQK_DEBUG=1 ${keep:+GGML_CPU_ARM_SVE128=1} \ + "${emu_cmd[@]}" "$cli" -m "$work/cwd/none.gguf" -p x -v 2>&1)" + rc=$? + set -e + loaded="$(printf '%s\n' "$out" | sed -n 's/.*loaded CPU backend from .*libggml-cpu-\(.*\)\.so$/\1/p' | head -n1 | tr -d '\r')" + bad="$(printf '%s\n' "$out" | grep -m1 -E 'SDE-ERROR|Illegal instruction|SIGILL|core dumped' || true)" + where="$(printf '%s\n' "$out" | sed -n 's/^Image: .*\/\([^/]*\)$/\1/p; s/^Function: \(.*\)$/in \1/p' | head -n2 | tr '\n' ' ')" + iqk="$(printf '%s\n' "$out" | sed -n 's/^ggml-cpu iqk gemm: \(built types [^ ]*\), \(enabled [^ ]*\).*/\1 \2/p' | head -n1)" + note="" + # rc 1 is sd-cli's own exit for the missing model; a signal is 128+n, SDE's chip check also exits 1. + if [ -n "$bad" ] || [ "$rc" -gt 1 ]; then + note="CRASH: ${bad:-exit $rc} ${where}" + elif [ -z "$loaded" ]; then + note="no CPU variant loaded" + else + needs=" $(variant_needs "$loaded") " + for f in $needs; do + case "$has" in *" $f "*) ;; *) note="$note needs $f";; esac + done + case "$cpu" in + *-sve128) + if [ -n "$keep" ]; then + case "$needs" in *" SVE "*) ;; *) note="$note expected an SVE variant with GGML_CPU_ARM_SVE128=1";; esac + elif [ "$loaded" != "armv8.2_2" ]; then + note="$note expected armv8.2_2 on 128-bit SVE" + fi + ;; + esac + built="$(printf '%s\n' "$iqk" | sed -n 's/^built types \([^ ,]*\).*/\1/p')" + enabled="$(printf '%s\n' "$iqk" | sed -n 's/.*enabled \([^ ,]*\).*/\1/p')" + if [ -n "$built" ] && [ "$built" != "0x0" ] && [ "$enabled" != "$built" ]; then + note="$note iqk built but not enabled by default" + fi + fi + printf '%-18s %-16s %-6s %s\n' "$entry" "${loaded:-none}" "$rc" "${note:-ok}${iqk:+ (iqk $iqk)}" + if [ -n "$note" ]; then + fail=1 + printf '%s\n' "$out" | tail -n 15 >&2 + fi +done +if [ "$fail" != 0 ]; then + echo "ERROR: a CPU variant crashed or was picked for a CPU that lacks its features (see above)" >&2 +fi +exit "$fail" diff --git a/scripts/unsloth/ggml-patches/0008-ggml-cpu-avx512-variants-need-avx-vnni.patch b/scripts/unsloth/ggml-patches/0008-ggml-cpu-avx512-variants-need-avx-vnni.patch index b3e8b9a812..6304b90ed3 100644 --- a/scripts/unsloth/ggml-patches/0008-ggml-cpu-avx512-variants-need-avx-vnni.patch +++ b/scripts/unsloth/ggml-patches/0008-ggml-cpu-avx512-variants-need-avx-vnni.patch @@ -2,14 +2,15 @@ diff --git a/src/ggml-cpu/arch/x86/cpu-feats.cpp b/src/ggml-cpu/arch/x86/cpu-fea index d775a036..a29335b9 100644 --- a/src/ggml-cpu/arch/x86/cpu-feats.cpp +++ b/src/ggml-cpu/arch/x86/cpu-feats.cpp -@@ -300,6 +300,10 @@ static int ggml_backend_cpu_x86_score() { +@@ -300,6 +300,11 @@ static int ggml_backend_cpu_x86_score() { if (!is.AVX512VL()) { return 0; } if (!is.AVX512DQ()) { return 0; } if (!is.AVX512BW()) { return 0; } -+ // Unsloth: the AVX512 variants are only offered on CPUs that also have AVX-VNNI (Zen 5, -+ // Sapphire Rapids and newer). On Zen 4, which has AVX512 but not AVX-VNNI, they ran -+ // stable-diffusion.cpp 15 to 50% slower than the AVX2 variant; the AVX2 one is picked there. -+ if (!is.AVX_VNNI()) { return 0; } ++ // Unsloth: on AMD the AVX512 variants are only offered with AVX-VNNI (Zen 5). On Zen 4, which ++ // has AVX512 but not AVX-VNNI, they ran stable-diffusion.cpp 15 to 50% slower than the AVX2 ++ // variant, so the AVX2 one is picked there. Intel AVX512 CPUs without AVX-VNNI (Skylake-SP, ++ // Cascade Lake, Ice Lake) keep them: cascadelake ran Z-Image 1.7x faster than haswell. ++ if (is.is_amd && !is.AVX_VNNI()) { return 0; } score += 1<<7; #endif #ifdef GGML_AVX512_VBMI diff --git a/scripts/unsloth/ggml-patches/0014-ggml-cpu-conv-direct-row-runs.patch b/scripts/unsloth/ggml-patches/0014-ggml-cpu-conv-direct-row-runs.patch new file mode 100644 index 0000000000..310e1feecc --- /dev/null +++ b/scripts/unsloth/ggml-patches/0014-ggml-cpu-conv-direct-row-runs.patch @@ -0,0 +1,231 @@ +diff --git a/include/ggml-cpu.h b/include/ggml-cpu.h +--- a/include/ggml-cpu.h ++++ b/include/ggml-cpu.h +@@ -3,6 +3,10 @@ + #include "ggml.h" + #include "ggml-backend.h" + ++// The direct conv_2d / conv_3d kernels fill and scatter their patches in row runs (bit-identical to the ++// per-element loops); stable-diffusion.cpp defaults CPU convolution to direct only when this is defined. ++#define GGML_CPU_CONV_DIRECT_ROW_RUNS 1 ++ + #ifdef __cplusplus + extern "C" { + #endif +diff --git a/src/ggml-cpu/ops.cpp b/src/ggml-cpu/ops.cpp +--- a/src/ggml-cpu/ops.cpp ++++ b/src/ggml-cpu/ops.cpp +@@ -7464,6 +7464,8 @@ + // ggml_compute_forward_conv_2d + + ++#define CONV_2D_RUN 64 ++ + static void ggml_compute_forward_conv_2d_impl(const ggml_compute_params * params, + const ggml_tensor * kernel, // [KW, KH, IC, OC] + const ggml_tensor * src, // [W, H, C, N] +@@ -7521,40 +7523,44 @@ + const int64_t patch_start = patch_start_batch + params->ith * patch_per_thread; + const int64_t patch_end = std::min(patch_start + patch_per_thread, patch_end_batch); + +- //im2col for a patch +- for (int64_t p = patch_start; p < patch_end; ++p) { +- const int64_t batch_n = p / (dst_w * dst_h); +- const int64_t src_x = (p / dst_w) % dst_h; +- const int64_t src_y = p % dst_w; ++ // im2col for this thread's patches, a run of up to CONV_2D_RUN consecutive patches of one output row at a ++ // time: each (ic, ky) source row is then read once, contiguously, for the whole run instead of once per patch ++ // (one patch touches c_in planes; one plane per patch was a cache and TLB miss per element). Same values, ++ // same conversions as the per-patch loop, so the result is bit-identical. ++ const size_t ts = traits->type_size; ++ for (int64_t p = patch_start; p < patch_end;) { ++ const int64_t n_idx = p / (dst_w * dst_h); ++ const int64_t oy = (p / dst_w) % dst_h; ++ const int64_t ox0 = p % dst_w; ++ const int64_t run = std::min({ patch_end - p, dst_w - ox0, (int64_t) CONV_2D_RUN }); + +- const float * src_base = (const float *)((const char *)src_data + batch_n * src->nb[3]); +- char * dst_row = (char *) tmp + (p % patches_per_batch) * knl_n * traits->type_size; ++ const char * src_base = (const char *) src_data + n_idx * src->nb[3]; ++ char * dst_base = (char *) tmp + (p % patches_per_batch) * knl_n * ts; + + for (int64_t ic = 0; ic < c_in; ++ic) { + for (int64_t ky = 0; ky < knl_h; ++ky) { ++ const int64_t sy = oy * stride_y + ky * dilation_y - pad_y; ++ const bool row = sy >= 0 && sy < src_h; ++ const char * src_row = row ? src_base + sy * src->nb[1] + ic * src->nb[2] : nullptr; + for (int64_t kx = 0; kx < knl_w; ++kx) { +- const int64_t sy = src_x * stride_y + ky * dilation_y - pad_y; +- const int64_t sx = src_y * stride_x + kx * dilation_x - pad_x; +- +- int64_t dst_idx = ic * (knl_h * knl_w) + ky * knl_w + kx; +- +- float src_val; +- if (sy < 0 || sy >= src_h || sx < 0 || sx >= src_w) { +- src_val = 0.0f; +- } else { +- const float * src_ptr = (const float *)((const char *)src_base + sx * src->nb[0] + sy * src->nb[1] + ic * src->nb[2]); +- src_val = *src_ptr; +- } +- +- char * element_ptr = dst_row + dst_idx * traits->type_size; +- if (kernel_type == GGML_TYPE_F32) { +- *(float *) element_ptr = src_val; +- } else if (kernel_type == GGML_TYPE_F16) { +- *(ggml_fp16_t *) element_ptr = GGML_CPU_FP32_TO_FP16(src_val); ++ char * dst_col = dst_base + (ic * (knl_h * knl_w) + ky * knl_w + kx) * ts; ++ for (int64_t j = 0; j < run; ++j) { ++ const int64_t sx = (ox0 + j) * stride_x + kx * dilation_x - pad_x; ++ float src_val = 0.0f; ++ if (row && sx >= 0 && sx < src_w) { ++ src_val = *(const float *) (src_row + sx * src->nb[0]); ++ } ++ char * element_ptr = dst_col + j * knl_n * ts; ++ if (kernel_type == GGML_TYPE_F32) { ++ *(float *) element_ptr = src_val; ++ } else { ++ *(ggml_fp16_t *) element_ptr = GGML_CPU_FP32_TO_FP16(src_val); ++ } + } + } + } + } ++ p += run; + } // patches handled by this thread + + ggml_barrier(params->threadpool); +@@ -7574,17 +7580,23 @@ + const int64_t permute_start = params->ith * permute_per_thread; + const int64_t permute_end = std::min(permute_start + permute_per_thread, patch_n); + +- for (int64_t i = permute_start; i < permute_end; ++i) { +- const int64_t p = patch_start_batch + i; +- const int64_t batch_n = p / (dst_w * dst_h); +- const int64_t dst_y = (p / dst_w) % dst_h; +- const int64_t dst_x = p % dst_w; ++ // Same run blocking for the transpose back: one output plane at a time per run, so each write is a contiguous ++ // row segment rather than one element in each of c_out planes. ++ for (int64_t i = permute_start; i < permute_end;) { ++ const int64_t p = patch_start_batch + i; ++ const int64_t n_idx = p / (dst_w * dst_h); ++ const int64_t oy = (p / dst_w) % dst_h; ++ const int64_t ox0 = p % dst_w; ++ const int64_t run = std::min({ permute_end - i, dst_w - ox0, (int64_t) CONV_2D_RUN }); + + for (int64_t oc = 0; oc < c_out; ++oc) { +- const float value = gemm_output[i * c_out + oc]; +- float * dst_ptr = (float *)((char *)dst_data + dst_x * dst->nb[0] + dst_y * dst->nb[1] + oc * dst->nb[2] + batch_n * dst->nb[3]); +- *dst_ptr = value; ++ char * dst_row = (char *) dst_data + ox0 * dst->nb[0] + oy * dst->nb[1] + oc * dst->nb[2] + n_idx * dst->nb[3]; ++ const float * col = gemm_output + i * c_out + oc; ++ for (int64_t j = 0; j < run; ++j) { ++ *(float *) (dst_row + j * dst->nb[0]) = col[j * c_out]; ++ } + } ++ i += run; + } + } + } +@@ -7662,45 +7674,48 @@ + const int64_t patch_start = patch_start_batch + params->ith * patch_per_thread; + const int64_t patch_end = std::min(patch_start + patch_per_thread, patch_end_batch); + +- for (int64_t p = patch_start; p < patch_end; ++p) { ++ // Runs of consecutive patches along one output row, as in conv_2d: each (ic, kz, ky) source row is read once ++ // per run. Same values and conversions as the per-patch loop (bit-identical). ++ const size_t ts = traits->type_size; ++ for (int64_t p = patch_start; p < patch_end;) { + const int64_t p_in_batch = p % (dst_w * dst_h * dst_d); + const int64_t p_in_depth = p_in_batch % (dst_w * dst_h); + const int64_t batch_idx = p / (dst_w * dst_h * dst_d); + const int64_t dst_z = p_in_batch / (dst_w * dst_h); + const int64_t dst_y = p_in_depth / dst_w; +- const int64_t dst_x = p_in_depth % dst_w; ++ const int64_t dst_x0 = p_in_depth % dst_w; ++ const int64_t run = std::min({ patch_end - p, dst_w - dst_x0, (int64_t) CONV_2D_RUN }); + +- char * dst_row = (char *) tmp + (p % patches_per_batch) * knl_n_total * traits->type_size; ++ char * dst_base = (char *) tmp + (p % patches_per_batch) * knl_n_total * ts; + + for (int64_t ic = 0; ic < c; ++ic) { ++ const int64_t cn_idx = batch_idx * c + ic; + for (int64_t kz = 0; kz < knl_d; ++kz) { ++ const int64_t sz = dst_z * s2 + kz * d2 - p2; + for (int64_t ky = 0; ky < knl_h; ++ky) { ++ const int64_t sy = dst_y * s1 + ky * d1 - p1; ++ const bool row = sz >= 0 && sz < src_d && sy >= 0 && sy < src_h; ++ const char * src_row = row ? (const char *) src_data + sy * src->nb[1] + sz * src->nb[2] + cn_idx * src->nb[3] : nullptr; + for (int64_t kx = 0; kx < knl_w; ++kx) { +- const int64_t sz = dst_z * s2 + kz * d2 - p2; +- const int64_t sy = dst_y * s1 + ky * d1 - p1; +- const int64_t sx = dst_x * s0 + kx * d0 - p0; +- +- int64_t dst_idx = ic * knl_n_per_channel + kz * (knl_h * knl_w) + ky * knl_w + kx; +- +- float src_val; +- if (sz < 0 || sz >= src_d || sy < 0 || sy >= src_h || sx < 0 || sx >= src_w) { +- src_val = 0.0f; +- } else { +- const int64_t cn_idx = batch_idx * c + ic; +- const float * src_ptr = (const float *)((const char *)src_data + sx*src->nb[0] + sy*src->nb[1] + sz*src->nb[2] + cn_idx*src->nb[3]); +- src_val = *src_ptr; +- } +- +- char * element_ptr = dst_row + dst_idx * traits->type_size; +- if (kernel_type == GGML_TYPE_F32) { +- *(float *)element_ptr = src_val; +- } else if (kernel_type == GGML_TYPE_F16) { +- *(ggml_fp16_t *)element_ptr = GGML_CPU_FP32_TO_FP16(src_val); ++ char * dst_col = dst_base + (ic * knl_n_per_channel + kz * (knl_h * knl_w) + ky * knl_w + kx) * ts; ++ for (int64_t j = 0; j < run; ++j) { ++ const int64_t sx = (dst_x0 + j) * s0 + kx * d0 - p0; ++ float src_val = 0.0f; ++ if (row && sx >= 0 && sx < src_w) { ++ src_val = *(const float *) (src_row + sx * src->nb[0]); ++ } ++ char * element_ptr = dst_col + j * knl_n_total * ts; ++ if (kernel_type == GGML_TYPE_F32) { ++ *(float *) element_ptr = src_val; ++ } else { ++ *(ggml_fp16_t *) element_ptr = GGML_CPU_FP32_TO_FP16(src_val); ++ } + } + } + } + } + } ++ p += run; + } + + ggml_barrier(params->threadpool); +@@ -7714,21 +7729,25 @@ + const int64_t permute_start = params->ith * permute_per_thread; + const int64_t permute_end = std::min(permute_start + permute_per_thread, patch_n_in_batch); + +- for (int64_t i = permute_start; i < permute_end; ++i) { ++ for (int64_t i = permute_start; i < permute_end;) { + const int64_t p = patch_start_batch + i; + const int64_t p_in_batch = p % (dst_w * dst_h * dst_d); + const int64_t p_in_depth = p_in_batch % (dst_w * dst_h); + const int64_t batch_idx = p / (dst_w * dst_h * dst_d); + const int64_t dst_z = p_in_batch / (dst_w * dst_h); + const int64_t dst_y = p_in_depth / dst_w; +- const int64_t dst_x = p_in_depth % dst_w; ++ const int64_t dst_x0 = p_in_depth % dst_w; ++ const int64_t run = std::min({ permute_end - i, dst_w - dst_x0, (int64_t) CONV_2D_RUN }); + + for (int64_t ioc = 0; ioc < oc; ++ioc) { +- const float value = gemm_output[i * oc + ioc]; + const int64_t ocn_idx = batch_idx * oc + ioc; +- float * dst_ptr = (float *)((char *)dst_data + dst_x*dst->nb[0] + dst_y*dst->nb[1] + dst_z*dst->nb[2] + ocn_idx*dst->nb[3]); +- *dst_ptr = value; ++ char * dst_row = (char *) dst_data + dst_x0 * dst->nb[0] + dst_y * dst->nb[1] + dst_z * dst->nb[2] + ocn_idx * dst->nb[3]; ++ const float * col = gemm_output + i * oc + ioc; ++ for (int64_t j = 0; j < run; ++j) { ++ *(float *) (dst_row + j * dst->nb[0]) = col[j * oc]; ++ } + } ++ i += run; + } + } + } diff --git a/scripts/unsloth/ggml-patches/0017-ggml-cpu-llamafile-sgemm-q8_0-only.patch b/scripts/unsloth/ggml-patches/0017-ggml-cpu-llamafile-sgemm-q8_0-only.patch new file mode 100644 index 0000000000..a46b48b33a --- /dev/null +++ b/scripts/unsloth/ggml-patches/0017-ggml-cpu-llamafile-sgemm-q8_0-only.patch @@ -0,0 +1,149 @@ +diff --git a/src/ggml-cpu/ggml-cpu.c b/src/ggml-cpu/ggml-cpu.c +index 30b8cbb..6f9319d 100644 +--- a/src/ggml-cpu/ggml-cpu.c ++++ b/src/ggml-cpu/ggml-cpu.c +@@ -49,6 +49,115 @@ + + #ifdef GGML_USE_LLAMAFILE + #include "llamafile/sgemm.h" ++#if defined(__AVX2__) && !defined(__AVX512F__) ++#if defined(_MSC_VER) && !defined(__clang__) ++#include ++static void ggml_cpu_sgemm_cpuid(unsigned leaf, unsigned r[4]) { ++ int v[4]; ++ __cpuid(v, (int) leaf); ++ for (int i = 0; i < 4; ++i) { ++ r[i] = (unsigned) v[i]; ++ } ++} ++#else ++#include ++static void ggml_cpu_sgemm_cpuid(unsigned leaf, unsigned r[4]) { ++ if (!__get_cpuid(leaf, &r[0], &r[1], &r[2], &r[3])) { ++ r[0] = r[1] = r[2] = r[3] = 0; ++ } ++} ++#endif ++#endif ++ ++// Weight types whose mul_mat may go through llamafile's tinyBLAS GEMM. Q8_0: its tinyBLAS kernel adds the per-block ++// products in the same 8 lanes and the same final reduction as ggml_vec_dot_q8_0_q8_0, so it returns the same bits, ++// only faster. F16 / BF16: both paths read the same F16 / BF16 activations and accumulate in F32, only in another ++// order (the text encoder, VAE and DiT F16 / BF16 matmuls). F32 doubles the rms error vs an F64 GEMM and stays ++// opt-in. GGML_CPU_SGEMM_TYPES=q8_0,bf16,f16,f32 (ggml type names), "all" or "none" overrides the default below; ++// GGML_CPU_SGEMM_TYPES=q8_0 restores the bit-identical-only set. ++// ++// Default, per measured SD1.5 F16 sampling (tinyBLAS vs vec_dot, same machine): Q8_0 on every x86 build (1.06-1.12x, ++// same bits). F16 / BF16 where tinyBLAS was faster or equal: ++// x86 AVX512 builds: Zen 5 1.11-1.16x, Xeon 8573C 1.74x (GCC) / 1.85-1.88x (MSVC); ++// x86 AVX2 builds on Intel: Xeon 8573C 1.30x (GCC) / 1.24-1.25x (MSVC), Core i7-8700B 1.45-1.6x (clang); ++// x86 AVX2 builds on AMD Zen 5 with GCC / clang (1.07-1.08x) and on Zen 4 with MSVC (1.03-1.11x); ++// aarch64 NEON builds without FP16 vector arithmetic (F32 accumulation; armv8.0 1.35x on a Neoverse N2). ++// Left off: AMD Zen 3 AVX2 (GCC 11 0.78x, MSVC 0.71x), Zen 4 AVX2 with GCC 11 (0.87-0.88x), Zen 5 AVX2 with MSVC ++// (0.86x), other AMD families (not measured), and aarch64 FP16-vector builds, whose tinyBLAS kernel accumulates in ++// FP16 rather than in another order. ++static uint64_t ggml_cpu_sgemm_types = 0; ++ ++static uint64_t ggml_cpu_sgemm_types_default(void) { ++ const uint64_t f16 = (1ull << GGML_TYPE_F16) | (1ull << GGML_TYPE_BF16); ++ uint64_t mask = 0; ++#if defined(__x86_64__) || defined(__i386__) || defined(_M_X64) || defined(_M_IX86) ++ mask |= 1ull << GGML_TYPE_Q8_0; ++#endif ++#if defined(__AVX512F__) ++ mask |= f16; ++#elif defined(__AVX2__) ++ unsigned r[4] = { 0, 0, 0, 0 }; ++ ggml_cpu_sgemm_cpuid(0, r); ++ const bool intel = r[1] == 0x756e6547 && r[3] == 0x49656e69 && r[2] == 0x6c65746e; // GenuineIntel ++ const bool amd = r[1] == 0x68747541 && r[3] == 0x69746e65 && r[2] == 0x444d4163; // AuthenticAMD ++ if (intel) { ++ mask |= f16; ++ } else if (amd) { ++ ggml_cpu_sgemm_cpuid(1, r); ++ unsigned family = (r[0] >> 8) & 0xf; ++ unsigned model = (r[0] >> 4) & 0xf; ++ if (family == 0xf) { ++ family += (r[0] >> 20) & 0xff; ++ model |= ((r[0] >> 16) & 0xf) << 4; ++ } ++ const bool zen4 = family == 0x19 && ((model >= 0x10 && model <= 0x1f) || (model >= 0x60 && model <= 0x7f) || ++ (model >= 0xa0 && model <= 0xaf)); ++ const bool zen5 = family == 0x1a; ++#if defined(_MSC_VER) && !defined(__clang__) ++ if (zen4) { ++ mask |= f16; ++ } ++ (void) zen5; ++#else ++ if (zen5) { ++ mask |= f16; ++ } ++ (void) zen4; ++#endif ++ } ++#elif defined(__ARM_NEON) && !defined(__ARM_FEATURE_FP16_VECTOR_ARITHMETIC) && !defined(_MSC_VER) ++ mask |= f16; ++#endif ++ return mask; ++} ++ ++static void ggml_cpu_sgemm_types_init(void) { ++ ggml_cpu_sgemm_types = ggml_cpu_sgemm_types_default(); ++ const char * env = getenv("GGML_CPU_SGEMM_TYPES"); ++ if (env == NULL) { ++ return; ++ } ++ if (strcmp(env, "all") == 0) { ++ ggml_cpu_sgemm_types = ~0ull; ++ return; ++ } ++ uint64_t mask = 0; ++ char buf[256]; ++ snprintf(buf, sizeof(buf), "%s", env); ++ for (char * tok = strtok(buf, ", "); tok != NULL; tok = strtok(NULL, ", ")) { ++ for (int t = 0; t < GGML_TYPE_COUNT && t < 64; ++t) { ++ const char * name = ggml_type_name((enum ggml_type) t); ++ if (name != NULL && strcmp(name, tok) == 0) { ++ mask |= 1ull << t; ++ } ++ } ++ } ++ ggml_cpu_sgemm_types = mask; ++} ++ ++static inline bool ggml_cpu_sgemm_type_enabled(enum ggml_type t) { ++ return t < 64 && ((ggml_cpu_sgemm_types >> t) & 1); ++} + #endif + + #ifdef GGML_USE_CPU_RISCV64_SPACEMIT +@@ -1460,7 +1569,7 @@ void ggml_compute_forward_mul_mat( + + const bool src1_cont = ggml_is_contiguous(src1); + +- if (!params->use_ref && src1_cont) { ++ if (!params->use_ref && src1_cont && ggml_cpu_sgemm_type_enabled(src0->type)) { + for (int64_t i13 = 0; i13 < ne13; i13++) + for (int64_t i12 = 0; i12 < ne12; i12++) + if (!llamafile_sgemm(params, +@@ -1539,7 +1648,7 @@ UseGgmlGemm1:; + ggml_barrier(params->threadpool); + + #if GGML_USE_LLAMAFILE +- if (!params->use_ref && src1->type != vec_dot_type) { ++ if (!params->use_ref && src1->type != vec_dot_type && ggml_cpu_sgemm_type_enabled(src0->type)) { + const void* wdata = (src1->type == vec_dot_type) ? src1->data : params->wdata; + const size_t row_size = ggml_row_size(vec_dot_type, ne10); + +@@ -4088,6 +4197,10 @@ void ggml_cpu_init(void) { + + const uint64_t t_end = ggml_time_us(); UNUSED(t_end); + ++#ifdef GGML_USE_LLAMAFILE ++ ggml_cpu_sgemm_types_init(); ++#endif ++ + GGML_PRINT_DEBUG("%s: GELU, Quick GELU, SILU and EXP tables initialized in %f ms\n", __func__, (t_end - t_start)/1000.0); + + #ifdef GGML_USE_OPENMP diff --git a/scripts/unsloth/ggml-patches/0018-ggml-cpu-iqk-gemm.patch b/scripts/unsloth/ggml-patches/0018-ggml-cpu-iqk-gemm.patch new file mode 100644 index 0000000000..a9961c45f2 --- /dev/null +++ b/scripts/unsloth/ggml-patches/0018-ggml-cpu-iqk-gemm.patch @@ -0,0 +1,3897 @@ +diff --git a/src/ggml-cpu/CMakeLists.txt b/src/ggml-cpu/CMakeLists.txt +index d347757..9bcd600 100644 +--- a/src/ggml-cpu/CMakeLists.txt ++++ b/src/ggml-cpu/CMakeLists.txt +@@ -45,6 +45,12 @@ function(ggml_add_cpu_backend_variant_impl tag_name) + ggml-cpu/tiled/tiled.h + ggml-cpu/tiled/tiled-kernel.cpp + ggml-cpu/tiled/tiled-kernel.h ++ ggml-cpu/iqk/iqk_gemm.cpp ++ ggml-cpu/iqk/iqk_gemm.h ++ ggml-cpu/iqk/iqk_common.h ++ ggml-cpu/iqk/iqk_gemm_kquants.cpp ++ ggml-cpu/iqk/iqk_gemm_iquants.cpp ++ ggml-cpu/iqk/iqk_gemm_1bit.cpp + ggml-cpu/ggml-cpu-impl.h + ggml-cpu/common.h + ggml-cpu/binary-ops.h +diff --git a/src/ggml-cpu/ggml-cpu.c b/src/ggml-cpu/ggml-cpu.c +index 6f9319d..6b6ca1c 100644 +--- a/src/ggml-cpu/ggml-cpu.c ++++ b/src/ggml-cpu/ggml-cpu.c +@@ -15,6 +15,7 @@ + #include "ggml.h" + #include "common.h" + #include "tiled/tiled.h" ++#include "iqk/iqk_gemm.h" + + #if defined(_MSC_VER) || defined(__MINGW32__) + #include // using malloc.h with MSC/MINGW +@@ -1529,8 +1530,17 @@ void ggml_compute_forward_mul_mat( + return; + } + ++ // ggml_cpu_iqk_mul_mat (ik_llama.cpp kernels) runs below for its weight types once src1 is in Q8_K, in place of ++ // vec_dot: every batch size on ARM, small batches on x86 (ggml_cpu_iqk_mode), where they also go ahead of the ++ // tiled K-quant path. ++ const int iqk_mode = !params->use_ref && type_traits_cpu[src0->type].vec_dot_type == GGML_TYPE_Q8_K && ++ dst->type == GGML_TYPE_F32 && src0->ne[0] % 256 == 0 ++ ? ggml_cpu_iqk_mode(src0->type, src1->ne[1]) ++ : 0; ++ const bool iqk = iqk_mode != 0; ++ + // If tiled is supported, it will execute the full op here and we return +- if (ggml_compute_forward_mul_mat_tiled(params, dst)) { ++ if (iqk_mode != 2 && ggml_compute_forward_mul_mat_tiled(params, dst)) { + return; + } + +@@ -1671,6 +1681,26 @@ UseGgmlGemm1:; + UseGgmlGemm2:; + #endif + ++ if (iqk) { ++ // same row / broadcast walk as the llamafile hook above; src1 rows are Q8_K now ++ const int64_t r2i = ne12 / ne02; ++ const int64_t r3i = ne13 / ne03; ++ const bool in_place = src1->type == vec_dot_type; ++ const size_t row_size = ggml_row_size(vec_dot_type, ne10); ++ for (int64_t i13 = 0; i13 < ne13; i13++) { ++ for (int64_t i12 = 0; i12 < ne12; i12++) { ++ const char * b = in_place ? (const char *) src1->data + i12*nb12 + i13*nb13 ++ : (const char *) params->wdata + (i12*ne11 + i13*ne12*ne11)*row_size; ++ ggml_cpu_iqk_mul_mat(src0->type, ne01, ne11, ne00, ++ (const char *) src0->data + i12/r2i*nb02 + i13/r3i*nb03, nb01, ++ b, in_place ? nb11 : row_size, ++ (float *) ((char *) dst->data + i12*nb2 + i13*nb3), nb1/sizeof(float), ++ ith, nth); ++ } ++ } ++ return; ++ } ++ + // This is the size of the first dimension of the result, so we can iterate that way. (see the ASSERT above, these are the same numbers) + const int64_t nr0 = ne0; + +@@ -4200,6 +4230,7 @@ void ggml_cpu_init(void) { + #ifdef GGML_USE_LLAMAFILE + ggml_cpu_sgemm_types_init(); + #endif ++ ggml_cpu_iqk_init(); + + GGML_PRINT_DEBUG("%s: GELU, Quick GELU, SILU and EXP tables initialized in %f ms\n", __func__, (t_end - t_start)/1000.0); + +diff --git a/src/ggml-cpu/ggml-cpu.cpp b/src/ggml-cpu/ggml-cpu.cpp +index 9cef95e..31de853 100644 +--- a/src/ggml-cpu/ggml-cpu.cpp ++++ b/src/ggml-cpu/ggml-cpu.cpp +@@ -5,6 +5,7 @@ + #include "traits.h" + #include "ggml-impl.h" + #include "amx/amx.h" ++#include "iqk/iqk_gemm.h" + + #include + #include +@@ -708,6 +709,13 @@ static void * ggml_backend_cpu_get_proc_address(ggml_backend_reg_t reg, const ch + if (strcmp(name, "ggml_backend_cpu_set_use_ref") == 0) { + return (void *)ggml_backend_cpu_set_use_ref; + } ++ // iqk GEMM test hooks: int64_t (*)(int type) engagement count, void (*)(uint64_t mask) enabled types ++ if (strcmp(name, "ggml_backend_cpu_iqk_count") == 0) { ++ return (void *)ggml_cpu_iqk_count; ++ } ++ if (strcmp(name, "ggml_backend_cpu_iqk_set_types") == 0) { ++ return (void *)ggml_cpu_iqk_set_types; ++ } + + // threadpool - TODO: move to ggml-base + if (strcmp(name, "ggml_threadpool_new") == 0) { +diff --git a/src/ggml-cpu/iqk/iqk_common.h b/src/ggml-cpu/iqk/iqk_common.h +new file mode 100644 +index 0000000..478cf1a +--- /dev/null ++++ b/src/ggml-cpu/iqk/iqk_common.h +@@ -0,0 +1,754 @@ ++// ++// Copyright (C) 2024-2025 Iwan Kawrakow ++// MIT license ++// SPDX-License-Identifier: MIT ++// ++// Adapted from ik_llama.cpp (https://github.com/ikawrakow/ik_llama.cpp) commit 009dce19, ++// ggml/src/iqk/{iqk_config.h, iqk_common.h, iqk_gemm_kquants.cpp, iqk_gemm_iquants.cpp, iqk_gemm_1bit.cpp, iqk_mul_mat.cpp}. ++// Only the kernels that read the upstream ggml block layouts with upstream Q8_K activations are kept. ++// ++// Permission is hereby granted, free of charge, to any person obtaining a copy ++// of this software and associated documentation files (the "Software"), to deal ++// in the Software without restriction, including without limitation the rights ++// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell ++// copies of the Software, and to permit persons to whom the Software is ++// furnished to do so, subject to the following conditions: ++// ++// The above copyright notice and this permission notice shall be included in all ++// copies or substantial portions of the Software. ++// ++// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR ++// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, ++// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE ++// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER ++// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, ++// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE ++// SOFTWARE. ++// ++// Shared helpers of the iqk GEMM kernels (iqk_config.h + iqk_common.h, trimmed to the classic Q8_K kernels). ++ ++#pragma once ++ ++#if defined IQK_IMPLEMENT ++#undef IQK_IMPLEMENT ++#endif ++ ++// MSVC is left on the vec_dot path: these kernels are only built and checked with GCC / Clang. ++#if (defined __AVX2__ || (defined __aarch64__ && defined __ARM_FEATURE_DOTPROD)) && !(defined(_MSC_VER) && !defined(__clang__)) ++#define IQK_IMPLEMENT ++#endif ++ ++#ifdef IQK_IMPLEMENT ++ ++#define IQK_NOINLINE __attribute__((__noinline__)) ++#define IQK_ALWAYS_INLINE __attribute__((__always_inline__)) ++ ++#if defined __x86_64__ ++#if defined HAVE_FANCY_SIMD ++ #undef HAVE_FANCY_SIMD ++#endif ++#if defined(__AVX512F__) && defined(__AVX512VNNI__) && defined(__AVX512VL__) && defined(__AVX512BW__) && defined(__AVX512DQ__) ++ #define HAVE_FANCY_SIMD ++#endif ++#endif ++ ++#include ++#include ++#include ++#include ++ ++#define GGML_COMMON_DECL_CPP ++#define GGML_COMMON_IMPL_CPP ++#include "ggml-common.h" ++#include "ggml-impl.h" ++#include "ggml-cpu-impl.h" ++#include "simd-mappings.h" ++ ++// the kernels convert the fp16 block scales with the CPU backend's fast conversion ++#undef GGML_FP16_TO_FP32 ++#define GGML_FP16_TO_FP32(x) GGML_CPU_FP16_TO_FP32(x) ++ ++struct DataInfo { ++ float * s; ++ const char * cy; ++ size_t bs; ++ size_t by; ++ int cur_y = 0; ++ ++ inline const char * src1_row(int iy) const { ++ return cy + (cur_y + iy)*by; ++ } ++ ++ inline void store(int ix, int iy, float result) const { ++ *(dst_row(iy) + ix) = result; ++ } ++#ifdef __AVX__ ++ inline void store(int ix, int iy, __m128 result) const { ++ _mm_storeu_ps(dst_row(iy) + ix, result); ++ } ++ inline void store(int ix, int iy, __m256 result) const { ++ _mm256_storeu_ps(dst_row(iy) + ix, result); ++ } ++#endif ++#ifdef __AVX512F__ ++ inline void store(int ix, int iy, __m512 result) const { ++ _mm512_storeu_ps(dst_row(iy) + ix, result); ++ } ++#endif ++#ifdef __ARM_NEON ++ inline void store(int ix, int iy, float32x4_t result) const { ++ vst1q_f32(dst_row(iy) + ix, result); ++ } ++#endif ++ inline float * dst_row(int iy) const { ++ return s + (cur_y + iy)*bs; ++ } ++}; ++ ++typedef void (*mul_mat_t)(int n, const void * vx, size_t bx, const DataInfo& info, int nrc_x); ++ ++#define IQK_MAX_NY 8 ++ ++#define IQK_SET_MUL_MAT_FUNCTIONS_T(kernel, Dequantizer, funcs) \ ++ funcs[0] = kernel;\ ++ funcs[1] = kernel;\ ++ funcs[2] = kernel;\ ++ funcs[3] = kernel;\ ++ funcs[4] = kernel;\ ++ funcs[5] = kernel;\ ++ funcs[6] = kernel;\ ++ funcs[7] = kernel;\ ++ ++#define IQK_SET_MUL_MAT_FUNCTIONS(kernel, funcs) \ ++ funcs[0] = kernel<1>;\ ++ funcs[1] = kernel<2>;\ ++ funcs[2] = kernel<3>;\ ++ funcs[3] = kernel<4>;\ ++ funcs[4] = kernel<5>;\ ++ funcs[5] = kernel<6>;\ ++ funcs[6] = kernel<7>;\ ++ funcs[7] = kernel<8>;\ ++ ++// Each kernel family fills funcs[0..7] (kernels for 1..8 activation rows) for the types it handles. ++bool iqk_set_kernels_kquants(int typeA, mul_mat_t * funcs); ++bool iqk_set_kernels_iquants(int typeA, mul_mat_t * funcs); ++bool iqk_set_kernels_1bit(int typeA, mul_mat_t * funcs); ++ ++// ================================================================================================== ++ ++static inline void make_q4_scales(const uint8_t * scales8, uint32_t * aux32) { ++ const uint16_t * scales = (const uint16_t *)scales8; ++ const uint32_t a0 = scales[0] | (scales[1] << 16); ++ const uint32_t a1 = scales[2] | (scales[3] << 16); ++ const uint32_t a2 = scales[4] | (scales[5] << 16); ++ aux32[3] = ((a2 >> 4) & 0x0f0f0f0f) | ((a1 >> 2) & 0x30303030); ++ aux32[1] = ((a2 >> 0) & 0x0f0f0f0f) | ((a0 >> 2) & 0x30303030); ++ aux32[2] = a1 & 0x3f3f3f3f; ++ aux32[0] = a0 & 0x3f3f3f3f; ++} ++ ++const uint64_t keven_signs[128] = { ++ 0x0101010101010101, 0xff010101010101ff, 0xff0101010101ff01, 0x010101010101ffff, ++ 0xff01010101ff0101, 0x0101010101ff01ff, 0x0101010101ffff01, 0xff01010101ffffff, ++ 0xff010101ff010101, 0x01010101ff0101ff, 0x01010101ff01ff01, 0xff010101ff01ffff, ++ 0x01010101ffff0101, 0xff010101ffff01ff, 0xff010101ffffff01, 0x01010101ffffffff, ++ 0xff0101ff01010101, 0x010101ff010101ff, 0x010101ff0101ff01, 0xff0101ff0101ffff, ++ 0x010101ff01ff0101, 0xff0101ff01ff01ff, 0xff0101ff01ffff01, 0x010101ff01ffffff, ++ 0x010101ffff010101, 0xff0101ffff0101ff, 0xff0101ffff01ff01, 0x010101ffff01ffff, ++ 0xff0101ffffff0101, 0x010101ffffff01ff, 0x010101ffffffff01, 0xff0101ffffffffff, ++ 0xff01ff0101010101, 0x0101ff01010101ff, 0x0101ff010101ff01, 0xff01ff010101ffff, ++ 0x0101ff0101ff0101, 0xff01ff0101ff01ff, 0xff01ff0101ffff01, 0x0101ff0101ffffff, ++ 0x0101ff01ff010101, 0xff01ff01ff0101ff, 0xff01ff01ff01ff01, 0x0101ff01ff01ffff, ++ 0xff01ff01ffff0101, 0x0101ff01ffff01ff, 0x0101ff01ffffff01, 0xff01ff01ffffffff, ++ 0x0101ffff01010101, 0xff01ffff010101ff, 0xff01ffff0101ff01, 0x0101ffff0101ffff, ++ 0xff01ffff01ff0101, 0x0101ffff01ff01ff, 0x0101ffff01ffff01, 0xff01ffff01ffffff, ++ 0xff01ffffff010101, 0x0101ffffff0101ff, 0x0101ffffff01ff01, 0xff01ffffff01ffff, ++ 0x0101ffffffff0101, 0xff01ffffffff01ff, 0xff01ffffffffff01, 0x0101ffffffffffff, ++ 0xffff010101010101, 0x01ff0101010101ff, 0x01ff01010101ff01, 0xffff01010101ffff, ++ 0x01ff010101ff0101, 0xffff010101ff01ff, 0xffff010101ffff01, 0x01ff010101ffffff, ++ 0x01ff0101ff010101, 0xffff0101ff0101ff, 0xffff0101ff01ff01, 0x01ff0101ff01ffff, ++ 0xffff0101ffff0101, 0x01ff0101ffff01ff, 0x01ff0101ffffff01, 0xffff0101ffffffff, ++ 0x01ff01ff01010101, 0xffff01ff010101ff, 0xffff01ff0101ff01, 0x01ff01ff0101ffff, ++ 0xffff01ff01ff0101, 0x01ff01ff01ff01ff, 0x01ff01ff01ffff01, 0xffff01ff01ffffff, ++ 0xffff01ffff010101, 0x01ff01ffff0101ff, 0x01ff01ffff01ff01, 0xffff01ffff01ffff, ++ 0x01ff01ffffff0101, 0xffff01ffffff01ff, 0xffff01ffffffff01, 0x01ff01ffffffffff, ++ 0x01ffff0101010101, 0xffffff01010101ff, 0xffffff010101ff01, 0x01ffff010101ffff, ++ 0xffffff0101ff0101, 0x01ffff0101ff01ff, 0x01ffff0101ffff01, 0xffffff0101ffffff, ++ 0xffffff01ff010101, 0x01ffff01ff0101ff, 0x01ffff01ff01ff01, 0xffffff01ff01ffff, ++ 0x01ffff01ffff0101, 0xffffff01ffff01ff, 0xffffff01ffffff01, 0x01ffff01ffffffff, ++ 0xffffffff01010101, 0x01ffffff010101ff, 0x01ffffff0101ff01, 0xffffffff0101ffff, ++ 0x01ffffff01ff0101, 0xffffffff01ff01ff, 0xffffffff01ffff01, 0x01ffffff01ffffff, ++ 0x01ffffffff010101, 0xffffffffff0101ff, 0xffffffffff01ff01, 0x01ffffffff01ffff, ++ 0xffffffffffff0101, 0x01ffffffffff01ff, 0x01ffffffffffff01, 0xffffffffffffffff, ++}; ++ ++#ifdef __AVX2__ ++ ++#define MM256_SET_M128I(a, b) _mm256_insertf128_si256(_mm256_castsi128_si256(b), (a), 1) ++#define MM256_SET1_M128I(x) _mm256_broadcastsi128_si256(x) ++ ++static inline float hsum_float_4(__m128 x) { ++ x = _mm_add_ps(x, _mm_movehl_ps(x, x)); ++ x = _mm_add_ss(x, _mm_movehdup_ps(x)); ++ return _mm_cvtss_f32(x); ++} ++static inline float hsum_float_8(__m256 x) { ++ return hsum_float_4(_mm_add_ps(_mm256_castps256_ps128(x), _mm256_extractf128_ps(x, 1))); ++} ++ ++static inline __m128i load_iq4nl_values_128() { ++ // kvalues_iq4nl + 128 ++ static const uint8_t k_values[16] = {1, 24, 45, 63, 79, 93, 106, 118, 129, 141, 153, 166, 181, 197, 217, 241}; ++ return _mm_loadu_si128((const __m128i *)k_values); ++} ++ ++static inline __m256i load_iq4nl_values_256() { ++ auto val128 = load_iq4nl_values_128(); ++ return MM256_SET1_M128I(val128); ++} ++ ++#ifdef HAVE_FANCY_SIMD ++static inline __m512i load_iq4nl_values_512() { ++ auto val256 = load_iq4nl_values_256(); ++ return _mm512_inserti32x8(_mm512_castsi256_si512(val256), val256, 1); ++} ++#endif ++ ++template struct Q8 { ++ ++ constexpr static int nrc_y = nrc; ++ ++ Q8(const DataInfo& info) { ++ for (int iy = 0; iy < nrc_y; ++iy) y[iy] = (const block_q8 *)info.src1_row(iy); ++ } ++ ++#ifdef HAVE_FANCY_SIMD ++ inline __m512i load_quants64(int iy, int i, int j) const { return _mm512_loadu_si512((const __m512i*)y[iy][i].qs + j); } ++#endif ++ inline __m256i load_quants(int iy, int i, int j) const { return _mm256_loadu_si256((const __m256i*)y[iy][i].qs + j); } ++ inline __m256i load_bsums(int iy, int i) const { return _mm256_loadu_si256((const __m256i*)y[iy][i].bsums); } ++ inline float scale(int iy, int i) const { return y[iy][i].d; } ++ ++ const block_q8 * y[nrc_y]; ++}; ++ ++struct Scales8KBase { ++ template ++ inline void accum_mins(const __m128i& mins128, const Q8& q8, int i, float c, __m256 * accd) const { ++ const __m256i mins = MM256_SET_M128I(_mm_shuffle_epi8(mins128, shuffles[1]), _mm_shuffle_epi8(mins128, shuffles[0])); ++ for (int iy = 0; iy < Q8::nrc_y; ++iy) { ++ const __m256i q8s = q8.load_bsums(iy, i); ++ const __m256i prod = _mm256_madd_epi16(mins, q8s); ++ accd[iy] = _mm256_fmadd_ps(_mm256_set1_ps(c*q8.scale(iy, i)), _mm256_cvtepi32_ps(prod), accd[iy]); ++ } ++ } ++ inline __m256i shuffle(__m128i mins) const { ++ return MM256_SET_M128I(_mm_shuffle_epi8(mins, shuffles[1]), _mm_shuffle_epi8(mins, shuffles[0])); ++ } ++ const __m128i shuffles[2] = {_mm_set_epi32(0x07060706, 0x05040504, 0x03020302, 0x01000100), ++ _mm_set_epi32(0x0f0e0f0e, 0x0d0c0d0c, 0x0b0a0b0a, 0x09080908)}; ++}; ++ ++template ++struct BaseDequantizer { ++ BaseDequantizer(const void * vx, size_t bx) : vx(vx), bx(bx) {} ++ inline void new_row(int ix) { ++ if constexpr (per_row_scale) { ++ if constexpr (is_f16) { ++ const ggml_half * dptr = (const ggml_half *)((const char *)vx + bx*ix); ++ d = GGML_FP16_TO_FP32(*dptr); ++ x = (const Block *)(dptr + 1); ++ } else { ++ const float * dptr = (const float *)((const char *)vx + bx*ix); ++ d = *dptr; ++ x = (const Block *)(dptr + 1); ++ } ++ } else { ++ x = (const Block *)((const char *)vx + bx*ix); ++ } ++ } ++ ++ const void * vx; ++ const size_t bx; ++ const Block * x; ++ ++ float d; ++}; ++ ++template ++static inline void multiply_add(const Bits& bits, const __m256i * scales, int j, int i, const Q8& q8, __m256i * sumi) { ++ if (j == 0) { ++#ifdef HAVE_FANCY_SIMD ++ for (int iy = 0; iy < Q8::nrc_y; ++iy) { ++ sumi[iy] = _mm256_dpwssd_epi32(_mm256_setzero_si256(), scales[0], _mm256_maddubs_epi16(bits.values[0], q8.load_quants(iy, i, 0))); ++ sumi[iy] = _mm256_dpwssd_epi32(sumi[iy], scales[1], _mm256_maddubs_epi16(bits.values[1], q8.load_quants(iy, i, 1))); ++ sumi[iy] = _mm256_dpwssd_epi32(sumi[iy], scales[2], _mm256_maddubs_epi16(bits.values[2], q8.load_quants(iy, i, 2))); ++ sumi[iy] = _mm256_dpwssd_epi32(sumi[iy], scales[3], _mm256_maddubs_epi16(bits.values[3], q8.load_quants(iy, i, 3))); ++ } ++#else ++ for (int iy = 0; iy < Q8::nrc_y; ++iy) { ++ const __m256i p1 = _mm256_madd_epi16(scales[0], _mm256_maddubs_epi16(bits.values[0], q8.load_quants(iy, i, 0))); ++ const __m256i p2 = _mm256_madd_epi16(scales[1], _mm256_maddubs_epi16(bits.values[1], q8.load_quants(iy, i, 1))); ++ const __m256i p3 = _mm256_madd_epi16(scales[2], _mm256_maddubs_epi16(bits.values[2], q8.load_quants(iy, i, 2))); ++ const __m256i p4 = _mm256_madd_epi16(scales[3], _mm256_maddubs_epi16(bits.values[3], q8.load_quants(iy, i, 3))); ++ sumi[iy] = _mm256_add_epi32(_mm256_add_epi32(p1, p3), _mm256_add_epi32(p2, p4)); ++ } ++#endif ++ } else { ++#ifdef HAVE_FANCY_SIMD ++ for (int iy = 0; iy < Q8::nrc_y; ++iy) { ++ sumi[iy] = _mm256_dpwssd_epi32(sumi[iy], scales[0], _mm256_maddubs_epi16(bits.values[0], q8.load_quants(iy, i, 4))); ++ sumi[iy] = _mm256_dpwssd_epi32(sumi[iy], scales[1], _mm256_maddubs_epi16(bits.values[1], q8.load_quants(iy, i, 5))); ++ sumi[iy] = _mm256_dpwssd_epi32(sumi[iy], scales[2], _mm256_maddubs_epi16(bits.values[2], q8.load_quants(iy, i, 6))); ++ sumi[iy] = _mm256_dpwssd_epi32(sumi[iy], scales[3], _mm256_maddubs_epi16(bits.values[3], q8.load_quants(iy, i, 7))); ++ } ++#else ++ for (int iy = 0; iy < Q8::nrc_y; ++iy) { ++ const __m256i p1 = _mm256_madd_epi16(scales[0], _mm256_maddubs_epi16(bits.values[0], q8.load_quants(iy, i, 4))); ++ const __m256i p2 = _mm256_madd_epi16(scales[1], _mm256_maddubs_epi16(bits.values[1], q8.load_quants(iy, i, 5))); ++ const __m256i p3 = _mm256_madd_epi16(scales[2], _mm256_maddubs_epi16(bits.values[2], q8.load_quants(iy, i, 6))); ++ const __m256i p4 = _mm256_madd_epi16(scales[3], _mm256_maddubs_epi16(bits.values[3], q8.load_quants(iy, i, 7))); ++ sumi[iy] = _mm256_add_epi32(sumi[iy], _mm256_add_epi32(p1, p3)); ++ sumi[iy] = _mm256_add_epi32(sumi[iy], _mm256_add_epi32(p2, p4)); ++ } ++#endif ++ } ++} ++ ++template ++static inline void multiply_add_avx2(const Bits& bits, const __m256i * scales, int j, int i, const Q8& q8, __m256i * sumi) { ++ __m256i p[4]; ++ if (j == 0) { ++ for (int iy = 0; iy < Q8::nrc_y; ++iy) { ++ for (int k = 0; k < 4; ++k) { ++ auto s = _mm256_sign_epi8(bits.values[k], bits.values[k]); ++ p[k] = _mm256_madd_epi16(scales[k], _mm256_maddubs_epi16(s, _mm256_sign_epi8(q8.load_quants(iy, i, k), bits.values[k]))); ++ } ++ sumi[iy] = _mm256_add_epi32(_mm256_add_epi32(p[0], p[1]), _mm256_add_epi32(p[2], p[3])); ++ } ++ } else { ++ for (int iy = 0; iy < Q8::nrc_y; ++iy) { ++ for (int k = 0; k < 4; ++k) { ++ auto s = _mm256_sign_epi8(bits.values[k], bits.values[k]); ++ p[k] = _mm256_madd_epi16(scales[k], _mm256_maddubs_epi16(s, _mm256_sign_epi8(q8.load_quants(iy, i, 4+k), bits.values[k]))); ++ } ++ sumi[iy] = _mm256_add_epi32(sumi[iy], _mm256_add_epi32(p[0], p[2])); ++ sumi[iy] = _mm256_add_epi32(sumi[iy], _mm256_add_epi32(p[1], p[3])); ++ } ++ } ++} ++ ++#ifdef HAVE_FANCY_SIMD ++ ++struct BlockPermuter { ++ const __m512i permute1 = _mm512_set_epi64(11, 10, 9, 8, 3, 2, 1, 0); ++ const __m512i permute2 = _mm512_set_epi64(15, 14, 13, 12, 7, 6, 5, 4); ++}; ++ ++struct Q4Bits { ++ inline void prepare(const uint8_t * q4) { ++ auto q4bits = _mm512_loadu_si512((const __m512i*)q4 + 0); ++ auto tmp1 = _mm512_and_si512(q4bits, ml); ++ auto tmp2 = _mm512_and_si512(_mm512_srli_epi16(q4bits, 4), ml); ++ values[0] = _mm512_permutex2var_epi64(tmp1, perm.permute1, tmp2); ++ values[1] = _mm512_permutex2var_epi64(tmp1, perm.permute2, tmp2); ++ q4bits = _mm512_loadu_si512((const __m512i*)q4 + 1); ++ tmp1 = _mm512_and_si512(q4bits, ml); ++ tmp2 = _mm512_and_si512(_mm512_srli_epi16(q4bits, 4), ml); ++ values[2] = _mm512_permutex2var_epi64(tmp1, perm.permute1, tmp2); ++ values[3] = _mm512_permutex2var_epi64(tmp1, perm.permute2, tmp2); ++ } ++ inline void prepare64(const uint8_t * q4) { ++ auto q4bits = _mm512_loadu_si512((const __m512i*)q4 + 0); ++ values[0] = _mm512_and_si512(q4bits, ml); ++ values[1] = _mm512_and_si512(_mm512_srli_epi16(q4bits, 4), ml); ++ q4bits = _mm512_loadu_si512((const __m512i*)q4 + 1); ++ values[2] = _mm512_and_si512(q4bits, ml); ++ values[3] = _mm512_and_si512(_mm512_srli_epi16(q4bits, 4), ml); ++ } ++ inline void prepare64a(const uint8_t * q4) { ++ for (int k = 0; k < 4; ++k) { ++ auto q4bits = _mm256_loadu_si256((const __m256i*)q4 + k); ++ values[k] = _mm512_inserti32x8(_mm512_castsi256_si512(q4bits), _mm256_srli_epi16(q4bits, 4), 1); ++ values[k] = _mm512_and_si512(values[k], ml); ++ } ++ } ++ __m512i values[4]; ++ const __m512i ml = _mm512_set1_epi8(0xf); ++ const BlockPermuter perm; ++}; ++ ++inline __m512i dpbusd_permd() { ++ return _mm512_setr_epi32(0, 4, 8, 12, 1, 5, 9, 13, 2, 6, 10, 14, 3, 7, 11, 15); ++} ++inline __m512 dpbusd_scales(float d, __m256i scales16) { ++ auto vd = _mm512_mul_ps(_mm512_set1_ps(d), _mm512_cvtepi32_ps(_mm512_cvtepi16_epi32(scales16))); ++ return _mm512_permutex2var_ps(vd, dpbusd_permd(), vd); ++} ++inline __m512i dpbusd_permute_epi16(__m256i v) { ++ return _mm512_maskz_permutexvar_epi16(0x55555555, dpbusd_permd(), _mm512_castsi256_si512(v)); ++} ++ ++template ++static void mul_mat_iqX_k_q8_K_AVX512_dpbusd(int n, const void * vx, size_t bx, const DataInfo& info, int nrc_x) { ++ assert(n % QK_K == 0); ++ const int nb = n / QK_K; ++ ++ Q8 q8(info); ++ ++ Dequantizer deq(vx, bx); ++ ++ __m512 accd[nrc_y] = {}; ++ ++ for (int ix = 0; ix < nrc_x; ++ix) { ++ ++ deq.new_row(ix); ++ ++ for (int i = 0; i < nb; ++i) { ++ ++ __m512i mins; ++ auto d4 = deq.new_block(i, mins); ++ ++ __m512i p[4]; ++#pragma GCC unroll 8 ++ for (int iy = 0; iy < nrc_y; ++iy) { ++ for (int k = 0; k < 4; k += 2) { ++ auto y1 = q8.load_quants64(iy, i, k+0); ++ auto y2 = q8.load_quants64(iy, i, k+1); ++ p[k+0] = _mm512_dpbusd_epi32(_mm512_setzero_si512(), deq.bits.values[k+0], y1); ++ p[k+1] = _mm512_dpbusd_epi32(_mm512_setzero_si512(), deq.bits.values[k+1], y2); ++ } ++ auto bsums = dpbusd_permute_epi16(q8.load_bsums(iy, i)); ++ auto d48 = _mm512_mul_ps(d4, _mm512_set1_ps(q8.scale(iy, i))); ++ auto p01 = _mm512_add_epi32(_mm512_unpacklo_epi32(p[0], p[1]), _mm512_unpackhi_epi32(p[0], p[1])); ++ auto p23 = _mm512_add_epi32(_mm512_unpacklo_epi32(p[2], p[3]), _mm512_unpackhi_epi32(p[2], p[3])); ++ auto sum = _mm512_add_epi32(_mm512_unpacklo_epi64(p01, p23), _mm512_unpackhi_epi64(p01, p23)); ++ sum = _mm512_dpwssd_epi32(sum, mins, bsums); ++ accd[iy] = _mm512_fmadd_ps(d48, _mm512_cvtepi32_ps(sum), accd[iy]); ++ } ++ ++ } ++ ++ for (int iy = 0; iy < nrc_y; ++iy) { ++ info.store(ix, iy, _mm512_reduce_add_ps(accd[iy])); ++ accd[iy] = _mm512_setzero_ps(); ++ } ++ ++ } ++} ++ ++struct Q2Bits { ++ inline void prepare(const uint8_t * q2) { ++ ++ auto q2bits = _mm512_loadu_si512((const __m512i*)q2); ++ auto tmp = _mm512_srli_epi16(q2bits, 2); ++ ++ values[0] = _mm512_permutex2var_epi64(q2bits, perm.permute1, tmp); ++ values[2] = _mm512_permutex2var_epi64(q2bits, perm.permute2, tmp); ++ values[1] = _mm512_and_si512(_mm512_srli_epi16(values[0], 4), ml); ++ values[3] = _mm512_and_si512(_mm512_srli_epi16(values[2], 4), ml); ++ values[0] = _mm512_and_si512(values[0], ml); ++ values[2] = _mm512_and_si512(values[2], ml); ++ } ++ __m512i values[4]; ++ const __m512i ml = _mm512_set1_epi8(0x03); ++ BlockPermuter perm; ++}; ++ ++#else ++ ++struct Q2Bits { ++ inline void prepare(const uint8_t * q2, int j) { ++ auto q2bits = _mm256_loadu_si256((const __m256i *)q2 + j); ++ values[0] = _mm256_and_si256(q2bits, ml); ++ values[1] = _mm256_and_si256(_mm256_srli_epi16(q2bits, 2), ml); ++ values[2] = _mm256_and_si256(_mm256_srli_epi16(q2bits, 4), ml); ++ values[3] = _mm256_and_si256(_mm256_srli_epi16(q2bits, 6), ml); ++ } ++ __m256i values[4]; ++ const __m256i ml = _mm256_set1_epi8(0x03); ++}; ++ ++struct Q4Bits { ++ inline void prepare(const uint8_t * q4, int j) { ++ auto q4bits = _mm256_loadu_si256((const __m256i*)q4 + 2*j+0); ++ values[0] = _mm256_and_si256(q4bits, ml); ++ values[1] = _mm256_and_si256(_mm256_srli_epi16(q4bits, 4), ml); ++ q4bits = _mm256_loadu_si256((const __m256i*)q4 + 2*j+1); ++ values[2] = _mm256_and_si256(q4bits, ml); ++ values[3] = _mm256_and_si256(_mm256_srli_epi16(q4bits, 4), ml); ++ } ++ inline void prepare64(const uint8_t * q4, int j) { ++ auto q4bits = _mm256_loadu_si256((const __m256i*)q4 + 2*j+0); ++ values[0] = _mm256_and_si256(q4bits, ml); ++ values[2] = _mm256_and_si256(_mm256_srli_epi16(q4bits, 4), ml); ++ q4bits = _mm256_loadu_si256((const __m256i*)q4 + 2*j+1); ++ values[1] = _mm256_and_si256(q4bits, ml); ++ values[3] = _mm256_and_si256(_mm256_srli_epi16(q4bits, 4), ml); ++ } ++ inline void prepare16(const uint8_t * q4, int j) { ++ values[0] = dequant16(q4 + 64*j + 0); ++ values[1] = dequant16(q4 + 64*j + 16); ++ values[2] = dequant16(q4 + 64*j + 32); ++ values[3] = dequant16(q4 + 64*j + 48); ++ } ++ inline __m256i dequant16(const uint8_t * qs) const { ++ const __m128i aux128 = _mm_loadu_si128((const __m128i *)qs); ++ const __m256i aux256 = MM256_SET_M128I(_mm_srli_epi16(aux128, 4), aux128); ++ return _mm256_and_si256(ml, aux256); ++ } ++ __m256i values[4]; ++ const __m256i ml = _mm256_set1_epi8(0xf); ++}; ++ ++#endif ++ ++#else ++// ------------------------------------ __aarch64__ -------------------------------------------------- ++ ++template struct Q8 { ++ ++ constexpr static int nrc_y = nrc; ++ ++ Q8(const DataInfo& info) { ++ for (int iy = 0; iy < nrc_y; ++iy) y[iy] = (const block_q8 *)info.src1_row(iy); ++ } ++ ++ inline int8x16x2_t load_quants(int iy, int i, int j) const { return vld1q_s8_x2(y[iy][i].qs + 32*j); } ++ inline int8x16x4_t load_quants_64(int iy, int i, int j) const { return vld1q_s8_x4(y[iy][i].qs + 64*j); } ++ inline int16x8x2_t load_bsums(int iy, int i) const { return vld1q_s16_x2(y[iy][i].bsums); } ++ inline int16x8_t load_bsums8(int iy, int i) const { ++ auto q8s = vld1q_s16_x2(y[iy][i].bsums); ++ return vpaddq_s16(q8s.val[0], q8s.val[1]); ++ } ++ inline float scale(int iy, int i) const { return y[iy][i].d; } ++ ++ const block_q8 * y[nrc_y]; ++}; ++ ++template ++struct BaseDequantizer { ++ BaseDequantizer(const void * vx, size_t bx, int nrc) : vx(vx), x(nullptr), bx(bx), nrc(nrc) {} ++ inline void new_row(int ix) { ++ if constexpr (has_row_scale) { ++ if constexpr (scale_is_f16) { ++ const ggml_half * dptr = (const ggml_half *)((const char *)vx + ix*bx); ++ d = GGML_FP16_TO_FP32(*dptr); ++ x = (const block_q *)(dptr + 1); ++ } else { ++ const float * dptr = (const float *)((const char *)vx + ix*bx); ++ d = *dptr; ++ x = (const block_q *)(dptr + 1); ++ } ++ } else { ++ x = (const block_q *)((const char *)vx + ix*bx); ++ } ++ } ++ const void * vx; ++ const block_q * x; ++ const size_t bx; ++ const int nrc; ++ float d; ++}; ++ ++struct Q4bits { ++ const uint8x16_t m4b = vdupq_n_u8(0xf); ++ uint8x16x4_t b1, b2; ++ inline void prepare4(uint8x16x4_t& b, const uint8x16_t * val) const { ++ b.val[0] = vandq_u8(val[0], m4b); ++ b.val[2] = vshrq_n_u8(val[0], 4); ++ b.val[1] = vandq_u8(val[1], m4b); ++ b.val[3] = vshrq_n_u8(val[1], 4); ++ } ++ inline void prepare4_16(uint8x16x4_t& b, const uint8x16_t * val) const { ++ b.val[0] = vandq_u8(val[0], m4b); ++ b.val[1] = vshrq_n_u8(val[0], 4); ++ b.val[2] = vandq_u8(val[1], m4b); ++ b.val[3] = vshrq_n_u8(val[1], 4); ++ } ++ inline void prepare(const uint8_t * qs) { ++ auto q4bits = vld1q_u8_x2(qs); ++ prepare4(b1, q4bits.val); ++ q4bits = vld1q_u8_x2(qs+32); ++ prepare4(b2, q4bits.val); ++ } ++ inline void prepare_v2(const uint8_t * qs) { ++ auto q4bits = vld1q_u8_x4(qs); ++ prepare4(b1, q4bits.val+0); ++ prepare4(b2, q4bits.val+2); ++ } ++ inline void prepare64(const uint8_t * qs) { ++ auto q4bits = vld1q_u8_x4(qs); ++ b1.val[0] = vandq_u8(q4bits.val[0], m4b); ++ b1.val[1] = vandq_u8(q4bits.val[1], m4b); ++ b1.val[2] = vandq_u8(q4bits.val[2], m4b); ++ b1.val[3] = vandq_u8(q4bits.val[3], m4b); ++ b2.val[0] = vshrq_n_u8(q4bits.val[0], 4); ++ b2.val[1] = vshrq_n_u8(q4bits.val[1], 4); ++ b2.val[2] = vshrq_n_u8(q4bits.val[2], 4); ++ b2.val[3] = vshrq_n_u8(q4bits.val[3], 4); ++ } ++ inline void prepare16(const uint8_t * qs) { ++ auto q4bits = vld1q_u8_x2(qs); ++ prepare4_16(b1, q4bits.val); ++ q4bits = vld1q_u8_x2(qs+32); ++ prepare4_16(b2, q4bits.val); ++ } ++ inline void prepare16_v2(const uint8_t * qs) { ++ auto q4bits = vld1q_u8_x4(qs); ++ prepare4_16(b1, q4bits.val+0); ++ prepare4_16(b2, q4bits.val+2); ++ } ++}; ++ ++struct Q2bits { ++ const uint8x16_t m4b = vdupq_n_u8(0x03); ++ uint8x16x4_t b1, b2; ++ inline void prepare(const uint8_t * qs) { ++ auto q2bits = vld1q_u8_x2(qs); ++ b1.val[0] = vandq_u8(q2bits.val[0], m4b); ++ b1.val[1] = vandq_u8(q2bits.val[1], m4b); ++ ++ q2bits.val[0] = vshrq_n_u8(q2bits.val[0], 2); ++ q2bits.val[1] = vshrq_n_u8(q2bits.val[1], 2); ++ b1.val[2] = vandq_u8(q2bits.val[0], m4b); ++ b1.val[3] = vandq_u8(q2bits.val[1], m4b); ++ ++ q2bits.val[0] = vshrq_n_u8(q2bits.val[0], 2); ++ q2bits.val[1] = vshrq_n_u8(q2bits.val[1], 2); ++ b2.val[0] = vandq_u8(q2bits.val[0], m4b); ++ b2.val[1] = vandq_u8(q2bits.val[1], m4b); ++ ++ q2bits.val[0] = vshrq_n_u8(q2bits.val[0], 2); ++ q2bits.val[1] = vshrq_n_u8(q2bits.val[1], 2); ++ b2.val[2] = vandq_u8(q2bits.val[0], m4b); ++ b2.val[3] = vandq_u8(q2bits.val[1], m4b); ++ } ++}; ++ ++template ++static inline void compute_8_blocks(const uint8x16x4_t& qx_1, const uint8x16x4_t& qx_2, const Q8& q8, ++ const int32x4x2_t& scales, int iy, int i, int j, int32x4_t& sumi) { ++ auto mzero = vdupq_n_s32(0); ++ auto q8b_1 = q8.load_quants(iy, i, 4*j+0); ++ auto p1 = ggml_vdotq_s32(ggml_vdotq_s32(mzero, vreinterpretq_s8_u8(qx_1.val[0]), q8b_1.val[0]), ++ vreinterpretq_s8_u8(qx_1.val[1]), q8b_1.val[1]); // block 1 ++ auto q8b_2 = q8.load_quants(iy, i, 4*j+1); ++ auto p2 = ggml_vdotq_s32(ggml_vdotq_s32(mzero, vreinterpretq_s8_u8(qx_1.val[2]), q8b_2.val[0]), ++ vreinterpretq_s8_u8(qx_1.val[3]), q8b_2.val[1]); // block 2 ++ auto p12 = vpaddq_s32(p1, p2); ++ ++ auto q8b_3 = q8.load_quants(iy, i, 4*j+2); ++ auto p3 = ggml_vdotq_s32(ggml_vdotq_s32(mzero, vreinterpretq_s8_u8(qx_2.val[0]), q8b_3.val[0]), ++ vreinterpretq_s8_u8(qx_2.val[1]), q8b_3.val[1]); // block 1 ++ auto q8b_4 = q8.load_quants(iy, i, 4*j+3); ++ auto p4 = ggml_vdotq_s32(ggml_vdotq_s32(mzero, vreinterpretq_s8_u8(qx_2.val[2]), q8b_4.val[0]), ++ vreinterpretq_s8_u8(qx_2.val[3]), q8b_4.val[1]); // block 2 ++ auto p34 = vpaddq_s32(p3, p4); ++ ++ auto pall = vpaddq_s32(p12, p34); ++ sumi = vmlaq_s32(sumi, scales.val[j], pall); ++} ++ ++template ++static inline void compute_16_blocks(const uint8x16x4_t& qx_1, const uint8x16x4_t& qx_2, const Q8& q8, ++ const int32x4x4_t& scales, int iy, int i, int j, int32x4_t& sumi) { ++ ++ auto mzero = vdupq_n_s32(0); ++ auto q8b_1 = q8.load_quants(iy, i, 4*j+0); ++ auto p1 = vpaddq_s32(ggml_vdotq_s32(mzero, vreinterpretq_s8_u8(qx_1.val[0]), q8b_1.val[0]), ++ ggml_vdotq_s32(mzero, vreinterpretq_s8_u8(qx_1.val[1]), q8b_1.val[1])); // blocks 0, 0, 1, 1, ++ auto q8b_2 = q8.load_quants(iy, i, 4*j+1); ++ auto p2 = vpaddq_s32(ggml_vdotq_s32(mzero, vreinterpretq_s8_u8(qx_1.val[2]), q8b_2.val[0]), ++ ggml_vdotq_s32(mzero, vreinterpretq_s8_u8(qx_1.val[3]), q8b_2.val[1])); // blocks 3, 3, 4, 4, ++ auto p12 = vpaddq_s32(p1, p2); // blocks 0, 1, 2, 3 ++ sumi = vmlaq_s32(sumi, scales.val[2*j+0], p12); ++ ++ auto q8b_3 = q8.load_quants(iy, i, 4*j+2); ++ auto p3 = vpaddq_s32(ggml_vdotq_s32(mzero, vreinterpretq_s8_u8(qx_2.val[0]), q8b_3.val[0]), ++ ggml_vdotq_s32(mzero, vreinterpretq_s8_u8(qx_2.val[1]), q8b_3.val[1])); // block 4, 4, 5, 5, ++ auto q8b_4 = q8.load_quants(iy, i, 4*j+3); ++ auto p4 = vpaddq_s32(ggml_vdotq_s32(mzero, vreinterpretq_s8_u8(qx_2.val[2]), q8b_4.val[0]), ++ ggml_vdotq_s32(mzero, vreinterpretq_s8_u8(qx_2.val[3]), q8b_4.val[1])); // block 6, 6, 7, 7, ++ auto p34 = vpaddq_s32(p3, p4); // blocks 4, 5, 6, 7 ++ sumi = vmlaq_s32(sumi, scales.val[2*j+1], p34); ++} ++ ++struct SignHelper { ++ ++ inline void init() { shuffle = vcombine_u8(vdup_n_u8(0), vdup_n_u8(1)); } ++ ++ inline void apply_signs_1(uint8x16_t * b, const uint8x16_t& signs16) { ++ auto aux = vqtbl1q_u8(signs16, shuffle); ++ auto s = vreinterpretq_s8_u8(vorrq_u8(vceqq_u8(vandq_u8(aux, smask), smask), m1)); ++ b[0] = vreinterpretq_u8_s8(vmulq_s8(vreinterpretq_s8_u8(b[0]), s)); ++ shuffle = vaddq_u8(shuffle, step); ++ } ++ ++ const uint8x16_t smask = vreinterpretq_u8_u64(vdupq_n_u64(0x8040201008040201)); ++ const uint8x16_t m1 = vdupq_n_u8(1); ++ const uint8x16_t step = vdupq_n_u8(2); ++ uint8x16_t shuffle; ++}; ++ ++template ++static void mul_mat_qX_K_q8_K_T(int n, const void * vx, size_t bx, const DataInfo& info, int nrc_x) { ++ assert(n % QK_K == 0); ++ const int nb = n / QK_K; ++ ++ Q8 q8(info); ++ ++ Dequantizer deq(vx, bx, nrc_y); ++ ++ for (int ix = 0; ix < nrc_x; ++ix) { ++ ++ deq.new_row(ix); ++ ++ float32x4_t acc[nrc_y]; ++ for (int iy = 0; iy < nrc_y; ++iy) acc[iy] = vdupq_n_f32(0.f); ++ ++ for (int i = 0; i < nb; ++i) { ++ ++ int32x4_t sumi[nrc_y]; ++ for (int iy = 0; iy < nrc_y; ++iy) sumi[iy] = vdupq_n_s32(0); ++ ++ if constexpr (nrc_y > 1 && Dequantizer::should_scale_quants()) { ++ deq.process_scales(i, q8, acc); ++ deq.prepare(i, 0); ++ deq.compute(q8, i, 0, sumi); ++ deq.prepare(i, 1); ++ deq.compute(q8, i, 1, sumi); ++ } else { ++ if constexpr (Dequantizer::num_blocks() == 8) { ++ auto scales = deq.new_block(i, q8, acc); ++ deq.prepare(i, 0); ++ for (int iy = 0; iy < nrc_y; ++iy) compute_8_blocks(deq.bits.b1, deq.bits.b2, q8, scales, iy, i, 0, sumi[iy]); ++ deq.prepare(i, 1); ++ for (int iy = 0; iy < nrc_y; ++iy) compute_8_blocks(deq.bits.b1, deq.bits.b2, q8, scales, iy, i, 1, sumi[iy]); ++ } ++ else if constexpr (Dequantizer::num_blocks() == 16) { ++ auto scales = deq.new_block(i, q8, acc); ++ deq.prepare(i, 0); ++ for (int iy = 0; iy < nrc_y; ++iy) compute_16_blocks(deq.bits.b1, deq.bits.b2, q8, scales, iy, i, 0, sumi[iy]); ++ deq.prepare(i, 1); ++ for (int iy = 0; iy < nrc_y; ++iy) compute_16_blocks(deq.bits.b1, deq.bits.b2, q8, scales, iy, i, 1, sumi[iy]); ++ } ++ else { ++ GGML_ASSERT(false); ++ } ++ } ++ ++ for (int iy = 0; iy < nrc_y; ++iy) { ++ acc[iy] = vmlaq_f32(acc[iy], vcvtq_f32_s32(sumi[iy]), vdupq_n_f32(deq.d*q8.scale(iy, i))); ++ } ++ } ++ ++ for (int iy = 0; iy < nrc_y; ++iy) { ++ info.store(ix, iy, vaddvq_f32(acc[iy])); ++ } ++ } ++} ++ ++#endif ++ ++#endif // IQK_IMPLEMENT +diff --git a/src/ggml-cpu/iqk/iqk_gemm.cpp b/src/ggml-cpu/iqk/iqk_gemm.cpp +new file mode 100644 +index 0000000..1899fd7 +--- /dev/null ++++ b/src/ggml-cpu/iqk/iqk_gemm.cpp +@@ -0,0 +1,258 @@ ++// ++// Copyright (C) 2024-2025 Iwan Kawrakow ++// MIT license ++// SPDX-License-Identifier: MIT ++// ++// Adapted from ik_llama.cpp (https://github.com/ikawrakow/ik_llama.cpp) commit 009dce19, ++// ggml/src/iqk/{iqk_config.h, iqk_common.h, iqk_gemm_kquants.cpp, iqk_gemm_iquants.cpp, iqk_gemm_1bit.cpp, iqk_mul_mat.cpp}. ++// Only the kernels that read the upstream ggml block layouts with upstream Q8_K activations are kept. ++// ++// Permission is hereby granted, free of charge, to any person obtaining a copy ++// of this software and associated documentation files (the "Software"), to deal ++// in the Software without restriction, including without limitation the rights ++// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell ++// copies of the Software, and to permit persons to whom the Software is ++// furnished to do so, subject to the following conditions: ++// ++// The above copyright notice and this permission notice shall be included in all ++// copies or substantial portions of the Software. ++// ++// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR ++// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, ++// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE ++// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER ++// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, ++// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE ++// SOFTWARE. ++// ++// Driver of the iqk GEMM kernels (the tiling of MulMat::mul_mat_NxM in iqk_mul_mat.cpp). ++ ++#include "iqk_gemm.h" ++#include "iqk_common.h" ++ ++#include "ggml.h" ++ ++#include ++#include ++#include ++#include ++ ++namespace { ++ ++// the types this file can serve (whether or not this CPU variant has kernels for them) ++constexpr uint64_t iqk_type_bit(ggml_type t) { return 1ull << (int) t; } ++[[maybe_unused]] constexpr uint64_t k_iqk_all_types = ++ iqk_type_bit(GGML_TYPE_Q2_K) | iqk_type_bit(GGML_TYPE_Q3_K) | iqk_type_bit(GGML_TYPE_Q4_K) | ++ iqk_type_bit(GGML_TYPE_Q5_K) | iqk_type_bit(GGML_TYPE_Q6_K) | iqk_type_bit(GGML_TYPE_IQ4_XS) | ++ iqk_type_bit(GGML_TYPE_IQ2_XXS) | iqk_type_bit(GGML_TYPE_IQ2_XS) | iqk_type_bit(GGML_TYPE_IQ2_S) | ++ iqk_type_bit(GGML_TYPE_IQ3_XXS) | iqk_type_bit(GGML_TYPE_IQ3_S) | iqk_type_bit(GGML_TYPE_IQ1_S) | ++ iqk_type_bit(GGML_TYPE_IQ1_M); ++ ++#ifdef IQK_IMPLEMENT ++constexpr uint64_t k_iqk_built_types = k_iqk_all_types; ++#else ++constexpr uint64_t k_iqk_built_types = 0; ++#endif ++ ++std::atomic g_iqk_types{k_iqk_built_types}; ++std::atomic g_iqk_count[64]; ++int g_iqk_debug = 0; ++bool g_iqk_first = false; ++bool g_iqk_auto = false; ++ ++#ifdef IQK_IMPLEMENT ++ ++bool iqk_set_kernels(int type, mul_mat_t * funcs) { ++ for (int i = 0; i < IQK_MAX_NY; ++i) funcs[i] = nullptr; ++ return iqk_set_kernels_kquants(type, funcs) || iqk_set_kernels_iquants(type, funcs) || iqk_set_kernels_1bit(type, funcs); ++} ++ ++// rows [0, nrc_x) of vx times the activation rows [info.cur_y, nrc_y) ++void iqk_mul_mat_NxM(const mul_mat_t * funcs, int n, const void * vx, size_t bx, DataInfo & info, int nrc_x, int nrc_y) { ++ constexpr int k_x_step = 64; ++ int ny = IQK_MAX_NY; ++ while (ny > 0 && !funcs[ny-1]) --ny; ++ int n_left = nrc_y - info.cur_y; ++ int n_step = n_left/ny; ++ if (n_step > 0) { ++ if (n_step*ny != n_left) { ++ ++n_step; ++ int ny1 = n_left/n_step; ++ int ny2 = ny1 + 1; ++ int my1 = n_step*ny2 - n_left; ++ int my2 = n_step - my1; ++ for (int ix = 0; ix < nrc_x; ix += k_x_step) { ++ auto this_info = info; ++ this_info.s += ix; ++ int this_nrc_x = ix + k_x_step <= nrc_x ? k_x_step : nrc_x - ix; ++ for (int iy = 0; iy < my1; ++iy) { ++ funcs[ny1-1](n, (const void *)((const char *)vx + ix*bx), bx, this_info, this_nrc_x); ++ this_info.cur_y += ny1; ++ } ++ for (int iy = 0; iy < my2; ++iy) { ++ funcs[ny2-1](n, (const void *)((const char *)vx + ix*bx), bx, this_info, this_nrc_x); ++ this_info.cur_y += ny2; ++ } ++ } ++ info.cur_y += n_left; ++ } ++ else { ++ for (int ix = 0; ix < nrc_x; ix += k_x_step) { ++ auto this_info = info; ++ this_info.s += ix; ++ int this_nrc_x = ix + k_x_step <= nrc_x ? k_x_step : nrc_x - ix; ++ for (int iy = 0; iy < n_step; ++iy) { ++ funcs[ny-1](n, (const void *)((const char *)vx + ix*bx), bx, this_info, this_nrc_x); ++ this_info.cur_y += ny; ++ } ++ } ++ info.cur_y += ny * n_step; ++ } ++ } ++ n_left = nrc_y - info.cur_y; ++ if (n_left > 0) { ++ funcs[n_left-1](n, vx, bx, info, nrc_x); ++ } ++} ++ ++#endif // IQK_IMPLEMENT ++ ++} // namespace ++ ++extern "C" { ++ ++void ggml_cpu_iqk_init(void) { ++ uint64_t mask = k_iqk_built_types; ++ const char * on = getenv("GGML_CPU_IQK"); ++#if defined(__x86_64__) || defined(_M_X64) ++ // x86: by default the kernels take only the small batches where they measured faster (ggml_cpu_iqk_mode); the ++ // tiled K-quant path keeps the large ones, where it measured 1.2-5x faster. ++ g_iqk_auto = on == NULL; ++#endif ++ if (on != NULL && atoi(on) == 0) { ++ mask = 0; ++ } ++ g_iqk_first = on != NULL && atoi(on) == 2; ++ const char * env = getenv("GGML_CPU_IQK_TYPES"); ++ if (mask != 0 && env != NULL && strcmp(env, "all") != 0) { ++ uint64_t sel = 0; ++ char buf[256]; ++ snprintf(buf, sizeof(buf), "%s", env); ++ for (char * tok = strtok(buf, ", "); tok != NULL; tok = strtok(NULL, ", ")) { ++ for (int t = 0; t < GGML_TYPE_COUNT && t < 64; ++t) { ++ const char * name = ggml_type_name((enum ggml_type) t); ++ if (name != NULL && strcmp(name, tok) == 0) { ++ sel |= 1ull << t; ++ } ++ } ++ } ++ mask &= sel; ++ } ++ g_iqk_types.store(mask); ++ const char * dbg = getenv("GGML_CPU_IQK_DEBUG"); ++ g_iqk_debug = dbg != NULL ? atoi(dbg) : 0; ++ if (g_iqk_debug) { ++ fprintf(stderr, "ggml-cpu iqk gemm: built types 0x%llx, enabled 0x%llx (%s the tiled path):", (unsigned long long) k_iqk_built_types, ++ (unsigned long long) mask, g_iqk_auto ? "by batch size around" : g_iqk_first ? "ahead of" : "after"); ++ for (int t = 0; t < GGML_TYPE_COUNT && t < 64; ++t) { ++ if ((mask >> t) & 1) fprintf(stderr, " %s", ggml_type_name((enum ggml_type) t)); ++ } ++ fprintf(stderr, "\n"); ++ } ++} ++ ++bool ggml_cpu_iqk_type_enabled(int type) { ++ return type >= 0 && type < 64 && ((g_iqk_types.load(std::memory_order_relaxed) >> type) & 1); ++} ++ ++bool ggml_cpu_iqk_before_tiled(void) { ++ return g_iqk_first; ++} ++ ++// x86 default (GGML_CPU_IQK unset), from test-backend-ops MUL_MAT perf (m = 4096, k = 14336) on Zen 3 (EPYC 7763), ++// Zen 4 (EPYC 9V74), Zen 5 (EPYC 9V45), Xeon 8370C / 8573C and Core i7-8700B, iqk vs the path that runs otherwise: ++// ne11 = 1: vec_dot; iqk was slower for iq1_m / iq2_s / iq2_xxs (0.63-0.89x) and even for the others; ++// ne11 = 2..7: iqk, 1.05-3.2x faster than vec_dot (lowest: q6_K at ne11 = 2 on Zen 4, 0.96x); ++// ne11 = 8: iqk ahead of tiled, 1.02-2.9x, except iq2_xxs / iq3_s (0.84x / 0.91x on Zen 5); ++// ne11 = 9..12: iqk for q2_K..q6_K (1.01-2.3x); the I-quants measured slower than tiled there; ++// ne11 > 12: tiled (iqk 0.16-0.98x from ne11 = 16 up; ne11 = 512: 0.19-0.91x). ++int ggml_cpu_iqk_mode(int type, int64_t ne11) { ++ if (!ggml_cpu_iqk_type_enabled(type)) { ++ return 0; ++ } ++ if (!g_iqk_auto) { ++ return g_iqk_first ? 2 : 1; ++ } ++ if (ne11 < 2) { ++ return 0; ++ } ++ if (ne11 < 8) { ++ return 2; ++ } ++ if (ne11 == 8) { ++ return type == GGML_TYPE_IQ2_XXS || type == GGML_TYPE_IQ3_S ? 0 : 2; ++ } ++ const bool kquant = type == GGML_TYPE_Q2_K || type == GGML_TYPE_Q3_K || type == GGML_TYPE_Q4_K || ++ type == GGML_TYPE_Q5_K || type == GGML_TYPE_Q6_K; ++ return ne11 <= 12 && kquant ? 2 : 0; ++} ++ ++int64_t ggml_cpu_iqk_count(int type) { ++ return type >= 0 && type < 64 ? g_iqk_count[type].load() : 0; ++} ++ ++void ggml_cpu_iqk_set_types(uint64_t mask) { ++ g_iqk_types.store(mask & k_iqk_built_types); ++} ++ ++bool ggml_cpu_iqk_mul_mat(int type, int64_t Nx, int64_t Ny, int64_t ne00, ++ const void * A, size_t stride_A, const void * B, size_t stride_B, ++ float * C, size_t stride_C, int ith, int nth) { ++#ifdef IQK_IMPLEMENT ++ if (ne00 % QK_K != 0 || Nx <= 0 || Ny <= 0 || Nx > INT32_MAX || Ny > INT32_MAX || ne00 > INT32_MAX) { ++ return false; ++ } ++ mul_mat_t funcs[IQK_MAX_NY]; ++ if (!iqk_set_kernels(type, funcs)) { ++ return false; ++ } ++ if (ith == 0) { ++ const int64_t prev = g_iqk_count[type].fetch_add(1, std::memory_order_relaxed); ++ if (g_iqk_debug && prev == 0) { ++ fprintf(stderr, "ggml-cpu iqk gemm: first %s mul_mat (%lld x %lld x %lld, %d threads)\n", ++ ggml_type_name((enum ggml_type) type), (long long) Nx, (long long) Ny, (long long) ne00, nth); ++ } ++ } ++ ++ // Tiles of rows x activation columns: the rows are split over the threads first (each thread then walks all ++ // columns of its row range, 64 rows x up to 8 columns per kernel call); with fewer than 16 rows per thread the ++ // columns are split too. ++ const int nx = (int) Nx; ++ const int ny = (int) Ny; ++ int ntile_x = nth < nx/16 ? nth : (nx/16 > 0 ? nx/16 : 1); ++ int ntile_y = nth/ntile_x; ++ if (ntile_y > (ny + 7)/8) ntile_y = (ny + 7)/8; ++ if (ntile_y < 1) ntile_y = 1; ++ const int ntile = ntile_x*ntile_y; ++ for (int itile = ith; itile < ntile; itile += nth) { ++ const int jx = itile % ntile_x; ++ const int jy = itile / ntile_x; ++ const int ix0 = (int)((int64_t) jx * nx / ntile_x); ++ const int ix1 = (int)((int64_t)(jx + 1) * nx / ntile_x); ++ const int iy0 = (int)((int64_t) jy * ny / ntile_y); ++ const int iy1 = (int)((int64_t)(jy + 1) * ny / ntile_y); ++ if (ix1 <= ix0 || iy1 <= iy0) { ++ continue; ++ } ++ DataInfo info{C + ix0, (const char *) B, stride_C, stride_B, iy0}; ++ iqk_mul_mat_NxM(funcs, (int) ne00, (const char *) A + ix0*stride_A, stride_A, info, ix1 - ix0, iy1); ++ } ++ return true; ++#else ++ (void) type; (void) Nx; (void) Ny; (void) ne00; (void) A; (void) stride_A; (void) B; (void) stride_B; ++ (void) C; (void) stride_C; (void) ith; (void) nth; ++ return false; ++#endif ++} ++ ++} // extern "C" +diff --git a/src/ggml-cpu/iqk/iqk_gemm.h b/src/ggml-cpu/iqk/iqk_gemm.h +new file mode 100644 +index 0000000..a36b5d7 +--- /dev/null ++++ b/src/ggml-cpu/iqk/iqk_gemm.h +@@ -0,0 +1,45 @@ ++#pragma once ++ ++// GEMM kernels ported from ik_llama.cpp (iqk_gemm.cpp) for the k-quant / i-quant weight types whose ++// activations are Q8_K. They read the upstream block layouts and the upstream Q8_K activation rows and ++// form the same integer products as the ggml vec_dot kernels; only the float accumulation order differs. ++ ++#include ++#include ++#include ++ ++#ifdef __cplusplus ++extern "C" { ++#endif ++ ++// Reads once: GGML_CPU_IQK (0 disables; 1 enables, the default on ARM, while x86 defaults to small batches only; 2 also runs them ++// ahead of the x86 tiled K-quant path), ++// GGML_CPU_IQK_TYPES (ggml type names or "all"; the other types keep the tiled / vec_dot paths) and GGML_CPU_IQK_DEBUG ++// (1: print the enabled types and the first mul_mat per type to stderr). ++void ggml_cpu_iqk_init(void); ++ ++// true if this build has a kernel for the weight type and it is enabled ++bool ggml_cpu_iqk_type_enabled(int type); ++ ++// GGML_CPU_IQK=2: the iqk kernels take the op before the tiled path gets a chance ++bool ggml_cpu_iqk_before_tiled(void); ++ ++// For a weight type and src1 row count: 0 = not iqk, 1 = iqk after the tiled path, 2 = iqk ahead of it. ++// GGML_CPU_IQK=1 / 2 give 1 / 2 for every enabled type; unset on x86 picks by batch size (see iqk_gemm.cpp). ++int ggml_cpu_iqk_mode(int type, int64_t ne11); ++ ++// C[iy*stride_C + ix] = dot(A row ix, B row iy) for ix < Nx, iy < Ny; A rows of type `type` (stride_A bytes), ++// B rows are Q8_K (stride_B bytes), ne00 % 256 == 0. The work is split over the nth threads by ith. ++// Returns false (computing nothing) if the type is not handled. ++bool ggml_cpu_iqk_mul_mat(int type, int64_t Nx, int64_t Ny, int64_t ne00, ++ const void * A, size_t stride_A, const void * B, size_t stride_B, ++ float * C, size_t stride_C, int ith, int nth); ++ ++// test hooks (reached through the CPU backend's get_proc_address): number of mul_mat ops computed by ++// these kernels per weight type, and an override of the enabled type mask (bit = ggml type) ++int64_t ggml_cpu_iqk_count(int type); ++void ggml_cpu_iqk_set_types(uint64_t mask); ++ ++#ifdef __cplusplus ++} ++#endif +diff --git a/src/ggml-cpu/iqk/iqk_gemm_1bit.cpp b/src/ggml-cpu/iqk/iqk_gemm_1bit.cpp +new file mode 100644 +index 0000000..43ba745 +--- /dev/null ++++ b/src/ggml-cpu/iqk/iqk_gemm_1bit.cpp +@@ -0,0 +1,364 @@ ++// ++// Copyright (C) 2024-2025 Iwan Kawrakow ++// MIT license ++// SPDX-License-Identifier: MIT ++// ++// Adapted from ik_llama.cpp (https://github.com/ikawrakow/ik_llama.cpp) commit 009dce19, ++// ggml/src/iqk/{iqk_config.h, iqk_common.h, iqk_gemm_kquants.cpp, iqk_gemm_iquants.cpp, iqk_gemm_1bit.cpp, iqk_mul_mat.cpp}. ++// Only the kernels that read the upstream ggml block layouts with upstream Q8_K activations are kept. ++// ++// Permission is hereby granted, free of charge, to any person obtaining a copy ++// of this software and associated documentation files (the "Software"), to deal ++// in the Software without restriction, including without limitation the rights ++// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell ++// copies of the Software, and to permit persons to whom the Software is ++// furnished to do so, subject to the following conditions: ++// ++// The above copyright notice and this permission notice shall be included in all ++// copies or substantial portions of the Software. ++// ++// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR ++// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, ++// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE ++// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER ++// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, ++// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE ++// SOFTWARE. ++// ++// IQ1_S and IQ1_M times Q8_K (from iqk_gemm_1bit.cpp). ++ ++#include "iqk_common.h" ++ ++#ifdef IQK_IMPLEMENT ++ ++namespace { ++ ++// ik_llama.cpp ships this table as a literal: it is iq1s_grid with 1 added to every int8 lane, so the ++// lanes hold 0, 1, 2 instead of -1, 0, 1. It is derived here from the iq1s_grid of ggml-common.h. ++// Filled on first use, not by a static initializer: dlopen runs those in every CPU variant library while ggml ++// scores them, so an AVX512-vectorized initializer would SIGILL on CPUs without AVX512. ++uint64_t iq1s_grid_us[NGRID_IQ1S]; ++bool iq1s_grid_us_init() { ++ static const bool ready = [] { ++ for (int i = 0; i < NGRID_IQ1S; ++i) { ++ uint64_t v = 0; ++ for (int k = 0; k < 8; ++k) { ++ v |= (uint64_t)(uint8_t)((int8_t)(iq1s_grid[i] >> 8*k) + 1) << 8*k; ++ } ++ iq1s_grid_us[i] = v; ++ } ++ return true; ++ }(); ++ return ready; ++} ++ ++} ++ ++#ifdef __x86_64__ ++ ++namespace { ++ ++template ++static inline void prefetch_row_ahead(const Block * x) { ++ _mm_prefetch((const char *)x + 2048, _MM_HINT_T2); ++} ++ ++template ++void mul_mat_iq1_s_q8_K(int n, const void * vx, size_t bx, const DataInfo& info, int nrc_x) { ++ GGML_ASSERT(n%QK_K == 0); ++ int ix0 = 0; ++ Q8 q8(info); ++ __m256i qx[8]; ++ __m256i scales[4]; ++ __m256 acc[nrc_y] = {}; ++ auto delta_mask = _mm_set1_epi16(-32768); // to avoid stupid overflow warnings when using 0x8000 ++ __m256i shuffle0 = _mm256_set_epi64x(0x0302030203020302, 0x0100010001000100, 0x0302030203020302, 0x0100010001000100); ++ for (int ix = ix0; ix < nrc_x; ++ix) { ++ auto iq1s = (const block_iq1_s *)((const char *)vx + ix*bx); ++ for (int ibl = 0; ibl < n/QK_K; ++ibl) { ++ prefetch_row_ahead(iq1s + ibl); ++ float d = GGML_FP16_TO_FP32(iq1s[ibl].d); ++ auto qhb = _mm_loadu_si128((const __m128i *)iq1s[ibl].qh); ++ auto scales128 = _mm_and_si128(_mm_srli_epi16(qhb, 12), _mm_set1_epi16(7)); ++ scales128 = _mm_add_epi16(_mm_slli_epi16(scales128, 1), _mm_set1_epi16(1)); ++#ifdef HAVE_FANCY_SIMD ++ auto mask = _mm_cmpeq_epi16_mask(_mm_and_si128(qhb, delta_mask), delta_mask); ++ auto deltas128 = _mm_mask_blend_epi16(mask, _mm_set1_epi16(-7), _mm_set1_epi16(-9)); ++#else ++ auto mask = _mm_cmpeq_epi16(_mm_and_si128(qhb, delta_mask), delta_mask); ++ auto deltas128 = _mm_or_si128(_mm_and_si128(mask, _mm_set1_epi16(-9)), _mm_andnot_si128(mask, _mm_set1_epi16(-7))); ++#endif ++ deltas128 = _mm_mullo_epi16(scales128, deltas128); ++ scales128 = _mm_slli_epi16(scales128, 3); ++ auto deltas_l = _mm_unpacklo_epi16(deltas128, deltas128); ++ auto deltas_h = _mm_unpackhi_epi16(deltas128, deltas128); ++ auto deltas = MM256_SET_M128I(deltas_h, deltas_l); // blocks 0,0, 1,1, 2,2, ..., 7,7 ++ auto all_scales = MM256_SET1_M128I(scales128); ++ auto shuffle = shuffle0; ++ for (int ib64 = 0; ib64 < QK_K/64; ++ib64) { ++ scales[ib64] = _mm256_shuffle_epi8(all_scales, shuffle); ++ shuffle = _mm256_add_epi8(shuffle, _mm256_set1_epi8(4)); ++ } ++ const uint8_t * qs = iq1s[ibl].qs; ++ const uint16_t * qh = iq1s[ibl].qh; ++ for (int ib = 0; ib < QK_K/32; ib += 2) { ++ qx[ib+0] = _mm256_set_epi64x(iq1s_grid_us[qs[3] | ((qh[ib+0] >> 1) & 0x700)], iq1s_grid_us[qs[2] | ((qh[ib+0] << 2) & 0x700)], ++ iq1s_grid_us[qs[1] | ((qh[ib+0] << 5) & 0x700)], iq1s_grid_us[qs[0] | ((qh[ib+0] << 8) & 0x700)]); ++ qx[ib+1] = _mm256_set_epi64x(iq1s_grid_us[qs[7] | ((qh[ib+1] >> 1) & 0x700)], iq1s_grid_us[qs[6] | ((qh[ib+1] << 2) & 0x700)], ++ iq1s_grid_us[qs[5] | ((qh[ib+1] << 5) & 0x700)], iq1s_grid_us[qs[4] | ((qh[ib+1] << 8) & 0x700)]); ++ qs += 8; ++ } ++ for (int iy = 0; iy < nrc_y; ++iy) { ++ auto bsums = q8.load_bsums(iy, ibl); ++ auto sumi = _mm256_setzero_si256(); ++ for (int ib64 = 0; ib64 < QK_K/64; ++ib64) { ++ auto qy1 = q8.load_quants(iy, ibl, 2*ib64+0); ++ auto qy2 = q8.load_quants(iy, ibl, 2*ib64+1); ++#ifdef HAVE_FANCY_SIMD ++ auto dot1 = _mm256_dpbusd_epi32(_mm256_setzero_si256(), qx[2*ib64+0], qy1); ++ auto dot2 = _mm256_dpbusd_epi32(_mm256_setzero_si256(), qx[2*ib64+1], qy2); ++ sumi = _mm256_dpwssd_epi32(sumi, scales[ib64], _mm256_packs_epi32(dot1, dot2)); ++#else ++ auto dot1 = _mm256_maddubs_epi16(qx[2*ib64+0], qy1); ++ auto dot2 = _mm256_maddubs_epi16(qx[2*ib64+1], qy2); ++ auto dot = _mm256_add_epi16(_mm256_unpacklo_epi64(dot1, dot2), _mm256_unpackhi_epi64(dot1, dot2)); ++ sumi = _mm256_add_epi32(sumi, _mm256_madd_epi16(scales[ib64], dot)); ++#endif ++ } ++#ifdef HAVE_FANCY_SIMD ++ sumi = _mm256_dpwssd_epi32(sumi, bsums, deltas); ++#else ++ sumi = _mm256_add_epi32(sumi, _mm256_madd_epi16(bsums, deltas)); ++#endif ++ acc[iy] = _mm256_fmadd_ps(_mm256_set1_ps(d*q8.scale(iy, ibl)), _mm256_cvtepi32_ps(sumi), acc[iy]); ++ } ++ } ++ for (int iy = 0; iy < nrc_y; ++iy) { ++ info.store(ix, iy, 0.125f*hsum_float_8(acc[iy])); ++ acc[iy] = _mm256_setzero_ps(); ++ } ++ } ++} ++ ++template ++void mul_mat_iq1_m_q8_K(int n, const void * vx, size_t bx, const DataInfo& info, int nrc_x) { ++ GGML_ASSERT(n%QK_K == 0); ++ int ix0 = 0; ++ Q8 q8(info); ++ __m256i qx[8]; ++ __m256 acc[nrc_y] = {}; ++ auto scale_shuffle = _mm256_set_epi64x(0x0706070607060706, 0x0504050405040504, 0x0302030203020302, 0x0100010001000100); ++ auto delta_mask = _mm256_set_epi64x(0x8000, 0x0800, 0x0080, 0x0008); ++ iq1m_scale_t scale; ++ union { __m256i vec; int16_t val[16]; } helper; ++ for (int ix = ix0; ix < nrc_x; ++ix) { ++ auto iq1m = (const block_iq1_m *)((const char *)vx + ix*bx); ++ for (int ibl = 0; ibl < n/QK_K; ++ibl) { ++ prefetch_row_ahead(iq1m + ibl); ++ const uint16_t * sc = (const uint16_t *)iq1m[ibl].scales; // 4 x uint16_t, each containing 4 scales ++ scale.u16 = (sc[0] >> 12) | ((sc[1] >> 8) & 0x00f0) | ((sc[2] >> 4) & 0x0f00) | (sc[3] & 0xf000); ++ float d = GGML_FP16_TO_FP32(scale.f16); ++ auto qs = iq1m[ibl].qs; ++ auto qh = iq1m[ibl].qh; ++ auto aux = _mm_loadl_epi64((const __m128i *)iq1m[ibl].scales); ++ auto sc16 = _mm256_shuffle_epi8(MM256_SET1_M128I(aux), scale_shuffle); ++ sc16 = _mm256_and_si256(sc16, _mm256_set1_epi64x(0x0e0001c000380007)); ++ sc16 = _mm256_mullo_epi16(sc16, _mm256_set1_epi64x(0x0001000800400200)); ++ helper.vec = _mm256_add_epi8(_mm256_srli_epi16(sc16, 8), _mm256_set1_epi16(1)); ++ for (int ib64 = 0; ib64 < QK_K/64; ++ib64) { ++ qx[2*ib64+0] = _mm256_set_epi64x(iq1s_grid_us[qs[3] | (((uint16_t)qh[1] << 4) & 0x700)], iq1s_grid_us[qs[2] | (((uint16_t)qh[1] << 8) & 0x700)], ++ iq1s_grid_us[qs[1] | (((uint16_t)qh[0] << 4) & 0x700)], iq1s_grid_us[qs[0] | (((uint16_t)qh[0] << 8) & 0x700)]); ++ qx[2*ib64+1] = _mm256_set_epi64x(iq1s_grid_us[qs[7] | (((uint16_t)qh[3] << 4) & 0x700)], iq1s_grid_us[qs[6] | (((uint16_t)qh[3] << 8) & 0x700)], ++ iq1s_grid_us[qs[5] | (((uint16_t)qh[2] << 4) & 0x700)], iq1s_grid_us[qs[4] | (((uint16_t)qh[2] << 8) & 0x700)]); ++ //auto delta1 = _mm256_set_epi64x(qh[1] & 0x80 ? 0x0909090909090909 : 0x0707070707070707, ++ // qh[1] & 0x08 ? 0x0909090909090909 : 0x0707070707070707, ++ // qh[0] & 0x80 ? 0x0909090909090909 : 0x0707070707070707, ++ // qh[0] & 0x08 ? 0x0909090909090909 : 0x0707070707070707); ++ //auto delta2 = _mm256_set_epi64x(qh[3] & 0x80 ? 0x0909090909090909 : 0x0707070707070707, ++ // qh[3] & 0x08 ? 0x0909090909090909 : 0x0707070707070707, ++ // qh[2] & 0x80 ? 0x0909090909090909 : 0x0707070707070707, ++ // qh[2] & 0x08 ? 0x0909090909090909 : 0x0707070707070707); ++ auto qh16 = (const uint16_t *)qh; ++ auto delta1 = _mm256_cmpeq_epi64(_mm256_and_si256(_mm256_set1_epi64x(qh16[0]), delta_mask), delta_mask); ++ auto delta2 = _mm256_cmpeq_epi64(_mm256_and_si256(_mm256_set1_epi64x(qh16[1]), delta_mask), delta_mask); ++ delta1 = _mm256_sub_epi8(_mm256_set1_epi8(8), _mm256_or_si256(delta1, _mm256_set1_epi8(1))); ++ delta2 = _mm256_sub_epi8(_mm256_set1_epi8(8), _mm256_or_si256(delta2, _mm256_set1_epi8(1))); ++ qx[2*ib64+0] = _mm256_sub_epi8(_mm256_slli_epi16(qx[2*ib64+0], 3), delta1); ++ qx[2*ib64+1] = _mm256_sub_epi8(_mm256_slli_epi16(qx[2*ib64+1], 3), delta2); ++ qs += 8; ++ qh += 4; ++ } ++ for (int iy = 0; iy < nrc_y; ++iy) { ++ auto sumi = _mm256_setzero_si256(); ++ for (int j = 0; j < 8; ++j) { ++ auto p = _mm256_maddubs_epi16(_mm256_sign_epi8(qx[j], qx[j]), _mm256_sign_epi8(q8.load_quants(iy, ibl, j), qx[j])); ++ sumi = _mm256_add_epi32(sumi, _mm256_madd_epi16(p, MM256_SET_M128I(_mm_set1_epi16(helper.val[2*j+1]), _mm_set1_epi16(helper.val[2*j+0])))); ++ } ++ acc[iy] = _mm256_fmadd_ps(_mm256_set1_ps(d*q8.scale(iy, ibl)), _mm256_cvtepi32_ps(sumi), acc[iy]); ++ } ++ } ++ for (int iy = 0; iy < nrc_y; ++iy) { ++ info.store(ix, iy, 0.125f*hsum_float_8(acc[iy])); ++ acc[iy] = _mm256_setzero_ps(); ++ } ++ } ++} ++ ++} // namespace ++ ++#else ++// --------------------------------- __aarch64__ -------------------------------------- ++ ++namespace { ++ ++template ++void mul_mat_iq1_s_q8_K(int n, const void * vx, size_t bx, const DataInfo& info, int nrc_x) { ++ GGML_ASSERT(n%QK_K == 0); ++ Q8 q8(info); ++ int8x16_t qx[16]; ++ int32x4_t scales[2]; ++ int16x4_t deltas[2]; ++ float32x4_t acc[nrc_y] = {}; ++ auto delta_mask = vdupq_n_u16(0x8000); ++ for (int ix = 0; ix < nrc_x; ++ix) { ++ auto iq1s = (const block_iq1_s *)((const char *)vx + ix*bx); ++ for (int ibl = 0; ibl < n/QK_K; ++ibl) { ++ float d = GGML_FP16_TO_FP32(iq1s[ibl].d); ++ auto qhb = vld1q_u16(iq1s[ibl].qh); ++ auto scales128 = vandq_u16(vshrq_n_u16(qhb, 12), vdupq_n_u16(7)); ++ scales128 = vaddq_u16(vshlq_n_u16(scales128, 1), vdupq_n_u16(1)); ++ auto mask = vceqq_u16(vandq_u16(qhb, delta_mask), delta_mask); ++ // Note: we explicitly assume IQ1S_DELTA = 0.125 ++ auto deltas128 = vsubq_s16(vreinterpretq_s16_u16(vbicq_u16(scales128, mask)), vreinterpretq_s16_u16(vandq_u16(scales128, mask))); ++ //auto deltas128 = vorrq_s16(vandq_s16(vdupq_n_s16(-1), mask), vbicq_s16(vdupq_n_s16(1), mask)); ++ //deltas128 = vmulq_s16(scales128, deltas128); ++ scales128 = vshlq_n_u16(scales128, 3); ++ auto qs = iq1s[ibl].qs; ++ auto qh = iq1s[ibl].qh; ++ for (int ib64 = 0; ib64 < QK_K/64; ++ib64) { ++ qx[4*ib64+0] = vreinterpretq_s8_u64(uint64x2_t{iq1s_grid[qs[0] | ((qh[2*ib64+0] << 8) & 0x700)], iq1s_grid[qs[1] | ((qh[2*ib64+0] << 5) & 0x700)]}); ++ qx[4*ib64+1] = vreinterpretq_s8_u64(uint64x2_t{iq1s_grid[qs[2] | ((qh[2*ib64+0] << 2) & 0x700)], iq1s_grid[qs[3] | ((qh[2*ib64+0] >> 1) & 0x700)]}); ++ qx[4*ib64+2] = vreinterpretq_s8_u64(uint64x2_t{iq1s_grid[qs[4] | ((qh[2*ib64+1] << 8) & 0x700)], iq1s_grid[qs[5] | ((qh[2*ib64+1] << 5) & 0x700)]}); ++ qx[4*ib64+3] = vreinterpretq_s8_u64(uint64x2_t{iq1s_grid[qs[6] | ((qh[2*ib64+1] << 2) & 0x700)], iq1s_grid[qs[7] | ((qh[2*ib64+1] >> 1) & 0x700)]}); ++ qs += 8; ++ } ++ scales[0] = vreinterpretq_s32_u32(vmovl_u16(vget_low_u16 (scales128))); ++ scales[1] = vreinterpretq_s32_u32(vmovl_u16(vget_high_u16(scales128))); ++ deltas[0] = vget_low_s16 (deltas128); ++ deltas[1] = vget_high_s16(deltas128); ++ for (int iy = 0; iy < nrc_y; ++iy) { ++ auto bsums = q8.load_bsums8(iy, ibl); ++ auto sumi = vdupq_n_s32(0); ++ sumi = vmlal_s16(sumi, deltas[0], vget_low_s16 (bsums)); ++ sumi = vmlal_s16(sumi, deltas[1], vget_high_s16(bsums)); ++ for (int k = 0; k < QK_K/128; ++k) { ++ auto qy = q8.load_quants_64(iy, ibl, 2*k+0); ++ auto dot1 = ggml_vdotq_s32(ggml_vdotq_s32(vdupq_n_s32(0), qx[8*k+0], qy.val[0]), qx[8*k+1], qy.val[1]); ++ auto dot2 = ggml_vdotq_s32(ggml_vdotq_s32(vdupq_n_s32(0), qx[8*k+2], qy.val[2]), qx[8*k+3], qy.val[3]); ++ auto dot12 = vpaddq_s32(dot1, dot2); ++ qy = q8.load_quants_64(iy, ibl, 2*k+1); ++ auto dot3 = ggml_vdotq_s32(ggml_vdotq_s32(vdupq_n_s32(0), qx[8*k+4], qy.val[0]), qx[8*k+5], qy.val[1]); ++ auto dot4 = ggml_vdotq_s32(ggml_vdotq_s32(vdupq_n_s32(0), qx[8*k+6], qy.val[2]), qx[8*k+7], qy.val[3]); ++ auto dot34 = vpaddq_s32(dot3, dot4); ++ auto dot = vpaddq_s32(dot12, dot34); ++ sumi = vmlaq_s32(sumi, dot, scales[k]); ++ } ++ acc[iy] = vfmaq_f32(acc[iy], vdupq_n_f32(d*q8.scale(iy, ibl)), vcvtq_f32_s32(sumi)); ++ } ++ } ++ for (int iy = 0; iy < nrc_y; ++iy) { ++ info.store(ix, iy, 0.125f*vaddvq_f32(acc[iy])); ++ acc[iy] = vdupq_n_f32(0); ++ } ++ } ++} ++ ++template ++void mul_mat_iq1_m_q8_K(int n, const void * vx, size_t bx, const DataInfo& info, int nrc_x) { ++ GGML_ASSERT(n%QK_K == 0); ++ Q8 q8(info); ++ int8x16x2_t qx[8]; ++ int32x4x4_t scales; ++ float32x4_t acc[nrc_y] = {}; ++ uint8x16x2_t scale_shuffle = {vreinterpretq_u8_u64(uint64x2_t{0x0100010001000100, 0x0302030203020302}), ++ vreinterpretq_u8_u64(uint64x2_t{0x0504050405040504, 0x0706070607060706})}; ++ uint64x2x2_t delta_mask = {uint64x2_t{0x0008, 0x0080}, uint64x2_t{0x0800, 0x8000}}; ++ iq1m_scale_t block_scale; ++ for (int ix = 0; ix < nrc_x; ++ix) { ++ auto iq1m = (const block_iq1_m *)((const char *)vx + ix*bx); ++ for (int ibl = 0; ibl < n/QK_K; ++ibl) { ++ const uint16_t * sc = (const uint16_t *)iq1m[ibl].scales; // 4 x uint16_t, each containing 4 scales ++ block_scale.u16 = (sc[0] >> 12) | ((sc[1] >> 8) & 0x00f0) | ((sc[2] >> 4) & 0x0f00) | (sc[3] & 0xf000); ++ float d = GGML_FP16_TO_FP32(block_scale.f16); ++ auto qs = iq1m[ibl].qs; ++ auto qh = iq1m[ibl].qh; ++ auto aux8 = vld1_u8(iq1m[ibl].scales); ++ auto aux16 = vcombine_u8(aux8, aux8); ++ uint16x8x2_t sc16 = { vreinterpretq_u16_u8(vqtbl1q_u8(aux16, scale_shuffle.val[0])), vreinterpretq_u16_u8(vqtbl1q_u8(aux16, scale_shuffle.val[1])) }; ++ sc16.val[0] = vmulq_u16(vandq_u16(sc16.val[0], vreinterpretq_u16_u64(vdupq_n_u64(0x0e0001c000380007))), vreinterpretq_u16_u64(vdupq_n_u64(0x0001000800400200))); ++ sc16.val[1] = vmulq_u16(vandq_u16(sc16.val[1], vreinterpretq_u16_u64(vdupq_n_u64(0x0e0001c000380007))), vreinterpretq_u16_u64(vdupq_n_u64(0x0001000800400200))); ++ sc16.val[0] = vaddq_u16(vshrq_n_u16(sc16.val[0], 8), vdupq_n_u16(1)); ++ sc16.val[1] = vaddq_u16(vshrq_n_u16(sc16.val[1], 8), vdupq_n_u16(1)); ++ scales.val[0] = vmovl_s16(vget_low_s16 (vreinterpretq_s16_u16(sc16.val[0]))); ++ scales.val[1] = vmovl_s16(vget_high_s16(vreinterpretq_s16_u16(sc16.val[0]))); ++ scales.val[2] = vmovl_s16(vget_low_s16 (vreinterpretq_s16_u16(sc16.val[1]))); ++ scales.val[3] = vmovl_s16(vget_high_s16(vreinterpretq_s16_u16(sc16.val[1]))); ++ for (int ib64 = 0; ib64 < QK_K/64; ++ib64) { ++ qx[2*ib64+0] = {vreinterpretq_s8_u64(uint64x2_t{iq1s_grid_us[qs[0] | (((uint16_t)qh[0] << 8) & 0x700)], ++ iq1s_grid_us[qs[1] | (((uint16_t)qh[0] << 4) & 0x700)]}), ++ vreinterpretq_s8_u64(uint64x2_t{iq1s_grid_us[qs[2] | (((uint16_t)qh[1] << 8) & 0x700)], ++ iq1s_grid_us[qs[3] | (((uint16_t)qh[1] << 4) & 0x700)]})}; ++ qx[2*ib64+1] = {vreinterpretq_s8_u64(uint64x2_t{iq1s_grid_us[qs[4] | (((uint16_t)qh[2] << 8) & 0x700)], ++ iq1s_grid_us[qs[5] | (((uint16_t)qh[2] << 4) & 0x700)]}), ++ vreinterpretq_s8_u64(uint64x2_t{iq1s_grid_us[qs[6] | (((uint16_t)qh[3] << 8) & 0x700)], ++ iq1s_grid_us[qs[7] | (((uint16_t)qh[3] << 4) & 0x700)]})}; ++ auto qh16 = (const uint16_t *)qh; ++ auto h1 = vdupq_n_u64(qh16[0]); ++ auto h2 = vdupq_n_u64(qh16[1]); ++ auto delta1 = vsubq_s8(vdupq_n_s8(8), vorrq_s8(vreinterpretq_s8_u64(vceqq_u64(vandq_u64(h1, delta_mask.val[0]), delta_mask.val[0])), vdupq_n_s8(1))); ++ auto delta2 = vsubq_s8(vdupq_n_s8(8), vorrq_s8(vreinterpretq_s8_u64(vceqq_u64(vandq_u64(h1, delta_mask.val[1]), delta_mask.val[1])), vdupq_n_s8(1))); ++ qx[2*ib64+0].val[0] = vsubq_s8(vshlq_n_s8(qx[2*ib64+0].val[0], 3), delta1); ++ qx[2*ib64+0].val[1] = vsubq_s8(vshlq_n_s8(qx[2*ib64+0].val[1], 3), delta2); ++ delta1 = vsubq_s8(vdupq_n_s8(8), vorrq_s8(vreinterpretq_s8_u64(vceqq_u64(vandq_u64(h2, delta_mask.val[0]), delta_mask.val[0])), vdupq_n_s8(1))); ++ delta2 = vsubq_s8(vdupq_n_s8(8), vorrq_s8(vreinterpretq_s8_u64(vceqq_u64(vandq_u64(h2, delta_mask.val[1]), delta_mask.val[1])), vdupq_n_s8(1))); ++ qx[2*ib64+1].val[0] = vsubq_s8(vshlq_n_s8(qx[2*ib64+1].val[0], 3), delta1); ++ qx[2*ib64+1].val[1] = vsubq_s8(vshlq_n_s8(qx[2*ib64+1].val[1], 3), delta2); ++ qs += 8; ++ qh += 4; ++ } ++ for (int iy = 0; iy < nrc_y; ++iy) { ++ auto sumi = vdupq_n_s32(0); ++ for (int j = 0; j < 4; ++j) { ++ auto y1 = q8.load_quants(iy, ibl, 2*j+0); ++ auto dot1 = ggml_vdotq_s32(vdupq_n_s32(0), qx[2*j+0].val[0], y1.val[0]); ++ auto dot2 = ggml_vdotq_s32(vdupq_n_s32(0), qx[2*j+0].val[1], y1.val[1]); ++ auto y2 = q8.load_quants(iy, ibl, 2*j+1); ++ auto dot3 = ggml_vdotq_s32(vdupq_n_s32(0), qx[2*j+1].val[0], y2.val[0]); ++ auto dot4 = ggml_vdotq_s32(vdupq_n_s32(0), qx[2*j+1].val[1], y2.val[1]); ++ auto dot = vpaddq_s32(vpaddq_s32(dot1, dot2), vpaddq_s32(dot3, dot4)); ++ sumi = vmlaq_s32(sumi, dot, scales.val[j]); ++ } ++ acc[iy] = vfmaq_f32(acc[iy], vdupq_n_f32(d*q8.scale(iy, ibl)), vcvtq_f32_s32(sumi)); ++ } ++ } ++ for (int iy = 0; iy < nrc_y; ++iy) { ++ info.store(ix, iy, 0.125f*vaddvq_f32(acc[iy])); ++ acc[iy] = vdupq_n_f32(0.f); ++ } ++ } ++} ++ ++} // namespace ++ ++#endif ++ ++bool iqk_set_kernels_1bit(int typeA, mul_mat_t * funcs) { ++ switch (typeA) { ++ case GGML_TYPE_IQ1_S: IQK_SET_MUL_MAT_FUNCTIONS(mul_mat_iq1_s_q8_K, funcs); break; ++ case GGML_TYPE_IQ1_M: IQK_SET_MUL_MAT_FUNCTIONS(mul_mat_iq1_m_q8_K, funcs); break; ++ default: return false; ++ } ++ return iq1s_grid_us_init(); ++} ++ ++#endif // IQK_IMPLEMENT +diff --git a/src/ggml-cpu/iqk/iqk_gemm_iquants.cpp b/src/ggml-cpu/iqk/iqk_gemm_iquants.cpp +new file mode 100644 +index 0000000..f916fee +--- /dev/null ++++ b/src/ggml-cpu/iqk/iqk_gemm_iquants.cpp +@@ -0,0 +1,1169 @@ ++// ++// Copyright (C) 2024-2025 Iwan Kawrakow ++// MIT license ++// SPDX-License-Identifier: MIT ++// ++// Adapted from ik_llama.cpp (https://github.com/ikawrakow/ik_llama.cpp) commit 009dce19, ++// ggml/src/iqk/{iqk_config.h, iqk_common.h, iqk_gemm_kquants.cpp, iqk_gemm_iquants.cpp, iqk_gemm_1bit.cpp, iqk_mul_mat.cpp}. ++// Only the kernels that read the upstream ggml block layouts with upstream Q8_K activations are kept. ++// ++// Permission is hereby granted, free of charge, to any person obtaining a copy ++// of this software and associated documentation files (the "Software"), to deal ++// in the Software without restriction, including without limitation the rights ++// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell ++// copies of the Software, and to permit persons to whom the Software is ++// furnished to do so, subject to the following conditions: ++// ++// The above copyright notice and this permission notice shall be included in all ++// copies or substantial portions of the Software. ++// ++// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR ++// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, ++// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE ++// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER ++// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, ++// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE ++// SOFTWARE. ++// ++// IQ2_XXS, IQ2_XS, IQ2_S, IQ3_XXS and IQ3_S times Q8_K (from iqk_gemm_iquants.cpp). ++ ++#include "iqk_common.h" ++ ++#ifdef IQK_IMPLEMENT ++ ++#ifdef __x86_64__ ++ ++namespace { ++ ++inline __m256i get_scale_shuffle_8(int i) { ++ return _mm256_set1_epi16((2*i) | ((2*i+1) << 8)); ++} ++ ++inline void set_scales_8(const __m256i& all_scales, int j, __m256i * scales) { ++ scales[0] = _mm256_shuffle_epi8(all_scales, get_scale_shuffle_8(4*j+0)); ++ scales[1] = _mm256_shuffle_epi8(all_scales, get_scale_shuffle_8(4*j+1)); ++ scales[2] = _mm256_shuffle_epi8(all_scales, get_scale_shuffle_8(4*j+2)); ++ scales[3] = _mm256_shuffle_epi8(all_scales, get_scale_shuffle_8(4*j+3)); ++} ++ ++inline __m256i get_scale_shuffle_16(int i) { ++ static const uint8_t k_shuffle[128] = { ++ 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 2, 3, 2, 3, 2, 3, 2, 3, 2, 3, 2, 3, 2, 3, 2, 3, ++ 4, 5, 4, 5, 4, 5, 4, 5, 4, 5, 4, 5, 4, 5, 4, 5, 6, 7, 6, 7, 6, 7, 6, 7, 6, 7, 6, 7, 6, 7, 6, 7, ++ 8, 9, 8, 9, 8, 9, 8, 9, 8, 9, 8, 9, 8, 9, 8, 9, 10,11,10,11,10,11,10,11,10,11,10,11,10,11,10,11, ++ 12,13,12,13,12,13,12,13,12,13,12,13,12,13,12,13, 14,15,14,15,14,15,14,15,14,15,14,15,14,15,14,15, ++ }; ++ return _mm256_loadu_si256((const __m256i*)k_shuffle + i); ++} ++ ++inline void set_scales_16(const __m256i& all_scales, __m256i * scales) { ++ scales[0] = _mm256_shuffle_epi8(all_scales, get_scale_shuffle_16(0)); ++ scales[1] = _mm256_shuffle_epi8(all_scales, get_scale_shuffle_16(1)); ++ scales[2] = _mm256_shuffle_epi8(all_scales, get_scale_shuffle_16(2)); ++ scales[3] = _mm256_shuffle_epi8(all_scales, get_scale_shuffle_16(3)); ++} ++ ++// TODO: find the bug that causes this to be called without HAVE_FANCY_SIMD, which triggers ++// writing 4 vvalues into scales, which is of size 2. ++inline void set_scales_8_iq(int j, const __m256i& all_scales, __m256i * scales) { ++//#ifdef HAVE_FANCY_SIMD ++ auto shuffle = j == 0 ? _mm256_set_epi64x(0x0302030203020302, 0x0100010001000100, 0x0302030203020302, 0x0100010001000100) ++ : _mm256_set_epi64x(0x0b0a0b0a0b0a0b0a, 0x0908090809080908, 0x0b0a0b0a0b0a0b0a, 0x0908090809080908); ++ scales[0] = _mm256_shuffle_epi8(all_scales, shuffle); ++ scales[1] = _mm256_shuffle_epi8(all_scales, _mm256_add_epi8(shuffle, _mm256_set1_epi8(4))); ++//#else ++// set_scales_8(all_scales, j, scales); ++//#endif ++} ++ ++inline void set_scales_16_iq(const __m256i& all_scales, __m256i * scales) { ++#ifdef HAVE_FANCY_SIMD ++ auto shuffle = _mm256_set_epi64x(0x0706070607060706, 0x0302030203020302, 0x0504050405040504, 0x0100010001000100); ++ scales[0] = _mm256_shuffle_epi8(all_scales, shuffle); ++ scales[1] = _mm256_shuffle_epi8(all_scales, _mm256_add_epi8(shuffle, _mm256_set1_epi8(8))); ++#else ++ set_scales_16(all_scales, scales); ++#endif ++} ++ ++struct SimpleBits { ++ __m256i values[4]; ++}; ++ ++struct EvenSignHelper { ++#if defined HAVE_FANCY_SIMD && defined __AVX512VPOPCNTDQ__ ++ union sbits_t { ++ __m128i vec; ++ __mmask32 mask[4]; ++ }; ++ IQK_ALWAYS_INLINE void sign_2_values(__m256i aux, __m256i * values) const { ++ aux = _mm256_and_si256(_mm256_srlv_epi32(aux, shifts), mask); ++ auto pcnt = _mm256_popcnt_epi32(aux); ++ sbits_t sbits; ++ sbits.vec = _mm256_cvtepi32_epi8(_mm256_or_si256(aux, _mm256_slli_epi32(_mm256_and_si256(pcnt, mone), 7))); ++ values[0] = _mm256_mask_sub_epi8(values[0], sbits.mask[0], _mm256_setzero_si256(), values[0]); ++ values[1] = _mm256_mask_sub_epi8(values[1], sbits.mask[1], _mm256_setzero_si256(), values[1]); ++ //auto sign_bits = _mm256_cvtepi32_epi8(_mm256_or_si256(aux, _mm256_slli_epi32(_mm256_and_si256(pcnt, mone), 7))); ++ //const __mmask32 * m32 = (const __mmask32 *)&sign_bits; ++ //values[0] = _mm256_mask_sub_epi8(values[0], m32[0], _mm256_setzero_si256(), values[0]); ++ //values[1] = _mm256_mask_sub_epi8(values[1], m32[1], _mm256_setzero_si256(), values[1]); ++ } ++ const __m256i shifts = _mm256_set_epi32(21, 14, 7, 0, 21, 14, 7, 0); ++ const __m256i mask = _mm256_set1_epi32(127); ++ const __m256i mone = _mm256_set1_epi32(1); ++#endif ++ inline void sign_value(uint32_t aux32, __m256i& value) const { ++ auto signs = _mm256_set_epi64x(keven_signs[(aux32 >> 21) & 127], keven_signs[(aux32 >> 14) & 127], ++ keven_signs[(aux32 >> 7) & 127], keven_signs[(aux32 >> 0) & 127]); ++ value = _mm256_sign_epi8(value, signs); ++ } ++}; ++ ++struct SignHelper { ++ inline __m256i make_signs(uint32_t sign_bits) const { ++ auto aux256 = _mm256_set1_epi32(sign_bits); ++ aux256 = _mm256_and_si256(_mm256_shuffle_epi8(aux256, mask1), mask2); ++ return _mm256_or_si256(_mm256_cmpeq_epi8(aux256, mask2), mone); ++ } ++// inline __m256i make_signs(const uint16_t * sign_bits) const { ++//#ifdef HAVE_FANCY_SIMD ++//#else ++// return make_signs(sign_bits[0] | (sign_bits[1] << 16)); ++//#endif ++// } ++ inline __m256i sign_value(const uint16_t * sign_bits, const __m256i& value) const { ++#ifdef HAVE_FANCY_SIMD ++ const __mmask32 * mask = (const __mmask32 *)sign_bits; ++ return _mm256_mask_sub_epi8(value, mask[0], _mm256_setzero_si256(), value); ++#else ++ return _mm256_sign_epi8(value, make_signs(sign_bits[0] | (sign_bits[1] << 16))); ++#endif ++ } ++ IQK_ALWAYS_INLINE void sign_4_values(const uint16_t * sign_bits, __m256i * values) const { ++ // Somehow the FANCY_SIMD version has become 50% slower for TG??? ++#ifdef z_HAVE_FANCY_SIMD ++ //__mmask32 mask[4]; std::memcpy(mask, sign_bits, 4*sizeof(__mmask32)); ++ const __mmask32 * mask = (const __mmask32 *)sign_bits; ++ values[0] = _mm256_mask_sub_epi8(values[0], mask[0], _mm256_setzero_si256(), values[0]); ++ values[1] = _mm256_mask_sub_epi8(values[1], mask[1], _mm256_setzero_si256(), values[1]); ++ values[2] = _mm256_mask_sub_epi8(values[2], mask[2], _mm256_setzero_si256(), values[2]); ++ values[3] = _mm256_mask_sub_epi8(values[3], mask[3], _mm256_setzero_si256(), values[3]); ++#else ++ auto s128 = _mm_loadu_si128((const __m128i *)sign_bits); ++ auto s256 = MM256_SET1_M128I(s128); ++ __m256i aux256; ++ auto shuffle = mask1; ++ auto step = _mm256_set1_epi8(4); ++ aux256 = _mm256_and_si256(_mm256_shuffle_epi8(s256, shuffle), mask2); shuffle = _mm256_add_epi8(shuffle, step); ++ values[0] = _mm256_sign_epi8(values[0], _mm256_or_si256(_mm256_cmpeq_epi8(aux256, mask2), mone)); ++ aux256 = _mm256_and_si256(_mm256_shuffle_epi8(s256, shuffle), mask2); shuffle = _mm256_add_epi8(shuffle, step); ++ values[1] = _mm256_sign_epi8(values[1], _mm256_or_si256(_mm256_cmpeq_epi8(aux256, mask2), mone)); ++ aux256 = _mm256_and_si256(_mm256_shuffle_epi8(s256, shuffle), mask2); shuffle = _mm256_add_epi8(shuffle, step); ++ values[2] = _mm256_sign_epi8(values[2], _mm256_or_si256(_mm256_cmpeq_epi8(aux256, mask2), mone)); ++ aux256 = _mm256_and_si256(_mm256_shuffle_epi8(s256, shuffle), mask2); shuffle = _mm256_add_epi8(shuffle, step); ++ values[3] = _mm256_sign_epi8(values[3], _mm256_or_si256(_mm256_cmpeq_epi8(aux256, mask2), mone)); ++#endif ++ } ++ const __m256i mask1 = _mm256_set_epi64x(0x0303030303030303, 0x0202020202020202, 0x0101010101010101, 0x0000000000000000); ++ const __m256i mask2 = _mm256_set1_epi64x(0x8040201008040201ull); ++ const __m256i mone = _mm256_set1_epi8(1); ++}; ++ ++struct DequantizerIQ2XXS final : public BaseDequantizer { ++ DequantizerIQ2XXS(const void * vx, size_t bx) : BaseDequantizer(vx, bx) {} ++ ++ constexpr static int num_blocks = 8; ++ ++ union Data { ++ __m256i vec; ++ uint32_t val[8]; ++ }; ++ ++ inline __m128i load_scales(int i) { ++ d = 0.125f * GGML_FP16_TO_FP32(x[i].d); ++ const uint16_t * a16 = (const uint16_t *)x[i].qs; ++ auto scales = _mm_srli_epi16(_mm_set_epi16(a16[31], a16[27], a16[23], a16[19], a16[15], a16[11], a16[7], a16[3]), 12); ++ return _mm_or_si128(_mm_slli_epi16(scales, 1), _mm_set1_epi16(1)); ++ } ++ ++ inline void new_block(int i, __m256i * scales) { ++ auto sc16 = load_scales(i); ++ scales[0] = MM256_SET1_M128I(sc16); ++ } ++ inline void new_block_f(int i, __m256 * scales) { ++ auto sc16 = load_scales(i); ++ auto scf = _mm256_mul_ps(_mm256_set1_ps(d), _mm256_cvtepi32_ps(_mm256_cvtepi16_epi32(sc16))); ++ auto scf_l = _mm256_castps256_ps128(scf); ++ auto scf_h = _mm256_extractf128_ps(scf, 1); ++ scales[0] = _mm256_set_m128(scf_l, scf_l); ++ scales[1] = _mm256_set_m128(scf_h, scf_h); ++ scales[2] = _mm256_mul_ps(scf, _mm256_set1_ps(-minv)); ++ } ++ ++ inline float new_block(int i, __m256i * scales, __m256i& mins) { ++ auto sc16 = load_scales(i); ++ mins = scb.shuffle(sc16); ++ scales[0] = MM256_SET1_M128I(sc16); ++ return -d*minv; ++ } ++ ++ inline static void make4(const uint32_t * aux32, __m256i * values) { ++ const uint8_t * aux8 = (const uint8_t *)aux32; ++ values[0] = _mm256_set_epi64x(iq2xxs_grid[aux8[ 3]], iq2xxs_grid[aux8[ 2]], iq2xxs_grid[aux8[ 1]], iq2xxs_grid[aux8[ 0]]); ++ values[1] = _mm256_set_epi64x(iq2xxs_grid[aux8[11]], iq2xxs_grid[aux8[10]], iq2xxs_grid[aux8[ 9]], iq2xxs_grid[aux8[ 8]]); ++ values[2] = _mm256_set_epi64x(iq2xxs_grid[aux8[19]], iq2xxs_grid[aux8[18]], iq2xxs_grid[aux8[17]], iq2xxs_grid[aux8[16]]); ++ values[3] = _mm256_set_epi64x(iq2xxs_grid[aux8[27]], iq2xxs_grid[aux8[26]], iq2xxs_grid[aux8[25]], iq2xxs_grid[aux8[24]]); ++ } ++ ++ IQK_ALWAYS_INLINE void sign_values(const uint32_t * aux32, __m256i * values) const { ++#if defined z_HAVE_FANCY_SIMD && defined __AVX512VPOPCNTDQ__ ++ esh.sign_2_values(MM256_SET_M128I(_mm_set1_epi32(aux32[3]), _mm_set1_epi32(aux32[1])), values+0); ++ esh.sign_2_values(MM256_SET_M128I(_mm_set1_epi32(aux32[7]), _mm_set1_epi32(aux32[5])), values+2); ++#else ++ esh.sign_value(aux32[1], values[0]); ++ esh.sign_value(aux32[3], values[1]); ++ esh.sign_value(aux32[5], values[2]); ++ esh.sign_value(aux32[7], values[3]); ++#endif ++ } ++ inline void make4_signed(const uint32_t * aux32, const __m256i& min_value, __m256i * values) const { ++ make4(aux32, values); ++ sign_values(aux32, values); ++ for (int k = 0; k < 4; ++k) values[k] = _mm256_add_epi8(values[k], min_value); ++ } ++ inline void make4(const uint32_t * aux32, __m256i * values, __m256i * q8) const { ++ make4(aux32, values); ++ sign_values(aux32, q8); ++ } ++ inline void prepare(int i, int j) { ++ Data data; data.vec = _mm256_loadu_si256((const __m256i *)x[i].qs + j); ++ make4_signed(data.val, min_value, bits.values); ++ } ++ inline void prepare(int i, int j, const Q8<1>& q8, __m256i * q8_quants) { ++ for (int k = 0; k < 4; ++k) q8_quants[k] = q8.load_quants(0, i, 4*j+k); ++ Data data; data.vec = _mm256_loadu_si256((const __m256i *)x[i].qs + j); ++ make4(data.val, bits.values, q8_quants); ++ } ++ inline void prepare(int i, int j, __m256i * q8_quants) { ++ Data data; data.vec = _mm256_loadu_si256((const __m256i *)x[i].qs + j); ++ make4(data.val, bits.values, q8_quants); ++ } ++ ++ constexpr static int minv = 43; ++ SimpleBits bits; ++ Scales8KBase scb; ++ EvenSignHelper esh; ++ const __m256i min_value = _mm256_set1_epi8(minv); ++ const __m256i shuffle = _mm256_set_epi32(7, 5, 3, 1, 7, 5, 3, 1); ++}; ++ ++struct DequantizerIQ2XS final : public BaseDequantizer { ++ DequantizerIQ2XS(const void * vx, size_t bx) : BaseDequantizer(vx, bx) {} ++ ++ constexpr static int num_blocks = 16; ++ ++ static inline __m256i make_scales(const uint8_t * scales) { ++ auto tmp = _mm_loadl_epi64((const __m128i *)scales); ++ auto all = _mm_and_si128(_mm_unpacklo_epi8(tmp, _mm_srli_epi16(tmp, 4)), _mm_set1_epi8(0xf)); ++ auto scales8 = _mm_or_si128(_mm_slli_epi16(all, 1), _mm_set1_epi8(1)); ++ return _mm256_cvtepi8_epi16(scales8); ++ } ++ ++ inline __m256i load_scales(int i) { ++ d = 0.125f * GGML_FP16_TO_FP32(x[i].d); ++ return make_scales(x[i].scales); ++ } ++ inline static void prepare_scales(const __m256i& all, __m256i * scales) { ++ auto scales_l = _mm256_castsi256_si128(all); ++ auto scales_h = _mm256_extractf128_si256(all, 1); ++ scales[0] = MM256_SET1_M128I(scales_l); ++ scales[1] = MM256_SET1_M128I(scales_h); ++ } ++ ++ inline void new_block(int i, __m256i * scales) { ++ prepare_scales(load_scales(i), scales); ++ } ++ inline float new_block(int i, __m256i * scales, __m256i& mins) { ++ mins = load_scales(i); ++ prepare_scales(mins, scales); ++ return -d*minv; ++ } ++ ++ struct Helper { ++ const __m256i mone = _mm256_set1_epi8(1); ++ const __m256i mask = _mm256_set1_epi64x(0x8040201008040201); ++ //const __m256i bhelper = _mm256_set_epi64x(0x8000008000808000, 0x0080800080000080, 0x8000008000808000, 0x0080800080000080); ++ const __m256i bhelper = load_bhelper(); ++ const __m256i shuff1 = _mm256_set_epi64x(0x0606060606060606, 0x0404040404040404, 0x0202020202020202, 0x0000000000000000); ++ const __m256i shuff2 = _mm256_set_epi64x(0x0e0e0e0e0e0e0e0e, 0x0c0c0c0c0c0c0c0c, 0x0a0a0a0a0a0a0a0a, 0x0808080808080808); ++ static __m256i load_bhelper() { ++ static const uint8_t k_bit_helper[32] = { ++ 0x00, 0x80, 0x80, 0x00, 0x80, 0x00, 0x00, 0x80, 0x80, 0x00, 0x00, 0x80, 0x00, 0x80, 0x80, 0x00, ++ 0x00, 0x80, 0x80, 0x00, 0x80, 0x00, 0x00, 0x80, 0x80, 0x00, 0x00, 0x80, 0x00, 0x80, 0x80, 0x00, ++ }; ++ return _mm256_loadu_si256((const __m256i*)k_bit_helper); ++ } ++ }; ++ ++ union index_t { ++ __m256i vec; ++ uint16_t val[8]; ++ }; ++ ++ inline static void make4(const __m256i& data, const __m256i& mask, __m256i * values) { ++ index_t idx; ++ idx.vec = _mm256_and_si256(data, mask); ++ values[0] = _mm256_set_epi64x(iq2xs_grid[idx.val[ 3]], iq2xs_grid[idx.val[ 2]], iq2xs_grid[idx.val[ 1]], iq2xs_grid[idx.val[ 0]]); ++ values[1] = _mm256_set_epi64x(iq2xs_grid[idx.val[ 7]], iq2xs_grid[idx.val[ 6]], iq2xs_grid[idx.val[ 5]], iq2xs_grid[idx.val[ 4]]); ++ values[2] = _mm256_set_epi64x(iq2xs_grid[idx.val[11]], iq2xs_grid[idx.val[10]], iq2xs_grid[idx.val[ 9]], iq2xs_grid[idx.val[ 8]]); ++ values[3] = _mm256_set_epi64x(iq2xs_grid[idx.val[15]], iq2xs_grid[idx.val[14]], iq2xs_grid[idx.val[13]], iq2xs_grid[idx.val[12]]); ++ } ++ inline static void sign_value(const __m256i& sign_bits, const __m256i& shuffle, const __m256i& mask, ++ const __m256i& mone, __m256i& value) { ++ auto signs = _mm256_shuffle_epi8(sign_bits, shuffle); ++ signs = _mm256_cmpeq_epi8(_mm256_and_si256(signs, mask), mask); ++ value = _mm256_sign_epi8(value, _mm256_or_si256(signs, mone)); ++ } ++#if defined HAVE_FANCY_SIMD && defined __AVX512VPOPCNTDQ__ ++ static IQK_ALWAYS_INLINE inline void sign_values_popcnt(const __m256i& data, __m256i * values) { ++ auto partial_bits = _mm256_cvtepi16_epi8(_mm256_srli_epi16(data, 9)); ++ auto pcnt = _mm_popcnt_epi8(partial_bits); ++ auto full_bits = _mm_or_si128(partial_bits, _mm_slli_epi16(_mm_and_si128(pcnt, _mm_set1_epi8(1)), 7)); ++ const __mmask32 * m32 = (const __mmask32 *)&full_bits; ++ auto zero = _mm256_setzero_si256(); ++ values[0] = _mm256_mask_sub_epi8(values[0], m32[0], zero, values[0]); ++ values[1] = _mm256_mask_sub_epi8(values[1], m32[1], zero, values[1]); ++ values[2] = _mm256_mask_sub_epi8(values[2], m32[2], zero, values[2]); ++ values[3] = _mm256_mask_sub_epi8(values[3], m32[3], zero, values[3]); ++ } ++#else ++ static IQK_ALWAYS_INLINE inline void sign_values_helper(const __m256i& data, const Helper& helper, __m256i * values) { ++ auto psb1 = _mm256_srli_epi16(data, 9); ++ auto psb2 = _mm256_srli_epi16(data, 13); ++ auto psbc = _mm256_xor_si256(psb1, psb2); ++ auto oddb = _mm256_shuffle_epi8(helper.bhelper, psbc); ++ auto full = _mm256_or_si256(psb1, oddb); ++ auto full_l = _mm256_castsi256_si128(full); ++ auto full_h = _mm256_extractf128_si256(full, 1); ++ auto full_1 = MM256_SET1_M128I(full_l); ++ auto full_2 = MM256_SET1_M128I(full_h); ++ sign_value(full_1, helper.shuff1, helper.mask, helper.mone, values[0]); ++ sign_value(full_1, helper.shuff2, helper.mask, helper.mone, values[1]); ++ sign_value(full_2, helper.shuff1, helper.mask, helper.mone, values[2]); ++ sign_value(full_2, helper.shuff2, helper.mask, helper.mone, values[3]); ++ } ++#endif ++ IQK_ALWAYS_INLINE inline void sign_values(const __m256i& data, __m256i * values) const { ++#if defined HAVE_FANCY_SIMD && defined __AVX512VPOPCNTDQ__ ++ sign_values_popcnt(data, values); ++#else ++ sign_values_helper(data, helper, values); ++#endif ++ } ++ inline void make4_signed(const uint16_t * qs, const __m256i& m511, ++ const __m256i& min_value, __m256i * values) const { ++ auto q2 = _mm256_loadu_si256((const __m256i *)qs); ++ make4(q2, m511, values); ++ sign_values(q2, values); ++ for (int k = 0; k < 4; ++k) values[k] = _mm256_add_epi8(values[k], min_value); ++ } ++ inline void make4(const uint16_t * qs, const __m256i& m511, __m256i * values, __m256i * q8) const { ++ auto q2 = _mm256_loadu_si256((const __m256i *)qs); ++ make4(q2, m511, values); ++ sign_values(q2, q8); ++ } ++ ++ inline void prepare(int i, int j) { ++ make4_signed(x[i].qs + 16*j, idx_mask, min_value, bits.values); ++ } ++ inline void prepare(int i, int j, const Q8<1>& q8, __m256i * q8_quants) { ++ for (int k = 0; k < 4; ++k) q8_quants[k] = q8.load_quants(0, i, 4*j+k); ++ make4(x[i].qs + 16*j, idx_mask, bits.values, q8_quants); ++ } ++ inline void prepare_signed(int i, int j, __m256i * us) { ++ auto q2 = _mm256_loadu_si256((const __m256i *)x[i].qs+j); ++ make4(q2, idx_mask, us); ++ for (int k = 0; k < 4; ++k) bits.values[k] = us[k]; ++ sign_values(q2, bits.values); ++ } ++ IQK_ALWAYS_INLINE inline void prepare_signed(int i, int j, __m256i * us, __m256i * s) { ++ auto q2 = _mm256_loadu_si256((const __m256i *)x[i].qs+j); ++ make4(q2, idx_mask, us); ++ sign_values(q2, s); ++ } ++ ++ constexpr static int minv = 43; ++ ++ SimpleBits bits; ++#if !(defined HAVE_FANCY_SIMD && defined __AVX512VPOPCNTDQ__) ++ Helper helper; ++#endif ++ const __m256i idx_mask = _mm256_set1_epi16(511); ++ const __m256i min_value = _mm256_set1_epi8(minv); ++ ++}; ++ ++struct DequantizerIQ2S final : public BaseDequantizer { ++ DequantizerIQ2S(const void * vx, size_t bx) : BaseDequantizer(vx, bx) {} ++ ++ constexpr static int num_blocks = 16; ++ ++ static inline __m256i make_scales(const uint8_t * scales) { ++ auto tmp = _mm_loadl_epi64((const __m128i *)scales); ++ auto all = _mm_and_si128(_mm_unpacklo_epi8(tmp, _mm_srli_epi16(tmp, 4)), _mm_set1_epi8(0xf)); ++ auto scales8 = _mm_or_si128(_mm_slli_epi16(all, 1), _mm_set1_epi8(1)); ++ return _mm256_cvtepi8_epi16(scales8); ++ } ++ inline __m256i load_scales(int i) { ++ d = 0.125f * GGML_FP16_TO_FP32(x[i].d); ++ return make_scales(x[i].scales); ++ } ++ inline static void prepare_scales(const __m256i& all, __m256i * scales) { ++ auto scales_l = _mm256_castsi256_si128(all); ++ auto scales_h = _mm256_extractf128_si256(all, 1); ++ scales[0] = MM256_SET1_M128I(scales_l); ++ scales[1] = MM256_SET1_M128I(scales_h); ++ } ++ ++ inline void new_block(int i, __m256i * scales) { ++ prepare_scales(load_scales(i), scales); ++ } ++ inline float new_block(int i, __m256i * scales, __m256i& mins) { ++ mins = load_scales(i); ++ prepare_scales(mins, scales); ++ return -d*minv; ++ } ++ ++ union index_t { ++ __m256i vec; ++ uint32_t val[8]; ++ }; ++ ++ inline static void make2(const uint8_t * qs, const uint8_t * qh, const __m256i& idx_shift, const __m256i& idx_mask, __m256i * values) { ++ auto idx_l = _mm256_cvtepu8_epi32(_mm_loadl_epi64((const __m128i *)qs)); ++ auto idx_h = MM256_SET_M128I(_mm_set1_epi32(qh[1]), _mm_set1_epi32(qh[0])); ++ index_t idx; ++ idx.vec = _mm256_or_si256(idx_l, _mm256_and_si256(_mm256_sllv_epi32(idx_h, idx_shift), idx_mask)); ++ values[0] = _mm256_set_epi64x(iq2s_grid[idx.val[3]], iq2s_grid[idx.val[2]], iq2s_grid[idx.val[1]], iq2s_grid[idx.val[0]]); ++ values[1] = _mm256_set_epi64x(iq2s_grid[idx.val[7]], iq2s_grid[idx.val[6]], iq2s_grid[idx.val[5]], iq2s_grid[idx.val[4]]); ++ } ++ inline static void make2_signed(const SignHelper& sh, const uint8_t * qs, const uint8_t * qh, const uint16_t * sidx, ++ const __m256i& idx_shift, const __m256i& idx_mask, const __m256i& min_value, __m256i * values) { ++ make2(qs, qh, idx_shift, idx_mask, values); ++ values[0] = _mm256_add_epi8(sh.sign_value(sidx+0, values[0]), min_value); ++ values[1] = _mm256_add_epi8(sh.sign_value(sidx+2, values[1]), min_value); ++ } ++ ++ inline void prepare(int i, int j) { ++ auto qs = x[i].qs + 16*j; ++ auto qh = x[i].qh + 4*j; ++ const uint16_t * signs = (const uint16_t *)(x[i].qs + QK_K/8) + 8*j; ++ make2_signed(sh, qs+0, qh+0, signs+0, idx_shift, idx_mask, min_value, bits.values+0); ++ make2_signed(sh, qs+8, qh+2, signs+4, idx_shift, idx_mask, min_value, bits.values+2); ++ } ++ inline void prepare(int i, int j, const Q8<1>& q8, __m256i * q8_quants) { ++ auto qs = x[i].qs + 16*j; ++ auto qh = x[i].qh + 4*j; ++ const uint16_t * signs = (const uint16_t *)(x[i].qs + QK_K/8) + 8*j; ++ make2(qs+0, qh+0, idx_shift, idx_mask, bits.values+0); ++ make2(qs+8, qh+2, idx_shift, idx_mask, bits.values+2); ++ q8_quants[0] = _mm256_sign_epi8(q8.load_quants(0, i, 4*j+0), sh.make_signs(signs[0] | (signs[1] << 16))); ++ q8_quants[1] = _mm256_sign_epi8(q8.load_quants(0, i, 4*j+1), sh.make_signs(signs[2] | (signs[3] << 16))); ++ q8_quants[2] = _mm256_sign_epi8(q8.load_quants(0, i, 4*j+2), sh.make_signs(signs[4] | (signs[5] << 16))); ++ q8_quants[3] = _mm256_sign_epi8(q8.load_quants(0, i, 4*j+3), sh.make_signs(signs[6] | (signs[7] << 16))); ++ } ++ static inline void prepare(const uint8_t * qs, const uint8_t * qh, const uint16_t * signs, const SignHelper& sh, __m256i * values) { ++ auto idx_shift = _mm256_set_epi32(2, 4, 6, 8, 2, 4, 6, 8); ++ auto idx_mask = _mm256_set1_epi32(0x300); ++ make2(qs+0, qh+0, idx_shift, idx_mask, values+0); ++ make2(qs+8, qh+2, idx_shift, idx_mask, values+2); ++ values[0] = _mm256_sign_epi8(values[0], sh.make_signs(signs[0] | (signs[1] << 16))); ++ values[1] = _mm256_sign_epi8(values[1], sh.make_signs(signs[2] | (signs[3] << 16))); ++ values[2] = _mm256_sign_epi8(values[2], sh.make_signs(signs[4] | (signs[5] << 16))); ++ values[3] = _mm256_sign_epi8(values[3], sh.make_signs(signs[6] | (signs[7] << 16))); ++ } ++ inline void prepare_signed(int i, int j, __m256i * us, __m256i * s) { ++ auto qs = x[i].qs + 16*j; ++ auto qh = x[i].qh + 4*j; ++ const uint16_t * signs = (const uint16_t *)(x[i].qs + QK_K/8) + 8*j; ++ make2(qs+0, qh+0, idx_shift, idx_mask, us+0); ++ make2(qs+8, qh+2, idx_shift, idx_mask, us+2); ++ s[0] = _mm256_sign_epi8(s[0], sh.make_signs(signs[0] | (signs[1] << 16))); ++ s[1] = _mm256_sign_epi8(s[1], sh.make_signs(signs[2] | (signs[3] << 16))); ++ s[2] = _mm256_sign_epi8(s[2], sh.make_signs(signs[4] | (signs[5] << 16))); ++ s[3] = _mm256_sign_epi8(s[3], sh.make_signs(signs[6] | (signs[7] << 16))); ++ } ++ inline void prepare_signed(int i, int j, __m256i * us) { ++ auto qs = x[i].qs + 16*j; ++ auto qh = x[i].qh + 4*j; ++ const uint16_t * signs = (const uint16_t *)(x[i].qs + QK_K/8) + 8*j; ++ make2(qs+0, qh+0, idx_shift, idx_mask, us+0); ++ make2(qs+8, qh+2, idx_shift, idx_mask, us+2); ++ bits.values[0] = _mm256_sign_epi8(us[0], sh.make_signs(signs[0] | (signs[1] << 16))); ++ bits.values[1] = _mm256_sign_epi8(us[1], sh.make_signs(signs[2] | (signs[3] << 16))); ++ bits.values[2] = _mm256_sign_epi8(us[2], sh.make_signs(signs[4] | (signs[5] << 16))); ++ bits.values[3] = _mm256_sign_epi8(us[3], sh.make_signs(signs[6] | (signs[7] << 16))); ++ } ++ ++ constexpr static int minv = 43; ++ ++ SimpleBits bits; ++ SignHelper sh; ++ const __m256i idx_shift = _mm256_set_epi32(2, 4, 6, 8, 2, 4, 6, 8); ++ const __m256i idx_mask = _mm256_set1_epi32(0x300); ++ const __m256i min_value = _mm256_set1_epi8(minv); ++ ++}; ++ ++struct DequantizerIQ3XXS final : public BaseDequantizer { ++ DequantizerIQ3XXS(const void * vx, size_t bx) : BaseDequantizer(vx, bx) {} ++ ++ constexpr static int num_blocks = 8; ++ ++ inline __m128i prepare_scales(int i) { ++ d = 0.25f * GGML_FP16_TO_FP32(x[i].d); ++ auto tmp = _mm256_loadu_si256((const __m256i *)(x[i].qs + QK_K/4)); ++ auto scales32 = _mm256_srli_epi32(tmp, 28); ++ scales32 = _mm256_or_si256(_mm256_slli_epi32(scales32, 1), _mm256_set1_epi32(1)); ++ return _mm_packs_epi32(_mm256_castsi256_si128(scales32), _mm256_extractf128_si256(scales32, 1)); ++ } ++ ++ inline void new_block(int i, __m256i * scales) { ++ auto scales16 = prepare_scales(i); ++ scales[0] = MM256_SET1_M128I(scales16); ++ } ++ inline void new_block_f(int i, __m256 * scales) { ++ auto sc16 = prepare_scales(i); ++ auto scf = _mm256_mul_ps(_mm256_set1_ps(d), _mm256_cvtepi32_ps(_mm256_cvtepi16_epi32(sc16))); ++ auto scf_l = _mm256_castps256_ps128(scf); ++ auto scf_h = _mm256_extractf128_ps(scf, 1); ++ scales[0] = _mm256_set_m128(scf_l, scf_l); ++ scales[1] = _mm256_set_m128(scf_h, scf_h); ++ scales[2] = _mm256_mul_ps(scf, _mm256_set1_ps(-minv)); ++ } ++ inline float new_block(int i, __m256i * scales, __m256i& mins) { ++ auto scales16 = prepare_scales(i); ++ mins = scb.shuffle(scales16); ++ scales[0] = MM256_SET1_M128I(scales16); ++ return -d*minv; ++ } ++ ++ inline static __m256i make_quants(const uint8_t * qs) { ++ return _mm256_set_epi32(iq3xxs_grid[qs[7]], iq3xxs_grid[qs[6]], iq3xxs_grid[qs[5]], iq3xxs_grid[qs[4]], ++ iq3xxs_grid[qs[3]], iq3xxs_grid[qs[2]], iq3xxs_grid[qs[1]], iq3xxs_grid[qs[0]]); ++ } ++ inline static void make4_unsigned(const uint8_t * qs, __m256i * values) { ++ values[0] = make_quants(qs+ 0); ++ values[1] = make_quants(qs+ 8); ++ values[2] = make_quants(qs+16); ++ values[3] = make_quants(qs+24); ++ } ++ ++ IQK_ALWAYS_INLINE void sign_2_values(const uint16_t * signs, __m256i * values) const { ++#if defined HAVE_FANCY_SIMD && defined __AVX512VPOPCNTDQ__ ++ esh.sign_2_values(MM256_SET_M128I(_mm_set1_epi32(signs[2] | (signs[3] << 16)), _mm_set1_epi32(signs[0] | (signs[1] << 16))), values); ++#else ++ esh.sign_value(signs[0] | (signs[1] << 16), values[0]); ++ esh.sign_value(signs[2] | (signs[3] << 16), values[1]); ++#endif ++ } ++ ++ inline void prepare(int i, int j) { ++ auto qs = x[i].qs + 32*j; ++ const uint16_t * signs = (const uint16_t *)(x[i].qs + QK_K/4) + 8*j; ++ make4_unsigned(qs, bits.values); ++ sign_2_values(signs+0, bits.values+0); ++ sign_2_values(signs+4, bits.values+2); ++ for (int k = 0; k < 4; ++k) bits.values[k] = _mm256_add_epi8(bits.values[k], min_value); ++ } ++ inline void prepare(int i, int j, const Q8<1>& q8, __m256i * q8_quants) { ++ for (int k = 0; k < 4; ++k) q8_quants[k] = q8.load_quants(0, i, 4*j+k); ++ auto qs = x[i].qs + 32*j; ++ const uint16_t * signs = (const uint16_t *)(x[i].qs + QK_K/4) + 8*j; ++ make4_unsigned(qs, bits.values); ++ sign_2_values(signs+0, q8_quants+0); ++ sign_2_values(signs+4, q8_quants+2); ++ } ++ inline void prepare(int i, int j, __m256i * q8_quants) { ++ auto qs = x[i].qs + 32*j; ++ const uint16_t * signs = (const uint16_t *)(x[i].qs + QK_K/4) + 8*j; ++ make4_unsigned(qs, bits.values); ++ sign_2_values(signs+0, q8_quants+0); ++ sign_2_values(signs+4, q8_quants+2); ++ } ++ ++ constexpr static int minv = 64; ++ ++ SimpleBits bits; ++ Scales8KBase scb; ++ EvenSignHelper esh; ++ const __m256i min_value = _mm256_set1_epi8(minv); ++ ++}; ++ ++#ifdef z_HAVE_FANCY_SIMD ++// Strangely enough, the following implementation makes PP ~6% slower and TG ~6% faster ++// compared to the vanilla AVX2 version below. ++struct IndexHelperIQ3S { ++ union index_t { ++ __m256i vec; ++ uint16_t val[16]; ++ }; ++ inline void make2(const uint8_t * qs, const uint8_t * qh, __m256i * values) const { ++ auto idx_l = _mm256_cvtepu8_epi16(_mm_loadu_si128((const __m128i *)qs)); ++ const __mmask16 * m16 = (const __mmask16 *)qh; ++ index_t idx; ++ idx.vec = _mm256_mask_add_epi16(idx_l, m16[0], idx_l, offset); ++ values[0] = _mm256_set_epi32(iq3s_grid[idx.val[ 7]], iq3s_grid[idx.val[ 6]], iq3s_grid[idx.val[ 5]], iq3s_grid[idx.val[ 4]], ++ iq3s_grid[idx.val[ 3]], iq3s_grid[idx.val[ 2]], iq3s_grid[idx.val[ 1]], iq3s_grid[idx.val[ 0]]); ++ values[1] = _mm256_set_epi32(iq3s_grid[idx.val[15]], iq3s_grid[idx.val[14]], iq3s_grid[idx.val[13]], iq3s_grid[idx.val[12]], ++ iq3s_grid[idx.val[11]], iq3s_grid[idx.val[10]], iq3s_grid[idx.val[ 9]], iq3s_grid[idx.val[ 8]]); ++ } ++ const __m256i offset = _mm256_set1_epi16(256); ++}; ++#else ++struct IndexHelperIQ3S { ++ union index_t { ++ __m256i vec; ++ uint32_t val[8]; ++ }; ++ inline void make2(const uint8_t * qs, const uint8_t * qh, __m256i * values) const { ++ index_t idx; ++ auto idx_l = _mm256_cvtepu8_epi32(_mm_loadl_epi64((const __m128i *)qs)); ++ auto idx_h = _mm256_and_si256(_mm256_sllv_epi32(_mm256_set1_epi32(qh[0]), idx_shift), idx_mask); ++ idx.vec = _mm256_or_si256(idx_h, idx_l); ++ values[0] = _mm256_set_epi32(iq3s_grid[idx.val[7]], iq3s_grid[idx.val[6]], iq3s_grid[idx.val[5]], iq3s_grid[idx.val[4]], ++ iq3s_grid[idx.val[3]], iq3s_grid[idx.val[2]], iq3s_grid[idx.val[1]], iq3s_grid[idx.val[0]]); ++ idx_l = _mm256_cvtepu8_epi32(_mm_loadl_epi64((const __m128i *)(qs+8))); ++ idx_h = _mm256_and_si256(_mm256_sllv_epi32(_mm256_set1_epi32(qh[1]), idx_shift), idx_mask); ++ idx.vec = _mm256_or_si256(idx_h, idx_l); ++ values[1] = _mm256_set_epi32(iq3s_grid[idx.val[7]], iq3s_grid[idx.val[6]], iq3s_grid[idx.val[5]], iq3s_grid[idx.val[4]], ++ iq3s_grid[idx.val[3]], iq3s_grid[idx.val[2]], iq3s_grid[idx.val[1]], iq3s_grid[idx.val[0]]); ++ } ++ const __m256i idx_mask = _mm256_set1_epi32(256); ++ const __m256i idx_shift = _mm256_set_epi32(1, 2, 3, 4, 5, 6, 7, 8); ++}; ++#endif ++ ++struct DequantizerIQ3S final : public BaseDequantizer { ++ DequantizerIQ3S(const void * vx, size_t bx) : BaseDequantizer(vx, bx) {} ++ ++ constexpr static int num_blocks = 8; ++ ++ inline __m128i make_scales(int i, float& dd) const { ++ dd = GGML_FP16_TO_FP32(x[i].d); ++ uint32_t aux32[2]; ++ std::memcpy(aux32, x[i].scales, 4); ++ aux32[1] = (aux32[0] >> 4) & 0x0f0f0f0f; ++ aux32[0] &= 0x0f0f0f0f; ++ auto scales8 = _mm_shuffle_epi8(_mm_loadl_epi64((const __m128i *)aux32), _mm_set1_epi64x(0x0703060205010400)); ++ auto scales16 = _mm256_castsi256_si128(_mm256_cvtepi8_epi16(scales8)); ++ return _mm_or_si128(_mm_slli_epi16(scales16, 1), _mm_set1_epi16(1)); ++ } ++ inline void new_block(int i, __m256i * scales) { ++ auto scales16 = make_scales(i, d); ++ scales[0] = MM256_SET1_M128I(scales16); ++ } ++ inline void new_block_f(int i, __m256 * scales) { ++ auto sc16 = make_scales(i, d); ++ auto scf = _mm256_mul_ps(_mm256_set1_ps(d), _mm256_cvtepi32_ps(_mm256_cvtepi16_epi32(sc16))); ++ auto scf_l = _mm256_castps256_ps128(scf); ++ auto scf_h = _mm256_extractf128_ps(scf, 1); ++ scales[0] = _mm256_set_m128(scf_l, scf_l); ++ scales[1] = _mm256_set_m128(scf_h, scf_h); ++ scales[2] = _mm256_mul_ps(scf, _mm256_set1_ps(-minv)); ++ } ++ inline float new_block(int i, __m256i * scales, __m256i& mins) { ++ auto scales16 = make_scales(i, d); ++ mins = scb.shuffle(scales16); ++ scales[0] = MM256_SET1_M128I(scales16); ++ return -minv*d; ++ } ++ ++ inline void prepare(int i, int j) { ++ prepare_unsigned(i, j); ++ sh.sign_4_values((const uint16_t *)x[i].signs + 8*j, bits.values); ++ for (int k = 0; k < 4; ++k) bits.values[k] = _mm256_add_epi8(bits.values[k], min_value); ++ } ++ inline void prepare(int i, int j, const Q8<1>& q8, __m256i * q8_quants) { ++ prepare_unsigned(i, j); ++ for (int k = 0; k < 4; ++k) q8_quants[k] = q8.load_quants(0, i, 4*j+k); ++ sh.sign_4_values((const uint16_t *)x[i].signs + 8*j, q8_quants); ++ } ++ inline void prepare(int i, int j, __m256i * q8_quants) { ++ prepare_unsigned(i, j); ++ sh.sign_4_values((const uint16_t *)x[i].signs + 8*j, q8_quants); ++ } ++ ++ inline void prepare_unsigned(int i, int j) { ++ auto qs = x[i].qs + 32*j; ++ auto qh = x[i].qh + 4*j; ++ helper.make2(qs+ 0, qh+0, bits.values+0); ++ helper.make2(qs+16, qh+2, bits.values+2); ++ } ++ ++ constexpr static int minv = 16; ++ ++ SimpleBits bits; ++ SignHelper sh; ++ Scales8KBase scb; ++ IndexHelperIQ3S helper; ++ const __m256i min_value = _mm256_set1_epi8(minv); ++ ++}; ++ ++template ++inline void multiply_add_1(int j, const Bits& bits, const __m256i * scales, const __m256i * q8, __m256i * sumi) { ++ if (j == 0) { ++#ifdef HAVE_FANCY_SIMD ++ auto p1 = _mm256_dpbusd_epi32(_mm256_setzero_si256(), bits.values[0], q8[0]); ++ auto p2 = _mm256_dpbusd_epi32(_mm256_setzero_si256(), bits.values[1], q8[1]); ++ auto p3 = _mm256_dpbusd_epi32(_mm256_setzero_si256(), bits.values[2], q8[2]); ++ auto p4 = _mm256_dpbusd_epi32(_mm256_setzero_si256(), bits.values[3], q8[3]); ++ sumi[0] = _mm256_dpwssd_epi32(_mm256_setzero_si256(), scales[0], _mm256_packs_epi32(p1, p2)); ++ sumi[1] = _mm256_dpwssd_epi32(_mm256_setzero_si256(), scales[1], _mm256_packs_epi32(p3, p4)); ++#else ++ const __m256i p1 = _mm256_madd_epi16(scales[0], _mm256_maddubs_epi16(bits.values[0], q8[0])); ++ const __m256i p2 = _mm256_madd_epi16(scales[1], _mm256_maddubs_epi16(bits.values[1], q8[1])); ++ const __m256i p3 = _mm256_madd_epi16(scales[2], _mm256_maddubs_epi16(bits.values[2], q8[2])); ++ const __m256i p4 = _mm256_madd_epi16(scales[3], _mm256_maddubs_epi16(bits.values[3], q8[3])); ++ sumi[0] = _mm256_add_epi32(p1, p3); ++ sumi[1] = _mm256_add_epi32(p2, p4); ++#endif ++ } else { ++#ifdef HAVE_FANCY_SIMD ++ auto p1 = _mm256_dpbusd_epi32(_mm256_setzero_si256(), bits.values[0], q8[0]); ++ auto p2 = _mm256_dpbusd_epi32(_mm256_setzero_si256(), bits.values[1], q8[1]); ++ auto p3 = _mm256_dpbusd_epi32(_mm256_setzero_si256(), bits.values[2], q8[2]); ++ auto p4 = _mm256_dpbusd_epi32(_mm256_setzero_si256(), bits.values[3], q8[3]); ++ sumi[0] = _mm256_dpwssd_epi32(sumi[0], scales[0], _mm256_packs_epi32(p1, p2)); ++ sumi[1] = _mm256_dpwssd_epi32(sumi[1], scales[1], _mm256_packs_epi32(p3, p4)); ++#else ++ const __m256i p1 = _mm256_madd_epi16(scales[0], _mm256_maddubs_epi16(bits.values[0], q8[0])); ++ const __m256i p2 = _mm256_madd_epi16(scales[1], _mm256_maddubs_epi16(bits.values[1], q8[1])); ++ const __m256i p3 = _mm256_madd_epi16(scales[2], _mm256_maddubs_epi16(bits.values[2], q8[2])); ++ const __m256i p4 = _mm256_madd_epi16(scales[3], _mm256_maddubs_epi16(bits.values[3], q8[3])); ++ sumi[0] = _mm256_add_epi32(sumi[0], _mm256_add_epi32(p1, p3)); ++ sumi[1] = _mm256_add_epi32(sumi[1], _mm256_add_epi32(p2, p4)); ++#endif ++ } ++} ++ ++template ++static void mul_mat_qX_K_q8_K_IQ_1(int n, const void * vx, size_t bx, const DataInfo& info, int nrc_x) { ++ const int nb = n / QK_K; ++ Q8<1> q8(info); ++ Dequantizer deq(vx, bx); ++ __m256i scales[2]; ++ __m256i q8_quants[4]; ++ for (int ix = 0; ix < nrc_x; ++ix) { ++ ++ __m256 accd = _mm256_setzero_ps(); ++ deq.new_row(ix); ++ ++ for (int i = 0; i < nb; ++i) { ++ ++ __m256i sumi[2], all_scales[Dequantizer::num_blocks/8]; ++ deq.new_block(i, all_scales); ++ ++ for (int j = 0; j < QK_K/128; ++j) { ++ deq.prepare(i, j, q8, q8_quants); ++ if constexpr (Dequantizer::num_blocks == 8) { ++ set_scales_8_iq(j, all_scales[0], scales); ++ } else { ++ set_scales_16_iq(all_scales[j], scales); ++ } ++ multiply_add_1(j, deq.bits, scales, q8_quants, sumi); ++ } ++ accd = _mm256_fmadd_ps(_mm256_set1_ps(deq.d*q8.scale(0, i)), _mm256_cvtepi32_ps(_mm256_add_epi32(sumi[0], sumi[1])), accd); ++ } ++ ++ info.store(ix, 0, hsum_float_8(accd)); ++ } ++} ++ ++// So, if I uncomment this function and the call to it in mul_mat_qX_K_q8_K_IQ_N() below, ++// PP performance improves by ~2-3% (when we have __AVX512VNNI__ and __AVX512VL__). ++// But TG performance for iq3_xs drops by 35%. Seriously? I mean, c'mon, ++// what does the compilation of mul_mat_qX_K_q8_K_IQ_1 (which gets invoked during TG) ++// have to do with the compilation of mul_mat_qX_K_q8_K_IQ_N (invoked during PP)? ++//template ++//inline void multiply_add_iq(const Bits& bits, const __m256i * scales, int j, int i, const Q8& q8, __m256i * sumi) { ++//#if defined(__AVX512VNNI__) && defined(__AVX512VL__) ++// for (int iy = 0; iy < Q8::nrc_y; ++iy) { ++// sumi[iy] = _mm256_dpwssd_epi32(sumi[iy], scales[0], _mm256_maddubs_epi16(bits.values[0], q8.load_quants(iy, i, 4*j+0))); ++// sumi[iy] = _mm256_dpwssd_epi32(sumi[iy], scales[1], _mm256_maddubs_epi16(bits.values[1], q8.load_quants(iy, i, 4*j+1))); ++// sumi[iy] = _mm256_dpwssd_epi32(sumi[iy], scales[2], _mm256_maddubs_epi16(bits.values[2], q8.load_quants(iy, i, 4*j+2))); ++// sumi[iy] = _mm256_dpwssd_epi32(sumi[iy], scales[3], _mm256_maddubs_epi16(bits.values[3], q8.load_quants(iy, i, 4*j+3))); ++// } ++//#else ++// for (int iy = 0; iy < Q8::nrc_y; ++iy) { ++// const __m256i p1 = _mm256_madd_epi16(scales[0], _mm256_maddubs_epi16(bits.values[0], q8.load_quants(iy, i, 4*j+0))); ++// const __m256i p2 = _mm256_madd_epi16(scales[1], _mm256_maddubs_epi16(bits.values[1], q8.load_quants(iy, i, 4*j+1))); ++// const __m256i p3 = _mm256_madd_epi16(scales[2], _mm256_maddubs_epi16(bits.values[2], q8.load_quants(iy, i, 4*j+2))); ++// const __m256i p4 = _mm256_madd_epi16(scales[3], _mm256_maddubs_epi16(bits.values[3], q8.load_quants(iy, i, 4*j+3))); ++// sumi[iy] = _mm256_add_epi32(sumi[iy], _mm256_add_epi32(p1, p3)); ++// sumi[iy] = _mm256_add_epi32(sumi[iy], _mm256_add_epi32(p2, p4)); ++// } ++//#endif ++//} ++ ++template ++static void mul_mat_qX_K_q8_K_IQ_N(int n, const void * vx, size_t bx, const DataInfo& info, int nrc_x) { ++ const int nb = n / QK_K; ++ Q8 q8(info); ++ Dequantizer deq(vx, bx); ++ __m256i scales[4]; ++ __m256 accd[nrc_y]; ++ ++ for (int ix = 0; ix < nrc_x; ++ix) { ++ ++ for (int iy = 0; iy < nrc_y; ++iy) accd[iy] = _mm256_setzero_ps(); ++ ++ deq.new_row(ix); ++ ++ for (int i = 0; i < nb; ++i) { ++ ++ __m256i sumi[nrc_y], all_scales[Dequantizer::num_blocks/8]; ++ //for (int iy = 0; iy < nrc_y; ++iy) sumi[iy] = _mm256_setzero_si256(); ++ __m256i mins; ++ float dmin = deq.new_block(i, all_scales, mins); ++ for (int iy = 0; iy < nrc_y; ++iy) { ++ auto bsums = q8.load_bsums(iy, i); ++ auto prod = _mm256_madd_epi16(mins, bsums); ++ accd[iy] = _mm256_fmadd_ps(_mm256_set1_ps(dmin*q8.scale(iy, i)), _mm256_cvtepi32_ps(prod), accd[iy]); ++ } ++ ++ for (int j = 0; j < QK_K/128; ++j) { ++ deq.prepare(i, j); ++ if constexpr (Dequantizer::num_blocks == 8) { ++ set_scales_8(all_scales[0], j, scales); ++ } else { ++ set_scales_16(all_scales[j], scales); ++ } ++ //multiply_add_iq(deq.bits, scales, j, i, q8, sumi); ++ multiply_add(deq.bits, scales, j, i, q8, sumi); ++ } ++ for (int iy = 0; iy < nrc_y; ++iy) { ++ const __m256 vd = _mm256_set1_ps(deq.d*q8.scale(iy, i)); ++ accd[iy] = _mm256_fmadd_ps(vd, _mm256_cvtepi32_ps(sumi[iy]), accd[iy]); ++ } ++ } ++ ++ for (int iy = 0; iy < nrc_y; ++iy) { ++ info.store(ix, iy, hsum_float_8(accd[iy])); ++ } ++ } ++} ++ ++template ++static void mul_mat_qX_K_q8_K_IQ(int n, const void * vx, size_t bx, const DataInfo& info, int nrc_x) { ++ assert(n % QK_K == 0); ++#ifdef HAVE_FANCY_SIMD ++ if constexpr (nrc_y == 1) { ++ mul_mat_qX_K_q8_K_IQ_1(n, vx, bx, info, nrc_x); ++ } else { ++ mul_mat_qX_K_q8_K_IQ_N(n, vx, bx, info, nrc_x); ++ } ++#else ++ mul_mat_qX_K_q8_K_IQ_N(n, vx, bx, info, nrc_x); ++#endif ++} ++ ++template void set_functions(mul_mat_t * funcs) { ++ funcs[0] = mul_mat_qX_K_q8_K_IQ; ++ funcs[1] = mul_mat_qX_K_q8_K_IQ; ++ funcs[2] = mul_mat_qX_K_q8_K_IQ; ++ funcs[3] = mul_mat_qX_K_q8_K_IQ; ++ funcs[4] = mul_mat_qX_K_q8_K_IQ; ++ funcs[5] = mul_mat_qX_K_q8_K_IQ; ++ funcs[6] = mul_mat_qX_K_q8_K_IQ; ++ funcs[7] = mul_mat_qX_K_q8_K_IQ; ++} ++ ++} // namespace ++ ++bool iqk_set_kernels_iquants(int typeA, mul_mat_t * kernels) { ++ switch (typeA) { ++ case GGML_TYPE_IQ2_XXS: set_functions(kernels); break; ++ case GGML_TYPE_IQ2_XS: set_functions(kernels); break; ++ case GGML_TYPE_IQ2_S: set_functions(kernels); break; ++ case GGML_TYPE_IQ3_XXS: set_functions(kernels); break; ++ case GGML_TYPE_IQ3_S: set_functions(kernels); break; ++ default: return false; ++ } ++ return true; ++} ++ ++#else ++// --------------------------------- __aarch64__ -------------------------------------- ++ ++namespace { ++ ++inline int32x4x4_t make_wider(const int16x8x2_t& scales16) { ++ int32x4x4_t scales = { ++ vmovl_s16(vget_low_s16 (scales16.val[0])), ++ vmovl_s16(vget_high_s16(scales16.val[0])), ++ vmovl_s16(vget_low_s16 (scales16.val[1])), ++ vmovl_s16(vget_high_s16(scales16.val[1])), ++ }; ++ return scales; ++} ++ ++struct SimpleBits { ++ uint8x16x4_t b1; ++ uint8x16x4_t b2; ++}; ++ ++inline int32x4x2_t prepare_scales_8(const uint32x4_t& v1, const uint32x4_t& v2) { ++ int32x4x2_t scales; ++ scales.val[0] = vreinterpretq_s32_u32(vorrq_u32(vshlq_n_u32(vshrq_n_u32(v1, 28), 1), vdupq_n_u32(1))); ++ scales.val[1] = vreinterpretq_s32_u32(vorrq_u32(vshlq_n_u32(vshrq_n_u32(v2, 28), 1), vdupq_n_u32(1))); ++ return scales; ++} ++ ++inline void apply_signs_2(uint8x16_t * b, const uint64_t * signs, uint32_t sidx) { ++ auto s1 = vcombine_s8(vld1_s8((const int8_t *)(signs + ((sidx >> 0) & 127))), vld1_s8((const int8_t *)(signs + ((sidx >> 7) & 127)))); ++ auto s2 = vcombine_s8(vld1_s8((const int8_t *)(signs + ((sidx >>14) & 127))), vld1_s8((const int8_t *)(signs + ((sidx >>21) & 127)))); ++ b[0] = vreinterpretq_u8_s8(vmulq_s8(vreinterpretq_s8_u8(b[0]), s1)); ++ b[1] = vreinterpretq_u8_s8(vmulq_s8(vreinterpretq_s8_u8(b[1]), s2)); ++} ++ ++struct DequantizerIQ2XXS final : public BaseDequantizer { ++ DequantizerIQ2XXS(const void * vx, size_t bx, int nrc) : BaseDequantizer(vx, bx, nrc) {} ++ ++ constexpr static int num_blocks() { return 8; } ++ constexpr static bool should_scale_quants() { return false; } ++ ++ template ++ inline int32x4x2_t new_block(int i, const Q8& /*q8*/, float32x4_t * /*acc*/) { ++ d = 0.125f * GGML_FP16_TO_FP32(x[i].d); ++ ++ auto tmp = vld1q_u32_x4((const uint32_t *)x[i].qs); ++ data.val[0] = vuzp1q_u32(tmp.val[0], tmp.val[1]); // codebook indices for blocks 0...3 ++ data.val[1] = vuzp2q_u32(tmp.val[0], tmp.val[1]); // scales and signs for blocks 0...3 ++ data.val[2] = vuzp1q_u32(tmp.val[2], tmp.val[3]); // codebook indices for blocks 4...7 ++ data.val[3] = vuzp2q_u32(tmp.val[2], tmp.val[3]); // scales and signs for blocks 4...7 ++ ++ return prepare_scales_8(data.val[1], data.val[3]); ++ } ++ ++ static inline void prepare2(uint8x16_t * b, const uint8_t * idx, const uint64_t * signs, uint32_t sidx) { ++ b[0] = vreinterpretq_u8_u64(uint64x2_t{iq2xxs_grid[idx[0]], iq2xxs_grid[idx[1]]}); ++ b[1] = vreinterpretq_u8_u64(uint64x2_t{iq2xxs_grid[idx[2]], iq2xxs_grid[idx[3]]}); ++ apply_signs_2(b, signs, sidx); ++ } ++ ++ inline void prepare(int /*i*/, int j) { ++ const uint8_t * idx = (const uint8_t *)(data.val + 2*j); ++ const uint32_t * sidx = (const uint32_t *)(data.val + 2*j+1); ++ prepare2(bits.b1.val + 0, idx, keven_signs, sidx[0]); idx += 4; ++ prepare2(bits.b1.val + 2, idx, keven_signs, sidx[1]); idx += 4; ++ prepare2(bits.b2.val + 0, idx, keven_signs, sidx[2]); idx += 4; ++ prepare2(bits.b2.val + 2, idx, keven_signs, sidx[3]); ++ } ++ ++ uint32x4x4_t data; ++ SimpleBits bits; ++ ++}; ++ ++inline int32x4x4_t prepare_4bit_scales16(const uint8_t * sc) { ++ auto aux = vld1_u8(sc); ++ auto scales_l = vand_u8(aux, vdup_n_u8(0xf)); ++ auto scales_h = vshr_n_u8(aux, 4); ++ auto aux1 = vcombine_u8(vzip1_u8(scales_l, scales_h), vzip2_u8(scales_l, scales_h)); ++ ++ auto scales8 = vreinterpretq_s8_u8(vorrq_u8(vshlq_n_u8(aux1, 1), vdupq_n_u8(1))); ++ int16x8x2_t scales16 = { vmovl_s8(vget_low_s8(scales8)), vmovl_s8(vget_high_s8(scales8)) }; ++ return make_wider(scales16); ++} ++ ++struct DequantizerIQ2XS final : public BaseDequantizer { ++ DequantizerIQ2XS(const void * vx, size_t bx, int nrc) : BaseDequantizer(vx, bx, nrc) {} ++ ++ constexpr static int num_blocks() { return 16; } ++ constexpr static bool should_scale_quants() { return false; } ++ ++ template ++ inline int32x4x4_t new_block(int i, const Q8& /*q8*/, float32x4_t * /*acc*/) { ++ d = 0.125f * GGML_FP16_TO_FP32(x[i].d); ++ return prepare_4bit_scales16(x[i].scales); ++ } ++ ++ inline static uint8x16_t make1(const uint16_t * qs) { ++ auto b = vcombine_u8(vld1_u8((const uint8_t *)(iq2xs_grid + (qs[0] & 511))), vld1_u8((const uint8_t *)(iq2xs_grid + (qs[1] & 511)))); ++ auto s = vcombine_s8(vld1_s8((const int8_t *)(keven_signs + (qs[0] >> 9))), vld1_s8((const int8_t *)(keven_signs + (qs[1] >> 9)))); ++ return vreinterpretq_u8_s8(vmulq_s8(vreinterpretq_s8_u8(b), s)); ++ } ++ ++ inline static void make4(const uint16_t * qs, uint8x16_t * b) { ++ b[0] = make1(qs + 0); ++ b[1] = make1(qs + 2); ++ b[2] = make1(qs + 4); ++ b[3] = make1(qs + 6); ++ } ++ ++ inline void prepare(int i, int j) { ++ make4(x[i].qs + 16*j + 0, bits.b1.val); ++ make4(x[i].qs + 16*j + 8, bits.b2.val); ++ } ++ ++ SimpleBits bits; ++ ++ ++}; ++ ++struct DequantizerIQ2S final : public BaseDequantizer { ++ DequantizerIQ2S(const void * vx, size_t bx, int nrc) : BaseDequantizer(vx, bx, nrc) {} ++ ++ constexpr static int num_blocks() { return 16; } ++ constexpr static bool should_scale_quants() { return false; } ++ ++ template ++ inline int32x4x4_t new_block(int i, const Q8& /*q8*/, float32x4_t * /*acc*/) { ++ d = 0.125f * GGML_FP16_TO_FP32(x[i].d); ++ return prepare_4bit_scales16(x[i].scales); ++ } ++ ++ static inline void make4(SignHelper& sh, const uint8x16_t& signs16, const uint8_t * qs, const uint8_t * qh, uint8x16_t * b) { ++ uint32_t aux32[2]; ++ const uint16_t * aux16 = (const uint16_t *)aux32; ++ for (int k = 0; k < 2; ++k) { ++ aux32[1] = (qh[k] << 4) | (qh[k] << 18); ++ aux32[0] = (aux32[1] << 4) & 0x03000300; ++ aux32[1] &= 0x03000300; ++ b[2*k+0] = vcombine_u8(vld1_u8((const uint8_t *)(iq2s_grid + (qs[4*k+0] | aux16[0]))), ++ vld1_u8((const uint8_t *)(iq2s_grid + (qs[4*k+1] | aux16[1])))); ++ sh.apply_signs_1(b+2*k+0, signs16); ++ ++ b[2*k+1] = vcombine_u8(vld1_u8((const uint8_t *)(iq2s_grid + (qs[4*k+2] | aux16[2]))), ++ vld1_u8((const uint8_t *)(iq2s_grid + (qs[4*k+3] | aux16[3])))); ++ sh.apply_signs_1(b+2*k+1, signs16); ++ } ++ } ++ ++ inline void prepare(int i, int j) { ++ ++ const auto * qs = x[i].qs + 16*j; ++ const auto * qh = x[i].qh + 4*j; ++ const auto signs16 = vld1q_u8(qs + QK_K/8); ++ ++ sh.init(); ++ make4(sh, signs16, qs+0, qh+0, bits.b1.val); ++ make4(sh, signs16, qs+8, qh+2, bits.b2.val); ++ } ++ ++ SimpleBits bits; ++ SignHelper sh; ++ ++ ++}; ++ ++struct DequantizerIQ3XXS final : public BaseDequantizer { ++ DequantizerIQ3XXS(const void * vx, size_t bx, int nrc) : BaseDequantizer(vx, bx, nrc) {} ++ ++ constexpr static int num_blocks() { return 8; } ++ constexpr static bool should_scale_quants() { return false; } ++ ++ template ++ inline int32x4x2_t new_block(int i, const Q8& /*q8*/, float32x4_t * /*acc*/) { ++ d = 0.25f * GGML_FP16_TO_FP32(x[i].d); ++ gas = vld1q_u32_x2((const uint32_t *)(x[i].qs + QK_K/4)); ++ return prepare_scales_8(gas.val[0], gas.val[1]); ++ } ++ ++ inline static void make2(const uint8_t * q3, uint32_t sidx, uint8x16_t * b) { ++ b[0] = vreinterpretq_u8_u32(uint32x4_t{iq3xxs_grid[q3[0]], iq3xxs_grid[q3[1]], iq3xxs_grid[q3[2]], iq3xxs_grid[q3[3]]}); ++ b[1] = vreinterpretq_u8_u32(uint32x4_t{iq3xxs_grid[q3[4]], iq3xxs_grid[q3[5]], iq3xxs_grid[q3[6]], iq3xxs_grid[q3[7]]}); ++ apply_signs_2(b, keven_signs, sidx); ++ } ++ inline void prepare(int i, int j) { ++ const auto * q3 = x[i].qs + 32*j; ++ const auto * signs = (const uint32_t *)(gas.val + j); ++ make2(q3, signs[0], bits.b1.val + 0); q3 += 8; ++ make2(q3, signs[1], bits.b1.val + 2); q3 += 8; ++ make2(q3, signs[2], bits.b2.val + 0); q3 += 8; ++ make2(q3, signs[3], bits.b2.val + 2); ++ } ++ ++ SimpleBits bits; ++ uint32x4x2_t gas; ++ ++}; ++ ++struct DequantizerIQ3S final : public BaseDequantizer { ++ DequantizerIQ3S(const void * vx, size_t bx, int nrc) : BaseDequantizer(vx, bx, nrc) {} ++ ++ constexpr static int num_blocks() { return 8; } ++ constexpr static bool should_scale_quants() { return false; } ++ ++ template ++ inline int32x4x2_t new_block(int i, const Q8& /*q8*/, float32x4_t * /*acc*/) { ++ d = GGML_FP16_TO_FP32(x[i].d); ++ uint32_t scales32[2]; ++ std::memcpy(scales32, x[i].scales, 4); ++ scales32[1] = (((scales32[0] >> 4) & 0x0f0f0f0f) << 1) | 0x01010101; ++ scales32[0] = ((scales32[0] & 0x0f0f0f0f) << 1) | 0x01010101; ++ auto scales8 = vld1_u8((const uint8_t *)scales32); // 0, 2, 4, 6, 1, 3, 5, 7 ++ scales8 = vtbl1_u8(scales8, vreinterpret_u8_u64(vdup_n_u64(0x0703060205010400))); ++ auto scales16 = vreinterpretq_s16_u16(vmovl_u8(scales8)); ++ int32x4x2_t scales; ++ scales.val[0] = vmovl_s16(vget_low_s16(scales16)); ++ scales.val[1] = vmovl_s16(vget_high_s16(scales16)); ++ return scales; ++ } ++ ++ static inline void make2(SignHelper& sh, const uint8x16_t& signs16, const uint16x8_t& idx_l, uint8_t qh, ++ const int16x8_t& hshift, uint8x16_t * b) { ++ auto vindex = vorrq_u16(idx_l, vandq_u16(vshlq_u16(vdupq_n_u16(qh), hshift), vdupq_n_u16(256))); ++ const uint16_t * idx = (const uint16_t *)&vindex; ++ b[0] = vreinterpretq_u8_u32(uint32x4_t{iq3s_grid[idx[0]], iq3s_grid[idx[1]], iq3s_grid[idx[2]], iq3s_grid[idx[3]]}); ++ b[1] = vreinterpretq_u8_u32(uint32x4_t{iq3s_grid[idx[4]], iq3s_grid[idx[5]], iq3s_grid[idx[6]], iq3s_grid[idx[7]]}); ++ sh.apply_signs_1(b+0, signs16); ++ sh.apply_signs_1(b+1, signs16); ++ } ++ static inline void make4(SignHelper& sh, const uint8x16_t& signs16, const uint8_t * qs, const uint8_t * qh, ++ const int16x8_t& hshift, uint8x16_t * b) { ++ auto idx_l = vld1q_u8(qs); ++ make2(sh, signs16, vmovl_u8(vget_low_u8 (idx_l)), qh[0], hshift, b+0); ++ make2(sh, signs16, vmovl_u8(vget_high_u8(idx_l)), qh[1], hshift, b+2); ++ } ++ ++ inline void prepare(int i, int j) { ++ ++ static const int16_t k_shift[8] = {8, 7, 6, 5, 4, 3, 2, 1}; ++ const auto hshift = vld1q_s16(k_shift); ++ ++ const auto * qs = x[i].qs + 32*j; ++ const auto * qh = x[i].qh + 4*j; ++ const auto signs16 = vld1q_u8(x[i].signs + 16*j); ++ ++ sh.init(); ++ make4(sh, signs16, qs+ 0, qh+0, hshift, bits.b1.val); ++ make4(sh, signs16, qs+16, qh+2, hshift, bits.b2.val); ++ } ++ ++ SimpleBits bits; ++ SignHelper sh; ++ uint32x4x2_t gas; ++ ++}; ++ ++} // namespace ++ ++bool iqk_set_kernels_iquants(int typeA, mul_mat_t * kernels) { ++ switch (typeA) { ++ case GGML_TYPE_IQ2_XXS: IQK_SET_MUL_MAT_FUNCTIONS_T(mul_mat_qX_K_q8_K_T, DequantizerIQ2XXS, kernels); break; ++ case GGML_TYPE_IQ2_XS: IQK_SET_MUL_MAT_FUNCTIONS_T(mul_mat_qX_K_q8_K_T, DequantizerIQ2XS, kernels); break; ++ case GGML_TYPE_IQ2_S: IQK_SET_MUL_MAT_FUNCTIONS_T(mul_mat_qX_K_q8_K_T, DequantizerIQ2S, kernels); break; ++ case GGML_TYPE_IQ3_XXS: IQK_SET_MUL_MAT_FUNCTIONS_T(mul_mat_qX_K_q8_K_T, DequantizerIQ3XXS, kernels); break; ++ case GGML_TYPE_IQ3_S: IQK_SET_MUL_MAT_FUNCTIONS_T(mul_mat_qX_K_q8_K_T, DequantizerIQ3S, kernels); break; ++ default: return false; ++ } ++ return true; ++} ++ ++#endif ++ ++#endif // IQK_IMPLEMENT +diff --git a/src/ggml-cpu/iqk/iqk_gemm_kquants.cpp b/src/ggml-cpu/iqk/iqk_gemm_kquants.cpp +new file mode 100644 +index 0000000..916de2a +--- /dev/null ++++ b/src/ggml-cpu/iqk/iqk_gemm_kquants.cpp +@@ -0,0 +1,1162 @@ ++// ++// Copyright (C) 2024-2025 Iwan Kawrakow ++// MIT license ++// SPDX-License-Identifier: MIT ++// ++// Adapted from ik_llama.cpp (https://github.com/ikawrakow/ik_llama.cpp) commit 009dce19, ++// ggml/src/iqk/{iqk_config.h, iqk_common.h, iqk_gemm_kquants.cpp, iqk_gemm_iquants.cpp, iqk_gemm_1bit.cpp, iqk_mul_mat.cpp}. ++// Only the kernels that read the upstream ggml block layouts with upstream Q8_K activations are kept. ++// ++// Permission is hereby granted, free of charge, to any person obtaining a copy ++// of this software and associated documentation files (the "Software"), to deal ++// in the Software without restriction, including without limitation the rights ++// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell ++// copies of the Software, and to permit persons to whom the Software is ++// furnished to do so, subject to the following conditions: ++// ++// The above copyright notice and this permission notice shall be included in all ++// copies or substantial portions of the Software. ++// ++// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR ++// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, ++// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE ++// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER ++// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, ++// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE ++// SOFTWARE. ++// ++// k-quants (Q2_K, Q3_K, Q4_K, Q5_K, Q6_K) and IQ4_XS times Q8_K (from iqk_gemm_kquants.cpp). ++ ++#include "iqk_common.h" ++ ++#ifdef IQK_IMPLEMENT ++ ++#ifdef __x86_64__ ++ ++namespace { ++ ++// Handles q4_K and q5_K scales/mins ++struct Scales8K { ++ template ++ inline __m256i process_mins_and_scales(const uint8_t * data, float c, int i, const Q8& q8, __m256 * accd) { ++ make_q4_scales(data, utmp); ++ const __m256i mins_and_scales = _mm256_cvtepu8_epi16(_mm_set_epi32(utmp[3], utmp[2], utmp[1], utmp[0])); ++ const __m128i mins128 = _mm256_extracti128_si256(mins_and_scales, 1); ++ accum_mins(mins128, q8, i, c, accd); ++ const __m128i sc128 = _mm256_extracti128_si256(mins_and_scales, 0); ++ return MM256_SET1_M128I(sc128); ++ } ++#ifdef HAVE_FANCY_SIMD ++ template ++ inline __m512i process_mins_and_scales_64(const uint8_t * data, float c, int i, const Q8& q8, __m256 * accd) { ++ auto scales = process_mins_and_scales(data, c, i, q8, accd); ++ return _mm512_inserti32x8(_mm512_castsi256_si512(scales), scales, 1); ++ } ++#endif ++ template ++ inline void accum_mins(const __m128i& mins128, const Q8& q8, int i, float c, __m256 * accd) const { ++ base.accum_mins(mins128, q8, i, c, accd); ++ } ++#ifdef HAVE_FANCY_SIMD ++ const __m512i shuffles512[2] = { ++ _mm512_set_epi64(0x0706070607060706, 0x0302030203020302, 0x0706070607060706, 0x0302030203020302, ++ 0x0504050405040504, 0x0100010001000100, 0x0504050405040504, 0x0100010001000100), ++ _mm512_set_epi64(0x0f0e0f0e0f0e0f0e, 0x0b0a0b0a0b0a0b0a, 0x0f0e0f0e0f0e0f0e, 0x0b0a0b0a0b0a0b0a, ++ 0x0d0c0d0c0d0c0d0c, 0x0908090809080908, 0x0d0c0d0c0d0c0d0c, 0x0908090809080908) ++ }; ++#endif ++ Scales8KBase base; ++ ++ uint32_t utmp[4]; ++}; ++ ++template ++inline void process_mins_16(const __m256i& all_scales, const Q8& q8, int i, float d, __m256 * accm) { ++ for (int iy = 0; iy < Q8::nrc_y; ++iy) { ++ const __m256i prod = _mm256_madd_epi16(all_scales, q8.load_bsums(iy, i)); ++ accm[iy] = _mm256_fmadd_ps(_mm256_set1_ps(d * q8.scale(iy, i)), _mm256_cvtepi32_ps(prod), accm[iy]); ++ } ++} ++inline void prepare_scales_16(const __m256i& all_scales, __m256i * scales) { ++ const __m128i l_scales = _mm256_extracti128_si256(all_scales, 0); ++ const __m128i h_scales = _mm256_extracti128_si256(all_scales, 1); ++ scales[0] = MM256_SET1_M128I(l_scales); ++ scales[1] = MM256_SET1_M128I(h_scales); ++} ++ ++// Handles q3_K scales ++struct ScaleQ3 { ++ inline __m128i make_scales(const uint16_t * s8) const { ++ const uint16_t * scales16 = (const uint16_t *)s8; ++ uint32_t aux0 = scales16[0] | (scales16[1] << 16); ++ uint32_t aux1 = scales16[2] | (scales16[3] << 16); ++ uint32_t aux2 = scales16[4] | (scales16[5] << 16); ++ __m128i scales128 = _mm_set_epi32( ++ ((aux1 >> 4) & 0x0f0f0f0f) | ((aux2 >> 2) & 0x30303030), ++ ((aux0 >> 4) & 0x0f0f0f0f) | ((aux2 >> 0) & 0x30303030), ++ (aux1 & 0x0f0f0f0f) | ((aux2 << 2) & 0x30303030), ++ (aux0 & 0x0f0f0f0f) | ((aux2 << 4) & 0x30303030)); ++ return _mm_add_epi8(scales128, m32); ++ } ++ const __m128i m32 = _mm_set1_epi8(-32); ++}; ++ ++struct Scale16 { ++ inline void make_scales(const __m128i& scales8, __m512i * scales) const { ++ auto all_scales8 = MM256_SET1_M128I(scales8); ++ auto scales1 = _mm256_shuffle_epi8(all_scales8, shuffle1); ++ auto scales2 = _mm256_shuffle_epi8(all_scales8, shuffle2); ++ scales[0] = _mm512_cvtepi8_epi16(scales1); ++ scales[1] = _mm512_cvtepi8_epi16(scales2); ++ } ++ template ++ inline void process_mins_and_scales(int i, float c, const __m128i& mins8, const __m128i& scales8, ++ const Q8& q8, __m256 * accm, __m512i * scales) const { ++ process_mins_16(_mm256_cvtepi8_epi16(mins8), q8, i, c, accm); ++ make_scales(scales8, scales); ++ } ++ const __m256i shuffle1 = _mm256_set_epi32(0x07070707, 0x03030303, 0x06060606, 0x02020202, ++ 0x05050505, 0x01010101, 0x04040404, 0x00000000); ++ const __m256i shuffle2 = _mm256_set_epi32(0x0f0f0f0f, 0x0b0b0b0b, 0x0e0e0e0e, 0x0a0a0a0a, ++ 0x0d0d0d0d, 0x09090909, 0x0c0c0c0c, 0x08080808); ++}; ++ ++template ++inline void process_mins_and_scales_16(const __m128i& scales128, const Q8& q8, int i, float d, ++ __m256 * accm, __m256i * scales) { ++ const __m256i all_scales = _mm256_cvtepi8_epi16(scales128); ++ process_mins_16(all_scales, q8, i, d, accm); ++ prepare_scales_16(all_scales, scales); ++} ++ ++inline __m256i get_scale_shuffle_8(int i) { ++ return _mm256_set1_epi16((2*i) | ((2*i+1) << 8)); ++} ++ ++inline void set_scales_8(const __m256i& all_scales, int j, __m256i * scales) { ++ scales[0] = _mm256_shuffle_epi8(all_scales, get_scale_shuffle_8(4*j+0)); ++ scales[1] = _mm256_shuffle_epi8(all_scales, get_scale_shuffle_8(4*j+1)); ++ scales[2] = _mm256_shuffle_epi8(all_scales, get_scale_shuffle_8(4*j+2)); ++ scales[3] = _mm256_shuffle_epi8(all_scales, get_scale_shuffle_8(4*j+3)); ++} ++ ++inline __m256i get_scale_shuffle_16(int i) { ++ static const uint8_t k_shuffle[128] = { ++ 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 2, 3, 2, 3, 2, 3, 2, 3, 2, 3, 2, 3, 2, 3, 2, 3, ++ 4, 5, 4, 5, 4, 5, 4, 5, 4, 5, 4, 5, 4, 5, 4, 5, 6, 7, 6, 7, 6, 7, 6, 7, 6, 7, 6, 7, 6, 7, 6, 7, ++ 8, 9, 8, 9, 8, 9, 8, 9, 8, 9, 8, 9, 8, 9, 8, 9, 10,11,10,11,10,11,10,11,10,11,10,11,10,11,10,11, ++ 12,13,12,13,12,13,12,13,12,13,12,13,12,13,12,13, 14,15,14,15,14,15,14,15,14,15,14,15,14,15,14,15, ++ }; ++ return _mm256_loadu_si256((const __m256i*)k_shuffle + i); ++} ++ ++inline void set_scales_16(const __m256i& all_scales, __m256i * scales) { ++ scales[0] = _mm256_shuffle_epi8(all_scales, get_scale_shuffle_16(0)); ++ scales[1] = _mm256_shuffle_epi8(all_scales, get_scale_shuffle_16(1)); ++ scales[2] = _mm256_shuffle_epi8(all_scales, get_scale_shuffle_16(2)); ++ scales[3] = _mm256_shuffle_epi8(all_scales, get_scale_shuffle_16(3)); ++} ++ ++struct ScaleIQ4XS { ++ inline __m128i make_scales(const uint32_t scales_l, const uint16_t scales_h) { ++ uint32_t tmp32 = scales_h | (scales_h << 14); ++ const __m128i sh = _mm_slli_epi16(_mm_and_si128(_mm_srlv_epi32(_mm_set1_epi32(tmp32), hshift), hmask), 4); ++ const __m128i sl = _mm_and_si128(_mm_srlv_epi32(_mm_set1_epi32(scales_l), lshift), lmask); ++ return _mm_add_epi16(_mm_or_si128(sh, _mm_cvtepi8_epi16(_mm_shuffle_epi8(sl, lshuffle))), m32); ++ } ++ const __m128i hshift = _mm_set_epi32(12, 8, 4, 0); ++ const __m128i lshift = _mm_set_epi32(4, 0, 4, 0); ++ const __m128i hmask = _mm_set1_epi16(0x03); ++ const __m128i lmask = _mm_set1_epi8(0xf); ++ const __m128i lshuffle = _mm_set_epi32(0x07030602, 0x05010400, 0x07030602, 0x05010400); ++ const __m128i m32 = _mm_set1_epi16(-32); ++}; ++ ++#ifdef HAVE_FANCY_SIMD ++//====================================== Zen4 ================================================== ++ ++struct HighBit5 { ++ inline void apply(const uint8_t * h, Q4Bits& bits) { ++ auto hbits256 = _mm256_loadu_si256((const __m256i *)h); ++ auto hbits = _mm512_inserti32x8(_mm512_castsi256_si512(hbits256), _mm256_srli_epi16(hbits256, 1), 1); ++ bits.values[0] = _mm512_or_si512(bits.values[0], _mm512_and_si512(_mm512_slli_epi16(hbits, 4), mh)); ++ bits.values[1] = _mm512_or_si512(bits.values[1], _mm512_and_si512(_mm512_slli_epi16(hbits, 2), mh)); ++ bits.values[2] = _mm512_or_si512(bits.values[2], _mm512_and_si512(hbits, mh)); ++ bits.values[3] = _mm512_or_si512(bits.values[3], _mm512_and_si512(_mm512_srli_epi16(hbits, 2), mh)); ++ } ++ const __m512i mh = _mm512_set1_epi8(0x10); ++}; ++ ++struct HighBit3 { ++ inline void apply(const uint8_t * h, Q2Bits& bits) { ++ auto hbits256 = _mm256_loadu_si256((const __m256i *)h); ++ auto hbits = _mm512_inserti32x8(_mm512_castsi256_si512(hbits256), _mm256_srli_epi16(hbits256, 1), 1); ++ bits.values[0] = _mm512_or_si512(bits.values[0], _mm512_and_si512(_mm512_slli_epi16(hbits, 2), mh)); ++ bits.values[1] = _mm512_or_si512(bits.values[1], _mm512_and_si512(hbits, mh)); ++ bits.values[2] = _mm512_or_si512(bits.values[2], _mm512_and_si512(_mm512_srli_epi16(hbits, 2), mh)); ++ bits.values[3] = _mm512_or_si512(bits.values[3], _mm512_and_si512(_mm512_srli_epi16(hbits, 4), mh)); ++ } ++ const __m512i mh = _mm512_set1_epi8(0x04); ++}; ++ ++ ++template ++inline void compute_block(int iy, int i, float d, const Q8& q8, const __m512i * values, const __m512i * scales, __m512 * accd) { ++ const __m512i p1 = _mm512_dpbusd_epi32(_mm512_setzero_si512(), values[0], q8.load_quants64(iy, i, 0)); ++ const __m512i p2 = _mm512_dpbusd_epi32(_mm512_setzero_si512(), values[1], q8.load_quants64(iy, i, 1)); ++ const __m512i p3 = _mm512_dpbusd_epi32(_mm512_setzero_si512(), values[2], q8.load_quants64(iy, i, 2)); ++ const __m512i p4 = _mm512_dpbusd_epi32(_mm512_setzero_si512(), values[3], q8.load_quants64(iy, i, 3)); ++ auto sumi = _mm512_dpwssd_epi32(_mm512_setzero_si512(), scales[0], _mm512_packs_epi32(p1, p2)); ++ sumi = _mm512_dpwssd_epi32(sumi, scales[1], _mm512_packs_epi32(p3, p4)); ++ accd[iy] = _mm512_fmadd_ps(_mm512_set1_ps(d*q8.scale(iy, i)), _mm512_cvtepi32_ps(sumi), accd[iy]); ++} ++ ++struct DequantizerQ2K final : public BaseDequantizer { ++ DequantizerQ2K(const void * vx, size_t bx) : BaseDequantizer(vx, bx) {} ++ template ++ inline void new_block(int i, const Q8& q8, __m256 * accm, __m512i * scales) { ++ d = GGML_FP16_TO_FP32(x[i].data.data.d); ++ bits.prepare(x[i].qs); ++ const __m128i mins_and_scales = _mm_loadu_si128((const __m128i*)x[i].scales); ++ const __m128i scales8 = _mm_and_si128(mins_and_scales, m4); ++ const __m128i mins8 = _mm_and_si128(_mm_srli_epi16(mins_and_scales, 4), m4); ++ sc16.process_mins_and_scales(i, -GGML_FP16_TO_FP32(x[i].data.data.dmin), mins8, scales8, q8, accm, scales); ++ } ++ ++ Q2Bits bits; ++ Scale16 sc16; ++ const __m128i m4 = _mm_set1_epi8(0xf); ++ ++}; ++ ++struct DequantizerQ3K final : public BaseDequantizer { ++ DequantizerQ3K(const void * vx, size_t bx) : BaseDequantizer(vx, bx) {} ++ template ++ inline void new_block(int i, const Q8& q8, __m256 * accm, __m512i * scales) { ++ d = GGML_FP16_TO_FP32(x[i].d); ++ bits.prepare(x[i].qs); ++ hbits.apply(x[i].hmask, bits); ++ auto scales128 = sc3.make_scales((const uint16_t *)x[i].scales); ++ sc16.process_mins_and_scales(i, -4.f*d, scales128, scales128, q8, accm, scales); ++ } ++ ++ Q2Bits bits; ++ HighBit3 hbits; ++ ScaleQ3 sc3; ++ Scale16 sc16; ++ const __m128i m4 = _mm_set1_epi8(0xf); ++ const __m128i m32 = _mm_set1_epi8(-32); ++}; ++ ++struct DequantizerQ4K final : public BaseDequantizer { ++ DequantizerQ4K(const void * vx, size_t bx) : BaseDequantizer(vx, bx) {} ++ template ++ inline void new_block(int i, const Q8& q8, __m256 * accd, __m512i * scales) { ++ d = GGML_FP16_TO_FP32(x[i].data.data.d); ++ bits.prepare(x[i].qs); ++ auto all_scales = s8k.process_mins_and_scales_64(x[i].scales, -GGML_FP16_TO_FP32(x[i].data.data.dmin), i, q8, accd); ++ scales[0] = _mm512_shuffle_epi8(all_scales, s8k.shuffles512[0]); ++ scales[1] = _mm512_shuffle_epi8(all_scales, s8k.shuffles512[1]); ++ } ++ ++ Q4Bits bits; ++ Scales8K s8k; ++}; ++ ++struct DequantizerQ5K final : public BaseDequantizer { ++ DequantizerQ5K(const void * vx, size_t bx) : BaseDequantizer(vx, bx) {} ++ template ++ inline void new_block(int i, const Q8& q8, __m256 * accd, __m512i * scales) { ++ d = GGML_FP16_TO_FP32(x[i].data.data.d); ++ bits.prepare(x[i].qs); ++ hbits.apply(x[i].qh, bits); ++ auto all_scales = s8k.process_mins_and_scales_64(x[i].scales, -GGML_FP16_TO_FP32(x[i].data.data.dmin), i, q8, accd); ++ scales[0] = _mm512_shuffle_epi8(all_scales, s8k.shuffles512[0]); ++ scales[1] = _mm512_shuffle_epi8(all_scales, s8k.shuffles512[1]); ++ } ++ ++ Q4Bits bits; ++ HighBit5 hbits; ++ Scales8K s8k; ++}; ++ ++struct DequantizerQ6K final : public BaseDequantizer { ++ DequantizerQ6K(const void * vx, size_t bx) : BaseDequantizer(vx, bx) {} ++ template ++ inline void new_block(int i, const Q8& q8, __m256 * accm, __m512i * scales) { ++ d = GGML_FP16_TO_FP32(x[i].d); ++ bits.prepare64(x[i].ql); ++ add_high_bits(x[i].qh, bits); ++ auto scales128 = _mm_loadu_si128((const __m128i *)x[i].scales); ++ sc16.process_mins_and_scales(i, -32.f*d, scales128, scales128, q8, accm, scales); ++ } ++ ++ inline void add_high_bits(const uint8_t * qh, Q4Bits& bits) const { ++ auto hbits = _mm512_loadu_si512((const __m512i *)qh); ++ auto tmp1 = _mm512_and_si512(_mm512_slli_epi16(hbits, 4), mh); ++ auto tmp2 = _mm512_and_si512(_mm512_slli_epi16(hbits, 2), mh); ++ bits.values[0] = _mm512_or_si512(bits.values[0], _mm512_permutex2var_epi64(tmp1, bits.perm.permute1, tmp2)); ++ bits.values[2] = _mm512_or_si512(bits.values[2], _mm512_permutex2var_epi64(tmp1, bits.perm.permute2, tmp2)); ++ tmp1 = _mm512_and_si512(hbits, mh); ++ tmp2 = _mm512_and_si512(_mm512_srli_epi16(hbits, 2), mh); ++ bits.values[1] = _mm512_or_si512(bits.values[1], _mm512_permutex2var_epi64(tmp1, bits.perm.permute1, tmp2)); ++ bits.values[3] = _mm512_or_si512(bits.values[3], _mm512_permutex2var_epi64(tmp1, bits.perm.permute2, tmp2)); ++ } ++ ++ Q4Bits bits; ++ HighBit3 hbits; ++ Scale16 sc16; ++ ++ const __m512i mh = _mm512_set1_epi8(0x30); ++ ++}; ++ ++struct DequantizerIQ4XS final : public BaseDequantizer { ++ DequantizerIQ4XS(const void * vx, size_t bx) : BaseDequantizer(vx, bx), values(load_iq4nl_values_512()) {} ++ inline __m512 new_block(int i, __m512i& mins) { ++ d = GGML_FP16_TO_FP32(x[i].d); ++ prepare(x[i].qs); ++ auto scales128 = siq4.make_scales(*(const uint32_t *)x[i].scales_l, x[i].scales_h); ++ mins = dpbusd_permute_epi16(_mm256_set1_epi16(-128)); ++ return dpbusd_scales(d, s8k.shuffle(scales128)); ++ } ++ inline void prepare(const uint8_t * q4) { ++ bits.prepare64(q4); ++ // We now have in bits.valuse[0]: 0...15, 32...47, 64...79, 96...111 ++ // bits.valuse[1]: 16..31, 48...63, 80...95, 112..127 ++ // etc. ++ auto tmp = _mm512_permutex2var_epi64(bits.values[0], permute1, bits.values[1]); ++ bits.values[1] = _mm512_shuffle_epi8(values, _mm512_permutex2var_epi64(bits.values[0], permute2, bits.values[1])); ++ bits.values[0] = _mm512_shuffle_epi8(values, tmp); ++ tmp = _mm512_permutex2var_epi64(bits.values[2], permute1, bits.values[3]); ++ bits.values[3] = _mm512_shuffle_epi8(values, _mm512_permutex2var_epi64(bits.values[2], permute2, bits.values[3])); ++ bits.values[2] = _mm512_shuffle_epi8(values, tmp); ++ } ++ ++ Q4Bits bits; ++ Scales8KBase s8k; ++ ScaleIQ4XS siq4; ++ const __m512i values; ++ const __m512i permute1 = _mm512_set_epi64(11, 10, 3, 2, 9, 8, 1, 0); ++ const __m512i permute2 = _mm512_set_epi64(15, 14, 7, 6, 13, 12, 5, 4); ++}; ++ ++template ++static void mul_mat_qX_K_q8_K_AVX512_1(int n, const void * vx, size_t bx, const DataInfo& info, int nrc_x) { ++ assert(n % QK_K == 0); ++ const int nb = n / QK_K; ++ ++ constexpr int k_nx = 2; ++ ++ Q8<1> q8(info); ++ ++ Dequantizer deq1(vx, bx); ++ Dequantizer deq2(vx, bx); ++ ++ Dequantizer * deq[k_nx]; ++ deq[0] = &deq1; ++ deq[1] = &deq2; ++ ++ __m512i scales[2*k_nx]; ++ ++ for (int ix = 0; ix < nrc_x; ++ix) { ++ ++ auto accd = _mm512_setzero_ps(); ++ auto accm = _mm256_setzero_ps(); ++ ++ for (int kx = 0; kx < k_nx; ++kx) deq[kx]->new_row(ix); ++ ++ for (int i = 0; i < nb/k_nx; ++i) { ++ ++ for (int kx = 0; kx < k_nx; ++kx) deq[kx]->new_block(k_nx*i+kx, q8, &accm, scales+2*kx); ++ ++ for (int kx = 0; kx < k_nx; ++kx) { ++ compute_block(0, k_nx*i+kx, deq[kx]->d, q8, deq[kx]->bits.values, scales+2*kx, &accd); ++ } ++ ++ } ++ if (2*(nb/2) < nb) { ++ int i0 = 2*(nb/2); ++ deq[0]->new_block(i0, q8, &accm, scales); ++ compute_block(0, i0, deq[0]->d, q8, deq[0]->bits.values, scales, &accd); ++ } ++ ++ auto sum256 = _mm256_add_ps(_mm512_castps512_ps256(accd), _mm512_extractf32x8_ps(accd, 1)); ++ info.store(ix, 0, hsum_float_8(_mm256_add_ps(accm, sum256))); ++ } ++} ++ ++template ++static void mul_mat_qX_K_q8_K_AVX512(int n, const void * vx, size_t bx, const DataInfo& info, int nrc_x) { ++ assert(n % QK_K == 0); ++ const int nb = n / QK_K; ++ ++ Q8 q8(info); ++ ++ Dequantizer deq(vx, bx); ++ ++ __m256 accm[nrc_y]; ++ __m512 accd[nrc_y]; ++ __m512i scales[2]; ++ ++ for (int ix = 0; ix < nrc_x; ++ix) { ++ ++ for (int iy = 0; iy < nrc_y; ++iy) accd[iy] = _mm512_setzero_ps(); ++ for (int iy = 0; iy < nrc_y; ++iy) accm[iy] = _mm256_setzero_ps(); ++ ++ deq.new_row(ix); ++ ++ for (int i = 0; i < nb; ++i) { ++ ++ deq.new_block(i, q8, accm, scales); ++ ++ for (int iy = 0; iy < nrc_y; ++iy) { ++ const __m512i p1 = _mm512_dpbusd_epi32(_mm512_setzero_si512(), deq.bits.values[0], q8.load_quants64(iy, i, 0)); ++ const __m512i p2 = _mm512_dpbusd_epi32(_mm512_setzero_si512(), deq.bits.values[1], q8.load_quants64(iy, i, 1)); ++ const __m512i p3 = _mm512_dpbusd_epi32(_mm512_setzero_si512(), deq.bits.values[2], q8.load_quants64(iy, i, 2)); ++ const __m512i p4 = _mm512_dpbusd_epi32(_mm512_setzero_si512(), deq.bits.values[3], q8.load_quants64(iy, i, 3)); ++ auto sumi = _mm512_dpwssd_epi32(_mm512_setzero_si512(), scales[0], _mm512_packs_epi32(p1, p2)); ++ sumi = _mm512_dpwssd_epi32(sumi, scales[1], _mm512_packs_epi32(p3, p4)); ++ accd[iy] = _mm512_fmadd_ps(_mm512_set1_ps(deq.d*q8.scale(iy, i)), _mm512_cvtepi32_ps(sumi), accd[iy]); ++ } ++ ++ } ++ ++ for (int iy = 0; iy < nrc_y; ++iy) { ++ auto sum256 = _mm256_add_ps(_mm512_castps512_ps256(accd[iy]), _mm512_extractf32x8_ps(accd[iy], 1)); ++ info.store(ix, iy, hsum_float_8(_mm256_add_ps(accm[iy], sum256))); ++ } ++ ++ } ++} ++ ++#else ++//====================================== AVX2 ================================================== ++ ++struct HighBit5 { ++ inline void load(const uint8_t * h) { hbits = _mm256_loadu_si256((const __m256i *)h); } ++ inline void apply(Q4Bits& bits, bool do_shift) { ++ bits.values[0] = _mm256_or_si256(bits.values[0], _mm256_and_si256(_mm256_slli_epi16(hbits, 4), mh)); ++ bits.values[1] = _mm256_or_si256(bits.values[1], _mm256_and_si256(_mm256_slli_epi16(hbits, 3), mh)); ++ bits.values[2] = _mm256_or_si256(bits.values[2], _mm256_and_si256(_mm256_slli_epi16(hbits, 2), mh)); ++ bits.values[3] = _mm256_or_si256(bits.values[3], _mm256_and_si256(_mm256_slli_epi16(hbits, 1), mh)); ++ if (do_shift) { ++ hbits = _mm256_srli_epi16(hbits, 4); ++ } ++ } ++ const __m256i mh = _mm256_set1_epi8(0x10); ++ __m256i hbits; ++}; ++ ++struct HighBit3 { ++ inline void load(const uint8_t * h) { hbits = _mm256_loadu_si256((const __m256i *)h); } ++ inline void apply(Q2Bits& bits, bool do_shift) { ++ bits.values[0] = _mm256_or_si256(bits.values[0], _mm256_and_si256(_mm256_slli_epi16(hbits, 2), mh)); ++ bits.values[1] = _mm256_or_si256(bits.values[1], _mm256_and_si256(_mm256_slli_epi16(hbits, 1), mh)); ++ bits.values[2] = _mm256_or_si256(bits.values[2], _mm256_and_si256(hbits, mh)); ++ bits.values[3] = _mm256_or_si256(bits.values[3], _mm256_and_si256(_mm256_srli_epi16(hbits, 1), mh)); ++ if (do_shift) { ++ hbits = _mm256_srli_epi16(hbits, 4); ++ } ++ } ++ const __m256i mh = _mm256_set1_epi8(0x04); ++ __m256i hbits; ++}; ++ ++template ++inline void compute_block(int iy, int i, float d, const Q8& q8, const __m512i * values, const __m512i * scales, __m512 * accd) { ++ const __m512i p1 = _mm512_dpbusd_epi32(_mm512_setzero_si512(), values[0], q8.load_quants64(iy, i, 0)); ++ const __m512i p2 = _mm512_dpbusd_epi32(_mm512_setzero_si512(), values[1], q8.load_quants64(iy, i, 1)); ++ const __m512i p3 = _mm512_dpbusd_epi32(_mm512_setzero_si512(), values[2], q8.load_quants64(iy, i, 2)); ++ const __m512i p4 = _mm512_dpbusd_epi32(_mm512_setzero_si512(), values[3], q8.load_quants64(iy, i, 3)); ++ auto sumi = _mm512_dpwssd_epi32(_mm512_setzero_si512(), scales[0], _mm512_packs_epi32(p1, p2)); ++ sumi = _mm512_dpwssd_epi32(sumi, scales[1], _mm512_packs_epi32(p3, p4)); ++ accd[iy] = _mm512_fmadd_ps(_mm512_set1_ps(d*q8.scale(iy, i)), _mm512_cvtepi32_ps(sumi), accd[iy]); ++} ++ ++struct DequantizerQ2K final : public BaseDequantizer { ++ DequantizerQ2K(const void * vx, size_t bx) : BaseDequantizer(vx, bx) {} ++ ++ template ++ inline void new_block(int i, const Q8& q8, __m256 * accm, __m256i * scales) { ++ d = GGML_FP16_TO_FP32(x[i].data.data.d); ++ const __m128i mins_and_scales = _mm_loadu_si128((const __m128i*)x[i].scales); ++ const __m128i scales8 = _mm_and_si128(mins_and_scales, m4); ++ const __m128i mins8 = _mm_and_si128(_mm_srli_epi16(mins_and_scales, 4), m4); ++ process_mins_16(_mm256_cvtepi8_epi16(mins8), q8, i, -GGML_FP16_TO_FP32(x[i].data.data.dmin), accm); ++ prepare_scales_16(_mm256_cvtepi8_epi16(scales8), scales); ++ } ++ inline void prepare(int i, int j) { ++ bits.prepare(x[i].qs, j); ++ } ++ ++ Q2Bits bits; ++ ++ const __m128i m4 = _mm_set1_epi8(0xf); ++}; ++ ++struct DequantizerQ3K final : public BaseDequantizer { ++ DequantizerQ3K(const void * vx, size_t bx) : BaseDequantizer(vx, bx) {} ++ ++ template ++ inline void new_block(int i, const Q8& q8, __m256 * accm, __m256i * scales) { ++ d = GGML_FP16_TO_FP32(x[i].d); ++ hbits.load(x[i].hmask); ++ process_mins_and_scales_16(sc3.make_scales((const uint16_t *)x[i].scales), q8, i, -4.f*d, accm, scales); ++ } ++ inline void prepare(int i, int j) { ++ bits.prepare(x[i].qs, j); ++ hbits.apply(bits, j == 0); ++ } ++ ++ Q2Bits bits; ++ HighBit3 hbits; ++ ScaleQ3 sc3; ++ ++ const __m128i m32 = _mm_set1_epi8(-32); ++}; ++ ++struct DequantizerQ4K final : public BaseDequantizer { ++ DequantizerQ4K(const void * vx, size_t bx) : BaseDequantizer(vx, bx) {} ++ template ++ inline __m256i new_block(int i, const Q8& q8, __m256 * accd) { ++ d = GGML_FP16_TO_FP32(x[i].data.data.d); ++ return s8k.process_mins_and_scales(x[i].scales, -GGML_FP16_TO_FP32(x[i].data.data.dmin), i, q8, accd); ++ } ++ inline void prepare(int i, int j) { ++ bits.prepare(x[i].qs, j); ++ } ++ ++ Q4Bits bits; ++ Scales8K s8k; ++}; ++ ++struct DequantizerQ5K final : public BaseDequantizer { ++ DequantizerQ5K(const void * vx, size_t bx) : BaseDequantizer(vx, bx) {} ++ template ++ inline __m256i new_block(int i, const Q8& q8, __m256 * accd) { ++ d = GGML_FP16_TO_FP32(x[i].data.data.d); ++ hbits.load(x[i].qh); ++ return s8k.process_mins_and_scales(x[i].scales, -GGML_FP16_TO_FP32(x[i].data.data.dmin), i, q8, accd); ++ } ++ inline void prepare(int i, int j) { ++ bits.prepare(x[i].qs, j); ++ hbits.apply(bits, j == 0); ++ } ++ ++ Q4Bits bits; ++ HighBit5 hbits; ++ Scales8K s8k; ++}; ++ ++struct DequantizerQ6K final : public BaseDequantizer { ++ DequantizerQ6K(const void * vx, size_t bx) : BaseDequantizer(vx, bx) {} ++ template ++ inline void new_block(int i, const Q8& q8, __m256 * accm, __m256i * scales) { ++ d = GGML_FP16_TO_FP32(x[i].d); ++ process_mins_and_scales_16(_mm_loadu_si128((const __m128i *)x[i].scales), q8, i, -32.f*d, accm, scales); ++ } ++ inline void prepare(int i, int j) { ++ bits.prepare64(x[i].ql, j); ++ auto hbits = _mm256_loadu_si256((const __m256i *)x[i].qh + j); ++ bits.values[0] = _mm256_or_si256(bits.values[0], _mm256_and_si256(_mm256_slli_epi16(hbits, 4), mh)); ++ bits.values[1] = _mm256_or_si256(bits.values[1], _mm256_and_si256(_mm256_slli_epi16(hbits, 2), mh)); ++ bits.values[2] = _mm256_or_si256(bits.values[2], _mm256_and_si256(hbits, mh)); ++ bits.values[3] = _mm256_or_si256(bits.values[3], _mm256_and_si256(_mm256_srli_epi16(hbits, 2), mh)); ++ } ++ ++ Q4Bits bits; ++ const __m256i mh = _mm256_set1_epi8(0x30); ++}; ++ ++struct DequantizerIQ4XS final : public BaseDequantizer { ++ DequantizerIQ4XS(const void * vx, size_t bx) : BaseDequantizer(vx, bx), ++ values(MM256_SET1_M128I(_mm_loadu_si128((const __m128i *)kvalues_iq4nl))) {} ++ template ++ inline __m256i new_block(int i, [[maybe_unused]] const Q8& q8, [[maybe_unused]] __m256 * accd) { ++ d = GGML_FP16_TO_FP32(x[i].d); ++ auto scales128 = siq4.make_scales(*(const uint32_t *)x[i].scales_l, x[i].scales_h); ++ return MM256_SET1_M128I(scales128); ++ } ++ inline void prepare(int i, int j) { ++ bits.prepare16(x[i].qs, j); ++ bits.values[0] = _mm256_shuffle_epi8(values, bits.values[0]); ++ bits.values[1] = _mm256_shuffle_epi8(values, bits.values[1]); ++ bits.values[2] = _mm256_shuffle_epi8(values, bits.values[2]); ++ bits.values[3] = _mm256_shuffle_epi8(values, bits.values[3]); ++ } ++ ++ Q4Bits bits; ++ ScaleIQ4XS siq4; ++ const __m256i values; ++}; ++ ++template ++static void mul_mat_qX_K_q8_K_T(int n, const void * vx, size_t bx, const DataInfo& info, int nrc_x) { ++ assert(n % QK_K == 0); ++ const int nb = n / QK_K; ++ ++ Q8 q8(info); ++ ++ Dequantizer deq(vx, bx); ++ ++ __m256 accd[nrc_y]; ++ __m256i scales[4]; ++ ++ for (int ix = 0; ix < nrc_x; ++ix) { ++ ++ for (int iy = 0; iy < nrc_y; ++iy) accd[iy] = _mm256_setzero_ps(); ++ ++ deq.new_row(ix); ++ ++ for (int i = 0; i < nb; ++i) { ++ ++ auto all_scales = deq.new_block(i, q8, accd); ++ ++ __m256i sumi[nrc_y]; ++ ++ for (int j = 0; j < QK_K/128; ++j) { ++ ++ deq.prepare(i, j); ++ ++ set_scales_8(all_scales, j, scales); ++ ++ if constexpr (std::is_same_v) { ++ multiply_add_avx2(deq.bits, scales, j, i, q8, sumi); ++ } else { ++ multiply_add(deq.bits, scales, j, i, q8, sumi); ++ } ++ ++ } ++ ++ for (int iy = 0; iy < nrc_y; ++iy) { ++ const __m256 vd = _mm256_set1_ps(deq.d*q8.scale(iy, i)); ++ accd[iy] = _mm256_fmadd_ps(vd, _mm256_cvtepi32_ps(sumi[iy]), accd[iy]); ++ } ++ ++ } ++ ++ for (int iy = 0; iy < nrc_y; ++iy) { ++ info.store(ix, iy, hsum_float_8(accd[iy])); ++ } ++ ++ } ++} ++ ++template ++static void mul_mat_qY_K_q8_K_T(int n, const void * vx, size_t bx, const DataInfo& info, int nrc_x) { ++ assert(n%QK_K == 0); ++ const int nb = n/QK_K; ++ ++ Q8 q8(info); ++ ++ __m256i all_scales[2]; ++ __m256i scales[4]; ++ __m256 accd[nrc_y]; ++ ++ Dequantizer deq(vx, bx); ++ ++ for (int ix = 0; ix < nrc_x; ++ix) { ++ ++ deq.new_row(ix); ++ ++ for (int iy = 0; iy < nrc_y; ++iy) accd[iy] = _mm256_setzero_ps(); ++ ++ for (int i = 0; i < nb; ++i) { ++ ++ deq.new_block(i, q8, accd, all_scales); ++ ++ __m256i sumi[nrc_y]; ++ ++ for (int j = 0; j < QK_K/128; ++j) { ++ deq.prepare(i, j); ++ set_scales_16(all_scales[j], scales); ++ multiply_add(deq.bits, scales, j, i, q8, sumi); ++ } ++ ++ for (int iy = 0; iy < nrc_y; ++iy) { ++ accd[iy] = _mm256_fmadd_ps(_mm256_set1_ps(deq.d*q8.scale(iy, i)), _mm256_cvtepi32_ps(sumi[iy]), accd[iy]); ++ } ++ ++ } ++ ++ for (int iy = 0; iy < nrc_y; ++iy) { ++ info.store(ix, iy, hsum_float_8(accd[iy])); ++ } ++ ++ } ++ ++} ++ ++#endif ++ ++template void set_functions(mul_mat_t * funcs) { ++#ifdef HAVE_FANCY_SIMD ++ if constexpr (std::is_same_v) { ++ IQK_SET_MUL_MAT_FUNCTIONS_T(mul_mat_iqX_k_q8_K_AVX512_dpbusd, Dequantizer, funcs) ++ } else { ++ IQK_SET_MUL_MAT_FUNCTIONS_T(mul_mat_qX_K_q8_K_AVX512, Dequantizer, funcs) ++ funcs[0] = mul_mat_qX_K_q8_K_AVX512_1; ++ } ++#else ++ if constexpr (std::is_same_v || ++ std::is_same_v || ++ std::is_same_v) { ++ IQK_SET_MUL_MAT_FUNCTIONS_T(mul_mat_qY_K_q8_K_T, Dequantizer, funcs) ++ } else { ++ IQK_SET_MUL_MAT_FUNCTIONS_T(mul_mat_qX_K_q8_K_T, Dequantizer, funcs) ++ } ++#endif ++} ++ ++} // namespace ++ ++bool iqk_set_kernels_kquants(int typeA, mul_mat_t * kernels) { ++ switch (typeA) { ++ case GGML_TYPE_Q2_K: set_functions(kernels); break; ++ case GGML_TYPE_Q3_K: set_functions(kernels); break; ++ case GGML_TYPE_Q4_K: set_functions(kernels); break; ++ case GGML_TYPE_Q5_K: set_functions(kernels); break; ++ case GGML_TYPE_Q6_K: set_functions(kernels); break; ++ case GGML_TYPE_IQ4_XS: set_functions(kernels); break; ++ default: return false; ++ } ++ return true; ++} ++ ++#else ++// --------------------------------- __aarch64__ -------------------------------------- ++ ++namespace { ++ ++template ++inline void accum_mins_8(const int16x8_t& mins, const Q8& q8, float32x4_t * acc, int i, float c) { ++ for (int iy = 0; iy < Q8::nrc_y; ++iy) { ++ auto q8s = q8.load_bsums8(iy, i); ++ int32x4_t b1 = vmull_s16(vget_low_s16(mins), vget_low_s16(q8s)); ++ int32x4_t b2 = vmull_s16(vget_high_s16(mins), vget_high_s16(q8s)); ++ float32x4_t prod = vcvtq_f32_s32(vaddq_s32(b1, b2)); ++ acc[iy] = vmlaq_f32(acc[iy], prod, vdupq_n_f32(c*q8.scale(iy, i))); ++ } ++} ++template ++inline void accum_mins_16(const int16x8x2_t& mins, const Q8& q8, float32x4_t * acc, int i, float c) { ++ for (int iy = 0; iy < Q8::nrc_y; ++iy) { ++ auto q8s = q8.load_bsums(iy, i); ++ int32x4_t b1 = vmull_s16(vget_low_s16 (mins.val[0]), vget_low_s16 (q8s.val[0])); ++ int32x4_t b2 = vmull_s16(vget_high_s16(mins.val[0]), vget_high_s16(q8s.val[0])); ++ int32x4_t b3 = vmull_s16(vget_low_s16 (mins.val[1]), vget_low_s16 (q8s.val[1])); ++ int32x4_t b4 = vmull_s16(vget_high_s16(mins.val[1]), vget_high_s16(q8s.val[1])); ++ float32x4_t prod = vcvtq_f32_s32(vaddq_s32(vaddq_s32(b1, b2), vaddq_s32(b3, b4))); ++ acc[iy] = vmlaq_f32(acc[iy], prod, vdupq_n_f32(c*q8.scale(iy, i))); ++ } ++} ++ ++struct Scales8 { ++ uint32_t utmp[4]; ++ const uint8_t * sc8 = (const uint8_t *)utmp; ++ template ++ inline int32x4x2_t process_scales_mins(const Qx& x, const Q8& q8, int i, float32x4_t * acc) { ++ make_q4_scales(x.scales, utmp); ++ int16x8_t mins = vmovl_s8(vld1_s8((const int8_t *)sc8 + 8)); ++ accum_mins_8(mins, q8, acc, i, -GGML_FP16_TO_FP32(x.data.data.dmin)); ++ ++ uint8x8_t scales8 = vld1_u8(sc8); ++ uint16x8_t scales16 = vmovl_u8(scales8); ++ int32x4x2_t scales = {vreinterpretq_s32_u32(vmovl_u16(vget_low_u16(scales16))), ++ vreinterpretq_s32_u32(vmovl_u16(vget_high_u16(scales16)))}; ++ return scales; ++ } ++ inline float32x4x4_t make_scales(float d, float m, const uint8_t * scales) { ++ make_q4_scales(scales, utmp); ++ auto d16 = vmovl_u8(vld1_u8(sc8+0)); ++ auto m16 = vmovl_u8(vld1_u8(sc8+8)); ++ auto vd = vdupq_n_f32(d); ++ auto vm = vdupq_n_f32(m); ++ return { vmulq_f32(vd, vcvtq_f32_u32(vmovl_u16(vget_low_u16 (d16)))), ++ vmulq_f32(vd, vcvtq_f32_u32(vmovl_u16(vget_high_u16(d16)))), ++ vmulq_f32(vm, vcvtq_f32_u32(vmovl_u16(vget_low_u16 (m16)))), ++ vmulq_f32(vm, vcvtq_f32_u32(vmovl_u16(vget_high_u16(m16)))) }; ++ } ++}; ++ ++struct DequantizerQ4K final : public BaseDequantizer { ++ DequantizerQ4K(const void * vx, size_t bx, int nrc) : BaseDequantizer(vx, bx, nrc) {} ++ ++ constexpr static int num_blocks() { return 8; } ++ constexpr static bool should_scale_quants() { return false; } ++ ++ template ++ inline int32x4x2_t new_block(int i, const Q8& q8, float32x4_t * acc) { ++ d = GGML_FP16_TO_FP32(x[i].data.data.d); ++ return s8.process_scales_mins(x[i], q8, i, acc); ++ } ++ inline float32x4x4_t new_block(int i) { ++ d = GGML_FP16_TO_FP32(x[i].data.data.d); ++ float m = -GGML_FP16_TO_FP32(x[i].data.data.dmin); ++ return s8.make_scales(d, m, x[i].scales); ++ } ++ inline void prepare(int i, int j) { ++ if (nrc == 1) bits.prepare_v2(x[i].qs+64*j); ++ else bits.prepare(x[i].qs+64*j); ++ } ++ ++ Q4bits bits; ++ Scales8 s8; ++ ++}; ++ ++struct HighBit5 { ++ const uint8x16_t mhb = vdupq_n_u8(0x10); ++ uint8x16x2_t bits; ++ inline void apply(uint8x16x4_t& b1, uint8x16x4_t& b2, bool do_shift) { ++ b1.val[0] = vorrq_u8(b1.val[0], vandq_u8(vshlq_n_u8(bits.val[0], 4), mhb)); ++ b1.val[1] = vorrq_u8(b1.val[1], vandq_u8(vshlq_n_u8(bits.val[1], 4), mhb)); ++ b1.val[2] = vorrq_u8(b1.val[2], vandq_u8(vshlq_n_u8(bits.val[0], 3), mhb)); ++ b1.val[3] = vorrq_u8(b1.val[3], vandq_u8(vshlq_n_u8(bits.val[1], 3), mhb)); ++ ++ b2.val[0] = vorrq_u8(b2.val[0], vandq_u8(vshlq_n_u8(bits.val[0], 2), mhb)); ++ b2.val[1] = vorrq_u8(b2.val[1], vandq_u8(vshlq_n_u8(bits.val[1], 2), mhb)); ++ b2.val[2] = vorrq_u8(b2.val[2], vandq_u8(vshlq_n_u8(bits.val[0], 1), mhb)); ++ b2.val[3] = vorrq_u8(b2.val[3], vandq_u8(vshlq_n_u8(bits.val[1], 1), mhb)); ++ ++ if (do_shift) { ++ bits.val[0] = vshrq_n_u8(bits.val[0], 4); ++ bits.val[1] = vshrq_n_u8(bits.val[1], 4); ++ } ++ } ++}; ++ ++struct HighBit3 { ++ const uint8x16_t mhb = vdupq_n_u8(0x04); ++ uint8x16x2_t bits; ++ inline void apply(uint8x16x4_t& b1, uint8x16x4_t& b2, bool do_shift) { ++ b1.val[0] = vorrq_u8(b1.val[0], vandq_u8(vshlq_n_u8(bits.val[0], 2), mhb)); ++ b1.val[1] = vorrq_u8(b1.val[1], vandq_u8(vshlq_n_u8(bits.val[1], 2), mhb)); ++ b1.val[2] = vorrq_u8(b1.val[2], vandq_u8(vshlq_n_u8(bits.val[0], 1), mhb)); ++ b1.val[3] = vorrq_u8(b1.val[3], vandq_u8(vshlq_n_u8(bits.val[1], 1), mhb)); ++ ++ b2.val[0] = vorrq_u8(b2.val[0], vandq_u8(bits.val[0], mhb)); ++ b2.val[1] = vorrq_u8(b2.val[1], vandq_u8(bits.val[1], mhb)); ++ b2.val[2] = vorrq_u8(b2.val[2], vandq_u8(vshrq_n_u8(bits.val[0], 1), mhb)); ++ b2.val[3] = vorrq_u8(b2.val[3], vandq_u8(vshrq_n_u8(bits.val[1], 1), mhb)); ++ ++ if (do_shift) { ++ bits.val[0] = vshrq_n_u8(bits.val[0], 4); ++ bits.val[1] = vshrq_n_u8(bits.val[1], 4); ++ } ++ } ++}; ++ ++struct DequantizerQ5K final : public BaseDequantizer { ++ DequantizerQ5K(const void * vx, size_t bx, int nrc) : BaseDequantizer(vx, bx, nrc) {} ++ ++ constexpr static int num_blocks() { return 8; } ++ constexpr static bool should_scale_quants() { return false; } ++ ++ template ++ inline int32x4x2_t new_block(int i, const Q8& q8, float32x4_t * acc) { ++ d = GGML_FP16_TO_FP32(x[i].data.data.d); ++ h.bits = vld1q_u8_x2(x[i].qh); ++ return s8.process_scales_mins(x[i], q8, i, acc); ++ } ++ inline float32x4x4_t new_block(int i) { ++ d = GGML_FP16_TO_FP32(x[i].data.data.d); ++ float m = -GGML_FP16_TO_FP32(x[i].data.data.dmin); ++ h.bits = vld1q_u8_x2(x[i].qh); ++ return s8.make_scales(d, m, x[i].scales); ++ } ++ inline void prepare(int i, int j) { ++ if (nrc == 1) bits.prepare_v2(x[i].qs+64*j); ++ else bits.prepare(x[i].qs+64*j); ++ h.apply(bits.b1, bits.b2, j == 0); ++ } ++ ++ Q4bits bits; ++ HighBit5 h; ++ Scales8 s8; ++ ++ uint8x16x2_t hbits; ++ ++}; ++ ++inline int32x4x4_t make_wider(const int16x8x2_t& scales16) { ++ int32x4x4_t scales = { ++ vmovl_s16(vget_low_s16 (scales16.val[0])), ++ vmovl_s16(vget_high_s16(scales16.val[0])), ++ vmovl_s16(vget_low_s16 (scales16.val[1])), ++ vmovl_s16(vget_high_s16(scales16.val[1])), ++ }; ++ return scales; ++} ++ ++template ++inline int32x4x4_t process_scales_mins_16(const int8x16_t& scales8, const Q8& q8, float32x4_t * acc, int i, float c) { ++ int16x8x2_t scales16; ++ scales16.val[0] = vmovl_s8(vget_low_s8(scales8)); ++ scales16.val[1] = vmovl_s8(vget_high_s8(scales8)); ++ accum_mins_16(scales16, q8, acc, i, c); ++ return make_wider(scales16); ++} ++ ++struct DequantizerQ6K final : public BaseDequantizer { ++ DequantizerQ6K(const void * vx, size_t bx, int nrc) : BaseDequantizer(vx, bx, nrc) {} ++ ++ constexpr static int num_blocks() { return 16; } ++ constexpr static bool should_scale_quants() { return false; } ++ ++ template ++ inline int32x4x4_t new_block(int i, const Q8& q8, float32x4_t * acc) { ++ d = GGML_FP16_TO_FP32(x[i].d); ++ return process_scales_mins_16(vld1q_s8(x[i].scales), q8, acc, i, -32.f*d); ++ } ++ inline float32x4x4_t new_block(int i) { ++ d = GGML_FP16_TO_FP32(x[i].d); ++ auto vd = vdupq_n_f32(d); ++ auto scales8 = vld1q_s8(x[i].scales); ++ auto scales16_1 = vmovl_s8(vget_low_s8 (scales8)); ++ auto scales16_2 = vmovl_s8(vget_high_s8(scales8)); ++ return { vmulq_f32(vd, vcvtq_f32_s32(vmovl_s16(vget_low_s16 (scales16_1)))), ++ vmulq_f32(vd, vcvtq_f32_s32(vmovl_s16(vget_high_s16(scales16_1)))), ++ vmulq_f32(vd, vcvtq_f32_s32(vmovl_s16(vget_low_s16 (scales16_2)))), ++ vmulq_f32(vd, vcvtq_f32_s32(vmovl_s16(vget_high_s16(scales16_2)))) }; ++ } ++ ++ inline void prepare(int i, int j) { ++ ++ auto hbits = vld1q_u8_x2(x[i].qh + 32*j); ++ ++ bits.prepare64(x[i].ql+64*j); ++ bits.b1.val[0] = vorrq_u8(bits.b1.val[0], vandq_u8(vshlq_n_u8(hbits.val[0], 4), mhb)); ++ bits.b1.val[1] = vorrq_u8(bits.b1.val[1], vandq_u8(vshlq_n_u8(hbits.val[1], 4), mhb)); ++ bits.b1.val[2] = vorrq_u8(bits.b1.val[2], vandq_u8(vshlq_n_u8(hbits.val[0], 2), mhb)); ++ bits.b1.val[3] = vorrq_u8(bits.b1.val[3], vandq_u8(vshlq_n_u8(hbits.val[1], 2), mhb)); ++ ++ bits.b2.val[0] = vorrq_u8(bits.b2.val[0], vandq_u8(hbits.val[0], mhb)); ++ bits.b2.val[1] = vorrq_u8(bits.b2.val[1], vandq_u8(hbits.val[1], mhb)); ++ bits.b2.val[2] = vorrq_u8(bits.b2.val[2], vandq_u8(vshrq_n_u8(hbits.val[0], 2), mhb)); ++ bits.b2.val[3] = vorrq_u8(bits.b2.val[3], vandq_u8(vshrq_n_u8(hbits.val[1], 2), mhb)); ++ ++ } ++ ++ inline void prepare_signed(int i, int j) { ++ prepare(i, j); ++ auto m32 = vdupq_n_s8(-32); ++ for (int k = 0; k < 4; ++k) { ++ bits.b1.val[k] = vreinterpretq_u8_s8(vaddq_s8(vreinterpretq_s8_u8(bits.b1.val[k]), m32)); ++ bits.b2.val[k] = vreinterpretq_u8_s8(vaddq_s8(vreinterpretq_s8_u8(bits.b2.val[k]), m32)); ++ } ++ } ++ ++ Q4bits bits; ++ ++ const uint8x16_t mhb = vdupq_n_u8(0x30); ++ ++}; ++ ++struct DequantizerQ3K final : public BaseDequantizer { ++ DequantizerQ3K(const void * vx, size_t bx, int nrc) : BaseDequantizer(vx, bx, nrc) {} ++ ++ constexpr static int num_blocks() { return 16; } ++ constexpr static bool should_scale_quants() { return false; } ++ ++ template ++ inline int32x4x4_t new_block(int i, const Q8& q8, float32x4_t * acc) { ++ d = GGML_FP16_TO_FP32(x[i].d); ++ h.bits = vld1q_u8_x2(x[i].hmask); ++ mask = vdupq_n_u8(0x01); ++ const uint16_t * sc16 = (const uint16_t *)x[i].scales; ++ uint32_t aux0 = sc16[0] | (sc16[1] << 16); ++ uint32_t aux1 = sc16[2] | (sc16[3] << 16); ++ uint32_t aux2 = sc16[4] | (sc16[5] << 16); ++ aux32[0] = (aux0 & 0x0f0f0f0f) | ((aux2 << 4) & 0x30303030); ++ aux32[1] = (aux1 & 0x0f0f0f0f) | ((aux2 << 2) & 0x30303030); ++ aux32[2] = ((aux0 >> 4) & 0x0f0f0f0f) | ((aux2 >> 0) & 0x30303030); ++ aux32[3] = ((aux1 >> 4) & 0x0f0f0f0f) | ((aux2 >> 2) & 0x30303030); ++ auto scales8 = vaddq_s8(vld1q_s8((const int8_t *)aux32), vdupq_n_s8(-32)); ++ if (nrc > 1) { ++ return process_scales_mins_16(scales8, q8, acc, i, -4.f*d); ++ } ++ int16x8x2_t scales16; ++ scales16.val[0] = vmovl_s8(vget_low_s8(scales8)); ++ scales16.val[1] = vmovl_s8(vget_high_s8(scales8)); ++ return make_wider(scales16); ++ } ++ ++ inline void prepare(int i, int j) { ++ bits.prepare(x[i].qs+32*j); ++ if (nrc > 1) { ++ h.apply(bits.b1, bits.b2, j == 0); ++ } else { ++ auto minus4 = vdupq_n_u8(0xfc); ++ auto zero = vdupq_n_u8(0); ++ bits.b1.val[0] = vorrq_u8(bits.b1.val[0], vandq_u8(minus4, vceqq_u8(vandq_u8(h.bits.val[0], mask), zero))); ++ bits.b1.val[1] = vorrq_u8(bits.b1.val[1], vandq_u8(minus4, vceqq_u8(vandq_u8(h.bits.val[1], mask), zero))); ++ mask = vshlq_n_u8(mask, 1); ++ bits.b1.val[2] = vorrq_u8(bits.b1.val[2], vandq_u8(minus4, vceqq_u8(vandq_u8(h.bits.val[0], mask), zero))); ++ bits.b1.val[3] = vorrq_u8(bits.b1.val[3], vandq_u8(minus4, vceqq_u8(vandq_u8(h.bits.val[1], mask), zero))); ++ mask = vshlq_n_u8(mask, 1); ++ bits.b2.val[0] = vorrq_u8(bits.b2.val[0], vandq_u8(minus4, vceqq_u8(vandq_u8(h.bits.val[0], mask), zero))); ++ bits.b2.val[1] = vorrq_u8(bits.b2.val[1], vandq_u8(minus4, vceqq_u8(vandq_u8(h.bits.val[1], mask), zero))); ++ mask = vshlq_n_u8(mask, 1); ++ bits.b2.val[2] = vorrq_u8(bits.b2.val[2], vandq_u8(minus4, vceqq_u8(vandq_u8(h.bits.val[0], mask), zero))); ++ bits.b2.val[3] = vorrq_u8(bits.b2.val[3], vandq_u8(minus4, vceqq_u8(vandq_u8(h.bits.val[1], mask), zero))); ++ mask = vshlq_n_u8(mask, 1); ++ } ++ } ++ ++ uint32_t aux32[4]; ++ ++ Q2bits bits; ++ ++ uint8x16_t mask; ++ HighBit3 h; ++ ++}; ++ ++struct DequantizerQ2K final : public BaseDequantizer { ++ DequantizerQ2K(const void * vx, size_t bx, int nrc) : BaseDequantizer(vx, bx, nrc) {} ++ ++ constexpr static int num_blocks() { return 16; } ++ constexpr static bool should_scale_quants() { return true; } ++ ++ template ++ inline void process_scales(int i, const Q8& q8, float32x4_t * acc) { ++ d = GGML_FP16_TO_FP32(x[i].data.data.d); ++ auto scales_and_mins = vld1q_u8(x[i].scales); ++ auto mins8 = vreinterpretq_s8_u8(vshrq_n_u8(scales_and_mins, 4)); ++ int16x8x2_t scales16; ++ scales16.val[0] = vmovl_s8(vget_low_s8(mins8)); ++ scales16.val[1] = vmovl_s8(vget_high_s8(mins8)); ++ accum_mins_16(scales16, q8, acc, i, -GGML_FP16_TO_FP32(x[i].data.data.dmin)); ++ ++ scales8 = vandq_u8(scales_and_mins, vdupq_n_u8(0xf)); ++ } ++ ++ template ++ inline int32x4x4_t new_block(int i, const Q8& q8, float32x4_t * acc) { ++ process_scales(i, q8, acc); ++ int16x8x2_t scales16; ++ scales16.val[0] = vmovl_s8(vget_low_s8(vreinterpretq_s8_u8(scales8))); ++ scales16.val[1] = vmovl_s8(vget_high_s8(vreinterpretq_s8_u8(scales8))); ++ return make_wider(scales16); ++ } ++ ++ template ++ inline void compute(const Q8& q8, int i, int j, int32x4_t * sumi) { ++ auto m1 = vdupq_n_u8(1); ++ auto shuffle = vdupq_n_u8(8*j); ++ bits.b1.val[0] = vmulq_u8(bits.b1.val[0], vqtbl1q_u8(scales8, shuffle)); shuffle = vaddq_u8(shuffle, m1); ++ bits.b1.val[1] = vmulq_u8(bits.b1.val[1], vqtbl1q_u8(scales8, shuffle)); shuffle = vaddq_u8(shuffle, m1); ++ bits.b1.val[2] = vmulq_u8(bits.b1.val[2], vqtbl1q_u8(scales8, shuffle)); shuffle = vaddq_u8(shuffle, m1); ++ bits.b1.val[3] = vmulq_u8(bits.b1.val[3], vqtbl1q_u8(scales8, shuffle)); shuffle = vaddq_u8(shuffle, m1); ++ bits.b2.val[0] = vmulq_u8(bits.b2.val[0], vqtbl1q_u8(scales8, shuffle)); shuffle = vaddq_u8(shuffle, m1); ++ bits.b2.val[1] = vmulq_u8(bits.b2.val[1], vqtbl1q_u8(scales8, shuffle)); shuffle = vaddq_u8(shuffle, m1); ++ bits.b2.val[2] = vmulq_u8(bits.b2.val[2], vqtbl1q_u8(scales8, shuffle)); shuffle = vaddq_u8(shuffle, m1); ++ bits.b2.val[3] = vmulq_u8(bits.b2.val[3], vqtbl1q_u8(scales8, shuffle)); shuffle = vaddq_u8(shuffle, m1); ++ for (int iy = 0; iy < Q8::nrc_y; ++iy) { ++ auto q8b_1 = q8.load_quants(iy, i, 4*j+0); ++ sumi[iy] = ggml_vdotq_s32(ggml_vdotq_s32(sumi[iy], vreinterpretq_s8_u8(bits.b1.val[0]), q8b_1.val[0]), ++ vreinterpretq_s8_u8(bits.b1.val[1]), q8b_1.val[1]); ++ ++ auto q8b_2 = q8.load_quants(iy, i, 4*j+1); ++ sumi[iy] = ggml_vdotq_s32(ggml_vdotq_s32(sumi[iy], vreinterpretq_s8_u8(bits.b1.val[2]), q8b_2.val[0]), ++ vreinterpretq_s8_u8(bits.b1.val[3]), q8b_2.val[1]); ++ ++ auto q8b_3 = q8.load_quants(iy, i, 4*j+2); ++ sumi[iy] = ggml_vdotq_s32(ggml_vdotq_s32(sumi[iy], vreinterpretq_s8_u8(bits.b2.val[0]), q8b_3.val[0]), ++ vreinterpretq_s8_u8(bits.b2.val[1]), q8b_3.val[1]); ++ ++ auto q8b_4 = q8.load_quants(iy, i, 4*j+3); ++ sumi[iy] = ggml_vdotq_s32(ggml_vdotq_s32(sumi[iy], vreinterpretq_s8_u8(bits.b2.val[2]), q8b_4.val[0]), ++ vreinterpretq_s8_u8(bits.b2.val[3]), q8b_4.val[1]); ++ } ++ } ++ ++ inline void prepare(int i, int j) { ++ bits.prepare(x[i].qs+32*j); ++ } ++ ++ uint32_t aux32[4]; ++ ++ uint8x16_t scales8; ++ ++ Q2bits bits; ++ ++}; ++ ++struct DequantizerIQ4XS final : public BaseDequantizer { ++ ++ static int8x16_t load_values() { ++ static const int8_t iq4nl_values[16] = {-127, -104, -83, -65, -49, -35, -22, -10, 1, 13, 25, 38, 53, 69, 89, 113}; ++ return vld1q_s8(iq4nl_values); ++ } ++ ++ DequantizerIQ4XS(const void * vx, size_t bx, int nrc) : BaseDequantizer(vx, bx, nrc), values(load_values()) {} ++ ++ constexpr static int num_blocks() { return 8; } ++ constexpr static bool should_scale_quants() { return false; } ++ ++ inline void new_row(int ix) { x = (const block_iq4_xs *)((const char *)vx + bx*ix); } ++ ++ template ++ inline int32x4x2_t new_block(int i, const Q8& q8, float32x4_t * acc) { ++ (void)q8; ++ (void)acc; ++ d = GGML_FP16_TO_FP32(x[i].d); ++ const uint16_t scales_h = x[i].scales_h; ++ const uint16_t * scales_l = (const uint16_t *)x[i].scales_l; ++ aux32[0] = scales_l[0] | (scales_l[1] << 16); ++ aux32[1] = aux32[0] >> 4; ++ // scl is ordered as 0, 2, 4, 6, 1, 3, 5, 7 ++ uint8x8_t scl8 = vand_u8(vld1_u8((const uint8_t *)aux32), vdup_n_u8(0xf)); ++ uint16_t * aux16 = (uint16_t *)aux32; ++ aux16[0] = scales_h << 4; aux16[1] = scales_h << 2; aux16[2] = scales_h; aux16[3] = scales_h >> 2; ++ // sch is ordered as 0, 4, 1, 5, 2, 6, 3, 7 ++ uint8x8_t sch8 = vand_u8(vld1_u8((const uint8_t *)aux16), vdup_n_u8(0x30)); ++ int8x8_t scales8 = vadd_s8(vreinterpret_s8_u8(vorr_u8(scl8, vtbl1_u8(sch8, vreinterpret_u8_u32(hshuff)))), vdup_n_s8(-32)); ++ // shuffle 0, 2, 4, 6, 1, 3, 5, 7 -> 0, 1, 2, 3, 4, 5, 6, 7 ++ scales8 = vtbl1_s8(scales8, vreinterpret_s8_u32(hshuff)); ++ int16x8_t scales16 = vmovl_s8(scales8); ++ int32x4x2_t scales = {vmovl_s16(vget_low_s16(scales16)), vmovl_s16(vget_high_s16(scales16))}; ++ return scales; ++ } ++ inline void prepare(int i, int j) { ++ bits.prepare16(x[i].qs+64*j); ++ //if (nrc == 1) { ++ // bits.prepare16_v2(x[i].qs+64*j); ++ //} else { ++ // bits.prepare16(x[i].qs+64*j); ++ //} ++ for (int k = 0; k < 4; ++k) { ++ bits.b1.val[k] = vreinterpretq_u8_s8(vqtbl1q_s8(values, bits.b1.val[k])); ++ bits.b2.val[k] = vreinterpretq_u8_s8(vqtbl1q_s8(values, bits.b2.val[k])); ++ } ++ } ++ ++ Q4bits bits; ++ const int8x16_t values; ++ uint32_t aux32[2]; ++ ++ constexpr static uint32x2_t hshuff = {0x05010400, 0x07030602}; ++ ++}; ++ ++} // namespace ++ ++bool iqk_set_kernels_kquants(int typeA, mul_mat_t * kernels) { ++ switch (typeA) { ++ case GGML_TYPE_Q2_K: IQK_SET_MUL_MAT_FUNCTIONS_T(mul_mat_qX_K_q8_K_T, DequantizerQ2K, kernels); break; ++ case GGML_TYPE_Q3_K: IQK_SET_MUL_MAT_FUNCTIONS_T(mul_mat_qX_K_q8_K_T, DequantizerQ3K, kernels); break; ++ case GGML_TYPE_Q4_K: IQK_SET_MUL_MAT_FUNCTIONS_T(mul_mat_qX_K_q8_K_T, DequantizerQ4K, kernels); break; ++ case GGML_TYPE_Q5_K: IQK_SET_MUL_MAT_FUNCTIONS_T(mul_mat_qX_K_q8_K_T, DequantizerQ5K, kernels); break; ++ case GGML_TYPE_Q6_K: IQK_SET_MUL_MAT_FUNCTIONS_T(mul_mat_qX_K_q8_K_T, DequantizerQ6K, kernels); break; ++ case GGML_TYPE_IQ4_XS: IQK_SET_MUL_MAT_FUNCTIONS_T(mul_mat_qX_K_q8_K_T, DequantizerIQ4XS, kernels); break; ++ default: return false; ++ } ++ return true; ++} ++ ++#endif ++ ++#endif // IQK_IMPLEMENT diff --git a/scripts/unsloth/ggml-patches/0020-ggml-cpu-arm-sve128-prefer-neon.patch b/scripts/unsloth/ggml-patches/0020-ggml-cpu-arm-sve128-prefer-neon.patch new file mode 100644 index 0000000000..0fde4d9738 --- /dev/null +++ b/scripts/unsloth/ggml-patches/0020-ggml-cpu-arm-sve128-prefer-neon.patch @@ -0,0 +1,34 @@ +diff --git a/src/ggml-cpu/arch/arm/cpu-feats.cpp b/src/ggml-cpu/arch/arm/cpu-feats.cpp +index 84a11ea..79a3b2c 100644 +--- a/src/ggml-cpu/arch/arm/cpu-feats.cpp ++++ b/src/ggml-cpu/arch/arm/cpu-feats.cpp +@@ -1,6 +1,11 @@ + #include "ggml-backend-impl.h" + #include "ggml-feats.h" + ++#if defined(__linux__) && defined(GGML_USE_SVE) ++#include ++#include ++#endif ++ + #if defined(__aarch64__) || defined(_M_ARM64) + + static int ggml_backend_cpu_aarch64_score() { +@@ -18,6 +23,17 @@ static int ggml_backend_cpu_aarch64_score() { + #endif + #ifdef GGML_USE_SVE + if (!af.has_sve) { return 0; } ++#if defined(__linux__) && defined(PR_SVE_GET_VL) ++ // Unsloth: with 128-bit SVE (Neoverse N2 / V2, Cortex-A710 and later) the SVE variants ran stable-diffusion.cpp ++ // slower than the NEON dotprod + FP16 one (armv8.2_2) on a Neoverse N2: 1.29-1.38x slower for armv8.6_2 and about ++ // 5x slower for armv8.2_3 / armv8.6_1, F16, Q8_0, Q3_K and Q4_K alike. So they are not offered there and the best ++ // NEON variant is picked. Wider SVE (Graviton 3, A64FX) keeps them. GGML_CPU_ARM_SVE128=1 offers them again. ++ { ++ const int vl = prctl(PR_SVE_GET_VL); ++ const char * keep = getenv("GGML_CPU_ARM_SVE128"); ++ if (vl > 0 && (vl & PR_SVE_VL_LEN_MASK) <= 16 && !(keep != NULL && keep[0] == '1')) { return 0; } ++ } ++#endif + score += 1<<3; + #endif + #ifdef GGML_USE_MATMUL_INT8 diff --git a/src/conditioning/conditioner.hpp b/src/conditioning/conditioner.hpp index 7bbc383484..a146716d93 100644 --- a/src/conditioning/conditioner.hpp +++ b/src/conditioning/conditioner.hpp @@ -2008,7 +2008,8 @@ struct LLMEmbedder : public Conditioner { arch = LLM::LLMArch::GPT_OSS_20B; } else if (sd_version_is_pid(version)) { arch = LLM::LLMArch::GEMMA2_2B; - } else if (version == VERSION_QWEN_IMAGE_2_1 || + } else if (version == VERSION_FIBO2 || + version == VERSION_QWEN_IMAGE_2_1 || sd_version_is_flux3(version) || sd_version_is_lingbot_video(version) || sd_version_is_ideogram4(version) || @@ -2147,10 +2148,11 @@ struct LLMEmbedder : public Conditioner { std::tuple, std::vector, std::vector> tokenize(std::string text, const std::pair& attn_range, - size_t min_length = 0, - size_t max_length = 100000000, - bool spell_quotes = false, - const std::string& suffix = "") { + size_t min_length = 0, + size_t max_length = 100000000, + bool spell_quotes = false, + const std::string& suffix = "", + const std::vector>>& image_embeds = {}) { std::vector> parsed_attention; if (attn_range.first >= 0 && attn_range.second > 0) { if (attn_range.first > 0) { @@ -2195,6 +2197,19 @@ struct LLMEmbedder : public Conditioner { weights.insert(weights.end(), curr_tokens.size(), curr_weight); } + for (const auto& image_embed : image_embeds) { + const int index = image_embed.first; + const int64_t count = image_embed.second.shape()[1]; + if (index < 0 || static_cast(index) >= tokens.size() || count < 1) { + LOG_ERROR("Invalid reference image token range"); + return {}; + } + const int token = tokens[index]; + const float weight = weights[index]; + tokens.insert(tokens.begin() + index, count - 1, token); + weights.insert(weights.begin() + index, count - 1, weight); + } + if (!suffix.empty()) { std::vector suffix_tokens; if (!tokenizer->encode(suffix, suffix_tokens) || suffix_tokens.size() >= max_length) { @@ -2215,6 +2230,17 @@ struct LLMEmbedder : public Conditioner { std::vector mask; tokenizer->pad_tokens(tokens, &weights, &mask, min_length, max_length); + // Keep masks and weights in expanded coordinates, but only look up one token per image. + for (auto it = image_embeds.rbegin(); it != image_embeds.rend(); ++it) { + const int64_t begin = it->first; + const int64_t end = begin + it->second.shape()[1]; + if (begin < 0 || end > static_cast(tokens.size())) { + LOG_ERROR("Reference image exceeds the tokenized context"); + return {}; + } + tokens.erase(tokens.begin() + begin + 1, tokens.begin() + end); + } + // for (int i = 0; i < tokens.size(); i++) { // std::cout << tokens[i] << ":" << weights[i] << ", " << i << std::endl; // } @@ -2237,7 +2263,7 @@ struct LLMEmbedder : public Conditioner { const std::vector& image_grids = {}, const std::string& prompt_suffix = "", sd::Tensor* output_mask = nullptr) { - auto tokens_weights_mask = tokenize(prompt, prompt_attn_range, min_length, max_length, spell_quotes, prompt_suffix); + auto tokens_weights_mask = tokenize(prompt, prompt_attn_range, min_length, max_length, spell_quotes, prompt_suffix, image_embeds); auto& tokens = std::get<0>(tokens_weights_mask); auto& weights = std::get<1>(tokens_weights_mask); auto& mask = std::get<2>(tokens_weights_mask); @@ -2267,10 +2293,10 @@ struct LLMEmbedder : public Conditioner { sd::Tensor position_ids; if (sd_version_is_flux3(version)) { - position_ids = sd::Tensor::zeros({static_cast(tokens.size() * 4)}); + position_ids = sd::Tensor::zeros({static_cast(weights.size() * 4)}); for (size_t axis = 0; axis < 4; ++axis) { - for (size_t i = 0; i < tokens.size(); ++i) { - position_ids[axis * tokens.size() + i] = mask[i] == 0.f ? 1 : static_cast(i); + for (size_t i = 0; i < weights.size(); ++i) { + position_ids[axis * weights.size() + i] = mask[i] == 0.f ? 1 : static_cast(i); } } } @@ -2367,6 +2393,50 @@ struct LLMEmbedder : public Conditioner { } } + bool expand_vision_marker(const std::string& prompt, + size_t& marker_start, + std::vector> outputs, + int grid_h, + int grid_w, + std::vector>>& image_embeds, + LLM::DeepStackImageEmbeds& deepstack_image_embeds, + std::vector& image_grids, + int max_length = std::numeric_limits::max()) { + if (outputs.empty()) { + return false; + } + const std::string placeholder = "<|image_pad|>"; + size_t position = prompt.find(placeholder, marker_start); + if (position == std::string::npos) { + LOG_ERROR("Reference image marker is missing"); + return false; + } + std::vector prefix_tokens; + if (!tokenizer->encode(prompt.substr(0, position), prefix_tokens, nullptr)) { + return false; + } + int index = static_cast(prefix_tokens.size()); + for (const auto& image_embed : image_embeds) { + index += static_cast(image_embed.second.shape()[1]) - 1; + } + int count = static_cast(outputs[0].shape()[1]); + if (index > max_length || count > max_length - index) { + LOG_ERROR("Reference images exceed the %d-token context", max_length); + return false; + } + if (deepstack_image_embeds.empty()) { + deepstack_image_embeds.resize(outputs.size() - 1); + } + GGML_ASSERT(deepstack_image_embeds.size() == outputs.size() - 1); + image_embeds.emplace_back(index, std::move(outputs[0])); + for (size_t layer = 1; layer < outputs.size(); ++layer) { + deepstack_image_embeds[layer - 1].emplace_back(index, std::move(outputs[layer])); + } + image_grids.push_back({index, count, grid_h, grid_w}); + marker_start = position + placeholder.size(); + return true; + } + SDCondition get_learned_condition(int n_threads, const ConditionerParams& conditioner_params) override { std::string prompt; @@ -2394,9 +2464,8 @@ struct LLMEmbedder : public Conditioner { prompt_attn_range = {0, 0}; if (llm->enable_vision) { - const std::string placeholder = "<|image_pad|>"; - const int patch_size = llm->config.vision.patch_size; - const int factor = patch_size * llm->config.vision.spatial_merge_size; + const int patch_size = llm->config.vision.patch_size; + const int factor = patch_size * llm->config.vision.spatial_merge_size; auto resize_for_vision = [&](const sd::Tensor& image) { int height = static_cast(image.shape()[1]); @@ -2424,29 +2493,11 @@ struct LLMEmbedder : public Conditioner { int grid_h, int grid_w) { GGML_ASSERT(image_outputs.size() == 4); - auto image_embed = std::move(image_outputs[0]); prompt += "<|vision_start|>"; - std::vector prefix_tokens; - if (!tokenizer->encode(prompt, prefix_tokens, nullptr)) { - return false; - } - int image_embed_idx = static_cast(prefix_tokens.size()); - image_embeds.emplace_back(image_embed_idx, image_embed); - if (deepstack_image_embeds.empty()) { - deepstack_image_embeds.resize(image_outputs.size() - 1); - } - for (size_t layer = 0; layer < deepstack_image_embeds.size(); ++layer) { - deepstack_image_embeds[layer].emplace_back(image_embed_idx, std::move(image_outputs[layer + 1])); - } - image_grids.push_back({image_embed_idx, - static_cast(image_embed.shape()[1]), - grid_h, - grid_w}); - for (int64_t i = 0; i < image_embed.shape()[1]; ++i) { - prompt += placeholder; - } - prompt += "<|vision_end|>"; - return true; + size_t marker_start = prompt.size(); + prompt += "<|image_pad|><|vision_end|>"; + return expand_vision_marker(prompt, marker_start, std::move(image_outputs), grid_h, grid_w, + image_embeds, deepstack_image_embeds, image_grids); }; const auto* references = conditioner_params.minimax_h3_references; @@ -2609,14 +2660,13 @@ struct LLMEmbedder : public Conditioner { return {}; } int image_embed_idx = static_cast(prefix_tokens.size()); + for (const auto& previous_image : image_embeds) { + image_embed_idx += static_cast(previous_image.second.shape()[1]) - 1; + } image_embeds.emplace_back(image_embed_idx, image_embed); img_prompt += "<|vision_start|>"; - int64_t num_image_tokens = image_embed.shape()[1]; - img_prompt.reserve(img_prompt.size() + static_cast(num_image_tokens) * placeholder.size() + 32); - for (int j = 0; j < num_image_tokens; j++) { - img_prompt += placeholder; - } + img_prompt += placeholder; img_prompt += "<|vision_end|>"; } prompt += img_prompt; @@ -2660,26 +2710,14 @@ struct LLMEmbedder : public Conditioner { return {}; } prompt += (i == 0 ? "" : " ") + std::string("<|vision_start|>"; - std::vector prefix_tokens; - if (!tokenizer->encode(prompt, prefix_tokens, nullptr)) { + size_t marker_start = prompt.size(); + prompt += "<|image_pad|><|vision_end|>"; + if (!expand_vision_marker(prompt, marker_start, std::move(outputs), + static_cast(height) / llm->config.vision.patch_size, + static_cast(width) / llm->config.vision.patch_size, + image_embeds, deepstack_image_embeds, image_grids)) { return {}; } - int index = static_cast(prefix_tokens.size()); - int count = static_cast(outputs[0].shape()[1]); - image_embeds.emplace_back(index, std::move(outputs[0])); - if (deepstack_image_embeds.empty()) { - deepstack_image_embeds.resize(outputs.size() - 1); - } - for (size_t layer = 1; layer < outputs.size(); ++layer) { - deepstack_image_embeds[layer - 1].emplace_back(index, std::move(outputs[layer])); - } - image_grids.push_back({index, count, - static_cast(height) / llm->config.vision.patch_size, - static_cast(width) / llm->config.vision.patch_size}); - for (int j = 0; j < count; ++j) { - prompt += "<|image_pad|>"; - } - prompt += "<|vision_end|>"; } } prompt_attn_range.first = static_cast(prompt.size()); @@ -2690,7 +2728,6 @@ struct LLMEmbedder : public Conditioner { if (llm->enable_vision && conditioner_params.ref_images != nullptr && !conditioner_params.ref_images->empty()) { LOG_INFO("%s", sd_version_is_mage_flow(version) ? "MageFlowEditPipeline" : "QwenImageEditPlusPipeline"); prompt_template_encode_start_idx = 64; - int image_embed_idx = 64 + 6; int min_pixels = conditioner_params.ref_image_params.vlm_min_size; if (min_pixels <= 0) { @@ -2707,6 +2744,7 @@ struct LLMEmbedder : public Conditioner { } } + prompt = "<|im_start|>system\nDescribe the key features of the input image (color, shape, size, texture, objects, background), then explain how the user's text instruction should alter or modify the image. Generate a new image that meets the user's requirements while maintaining consistency with the original input where appropriate.<|im_end|>\n<|im_start|>user\n"; std::string placeholder = "<|image_pad|>"; std::string img_prompt; @@ -2726,19 +2764,22 @@ struct LLMEmbedder : public Conditioner { auto image_embed = llm->encode_image(n_threads, resized_image, false); GGML_ASSERT(!image_embed.empty()); + std::string image_prefix = prompt + img_prompt + (sd_version_is_mage_flow(version) ? "Image " : "Picture ") + std::to_string(i + 1) + ": <|vision_start|>"; + std::vector prefix_tokens; + if (!tokenizer->encode(image_prefix, prefix_tokens, nullptr)) { + return {}; + } + int image_embed_idx = static_cast(prefix_tokens.size()); + for (const auto& previous_image : image_embeds) { + image_embed_idx += static_cast(previous_image.second.shape()[1]) - 1; + } image_embeds.emplace_back(image_embed_idx, image_embed); - image_embed_idx += 1 + static_cast(image_embed.shape()[1]) + 6; img_prompt += (sd_version_is_mage_flow(version) ? "Image " : "Picture ") + std::to_string(i + 1) + ": <|vision_start|>"; - int64_t num_image_tokens = image_embed.shape()[1]; - img_prompt.reserve(num_image_tokens * placeholder.size()); - for (int j = 0; j < num_image_tokens; j++) { - img_prompt += placeholder; - } + img_prompt += placeholder; img_prompt += "<|vision_end|>"; } - prompt = "<|im_start|>system\nDescribe the key features of the input image (color, shape, size, texture, objects, background), then explain how the user's text instruction should alter or modify the image. Generate a new image that meets the user's requirements while maintaining consistency with the original input where appropriate.<|im_end|>\n<|im_start|>user\n"; prompt += img_prompt; prompt_attn_range.first = static_cast(prompt.size()); @@ -2816,14 +2857,13 @@ struct LLMEmbedder : public Conditioner { return {}; } int image_embed_idx = static_cast(prefix_tokens.size()); + for (const auto& previous_image : image_embeds) { + image_embed_idx += static_cast(previous_image.second.shape()[1]) - 1; + } image_embeds.emplace_back(image_embed_idx, image_embed); img_prompt += "<|vision_start|>"; - int64_t num_image_tokens = image_embed.shape()[1]; - img_prompt.reserve(img_prompt.size() + static_cast(num_image_tokens) * placeholder.size() + 32); - for (int j = 0; j < num_image_tokens; j++) { - img_prompt += placeholder; - } + img_prompt += placeholder; img_prompt += "<|vision_end|>"; } @@ -2888,14 +2928,13 @@ struct LLMEmbedder : public Conditioner { return {}; } int image_embed_idx = static_cast(prefix_tokens.size()); + for (const auto& previous_image : image_embeds) { + image_embed_idx += static_cast(previous_image.second.shape()[1]) - 1; + } image_embeds.emplace_back(image_embed_idx, image_embed); img_prompt += "Picture " + std::to_string(i + 1) + ": <|vision_start|>"; - int64_t num_image_tokens = image_embed.shape()[1]; - img_prompt.reserve(img_prompt.size() + static_cast(num_image_tokens) * placeholder.size() + 32); - for (int j = 0; j < num_image_tokens; j++) { - img_prompt += placeholder; - } + img_prompt += placeholder; img_prompt += "<|vision_end|>"; } prompt += img_prompt; @@ -2913,7 +2952,6 @@ struct LLMEmbedder : public Conditioner { LOG_INFO("LongCatEditPipeline"); prompt_template_encode_start_idx = 67; min_length = 512 + prompt_template_encode_start_idx; - int image_embed_idx = 36 + 6; int min_pixels = conditioner_params.ref_image_params.vlm_min_size; if (min_pixels <= 0) { @@ -2930,6 +2968,7 @@ struct LLMEmbedder : public Conditioner { } } + prompt = "<|im_start|>system\nAs an image editing expert, first analyze the content and attributes of the input image(s). Then, based on the user's editing instructions, clearly and precisely determine how to modify the given image(s), ensuring that only the specified parts are altered and all other aspects remain consistent with the original(s).<|im_end|>\n<|im_start|>user\n"; std::string placeholder = "<|image_pad|>"; std::string img_prompt; @@ -2948,19 +2987,22 @@ struct LLMEmbedder : public Conditioner { auto resized_image = clip_preprocess(image, w_bar, h_bar); auto image_embed = llm->encode_image(n_threads, resized_image, false); GGML_ASSERT(!image_embed.empty()); + std::string image_prefix = prompt + img_prompt + "<|vision_start|>"; + std::vector prefix_tokens; + if (!tokenizer->encode(image_prefix, prefix_tokens, nullptr)) { + return {}; + } + int image_embed_idx = static_cast(prefix_tokens.size()); + for (const auto& previous_image : image_embeds) { + image_embed_idx += static_cast(previous_image.second.shape()[1]) - 1; + } image_embeds.emplace_back(image_embed_idx, image_embed); - image_embed_idx += 1 + static_cast(image_embed.shape()[1]) + 6; img_prompt += "<|vision_start|>"; - int64_t num_image_tokens = image_embed.shape()[1]; - img_prompt.reserve(num_image_tokens * placeholder.size()); - for (int j = 0; j < num_image_tokens; j++) { - img_prompt += placeholder; - } + img_prompt += placeholder; img_prompt += "<|vision_end|>"; } - prompt = "<|im_start|>system\nAs an image editing expert, first analyze the content and attributes of the input image(s). Then, based on the user's editing instructions, clearly and precisely determine how to modify the given image(s), ensuring that only the specified parts are altered and all other aspects remain consistent with the original(s).<|im_end|>\n<|im_start|>user\n"; prompt += img_prompt; } else { prompt_template_encode_start_idx = 36; @@ -3028,7 +3070,10 @@ struct LLMEmbedder : public Conditioner { } else if (sd_version_is_z_image(version) || sd_version_is_z_image_l2p(version)) { prompt_template_encode_start_idx = 0; out_layers = {35}; // -2 - + if (conditioner_params.ref_images != nullptr && !conditioner_params.ref_images->empty()) { + LOG_ERROR("Z-Image currently does not support reference-image conditioning"); + return {}; + } if (conditioner_params.ref_images != nullptr && !conditioner_params.ref_images->empty()) { LOG_INFO("ZImageOmniPipeline"); prompt = "<|im_start|>user\n<|vision_start|>"; @@ -3563,12 +3608,10 @@ struct LLaDAImageEmbedder : public Conditioner { auto inputs_embeds = llm->compute_input_embeds(n_threads, text_ids); auto query_embeds = query_former->compute(n_threads, inputs_embeds); - // splice_image_embeds() replaces tokens in place, so the query slots have to exist in - // input_ids; their ids are irrelevant because the embeddings are overwritten. - std::vector padded(tokens.begin(), tokens.end()); - padded.resize(static_cast(n_text + num_queries), tokenizer->PAD_TOKEN_ID); - int64_t n_total = static_cast(padded.size()); - sd::Tensor input_ids({n_total}, padded); + std::vector input_tokens(tokens.begin(), tokens.end()); + input_tokens.push_back(tokenizer->PAD_TOKEN_ID); + sd::Tensor input_ids({n_text + 1}, std::move(input_tokens)); + int64_t n_total = n_text + num_queries; // Bidirectional everywhere except that the text tokens must not see the appended // queries, matching backbone_attention_mask[:, :, :text_length, text_length:] = min. diff --git a/src/conditioning/fibo2.hpp b/src/conditioning/fibo2.hpp new file mode 100644 index 0000000000..43b2ee0845 --- /dev/null +++ b/src/conditioning/fibo2.hpp @@ -0,0 +1,182 @@ +#ifndef __SD_CONDITIONING_FIBO2_HPP__ +#define __SD_CONDITIONING_FIBO2_HPP__ + +#include "conditioning/conditioner.hpp" + +struct Fibo2Conditioner : public LLMEmbedder { + using JSON = nlohmann::ordered_json; + + Fibo2Conditioner(ggml_backend_t backend, const String2TensorStorage& weights, std::shared_ptr weight_manager, const TokenizerConfig& tokenizers) + : LLMEmbedder(backend, weights, VERSION_FIBO2, "", true, weight_manager, tokenizers) { + if (llm->config.hidden_size != 2560 || llm->config.num_layers != 36) { + throw std::runtime_error("Fibo 2.0 requires its Qwen3-VL-4B text encoder"); + } + } + + static bool empty_value(const JSON& value) { + return value.is_null() || + (value.is_string() && value.get_ref().empty()) || + ((value.is_array() || value.is_object()) && value.empty()); + } + + static JSON clean_caption(const JSON& value) { + if (value.is_object()) { + auto result = JSON::object(); + for (auto it = value.begin(); it != value.end(); ++it) { + auto cleaned = clean_caption(it.value()); + if (!empty_value(cleaned)) { + result[it.key()] = std::move(cleaned); + } + } + return result; + } + if (value.is_array()) { + auto result = JSON::array(); + for (const auto& item : value) { + auto cleaned = clean_caption(item); + if (!empty_value(cleaned)) { + result.push_back(std::move(cleaned)); + } + } + return result; + } + return value; + } + + static std::pair parse_prompt(const std::string& prompt) { + auto caption = JSON::parse(prompt, nullptr, false); + if (!caption.is_object()) { + return {prompt, nullptr}; + } + JSON instruction; + if (caption.contains("edit_instruction")) { + instruction = caption["edit_instruction"]; + caption.erase("edit_instruction"); + } + if (caption.contains("structured_caption")) { + caption = JSON(caption["structured_caption"]); + if (!caption.is_object()) { + return {caption, instruction}; + } + } + bool short_form = caption.size() == 2 && caption.contains("short_description") && caption.contains("text_render"); + caption = clean_caption(caption); + for (const char* score : {"aesthetic_score", "preference_score"}) { + caption.erase(score); + for (auto& value : caption) { + if (value.is_object()) { + value.erase(score); + } + } + } + caption = clean_caption(caption); + if (short_form) { + caption = JSON{{"short_description", caption.value("short_description", JSON(""))}, + {"text_render", caption.value("text_render", JSON::array())}}; + } + return {caption, instruction}; + } + + static std::string format_prompt(const std::string& text, int num_images) { + auto [caption, instruction] = parse_prompt(text); + if (num_images > 0) { + if (instruction.is_null()) { + caption = nullptr; + instruction = text; + } + auto edit = JSON::object(); + for (int i = 0; i < num_images; ++i) { + edit[num_images == 1 ? "image" : "image_" + std::to_string(i + 1)] = "<|vision_start|><|image_pad|><|vision_end|>"; + } + if (!empty_value(caption)) { + edit["structured_caption"] = caption; + } + edit["edit_instruction"] = instruction; + return edit.dump(); + } + if (text.empty()) { + return R"({"short_description":"","objects":[],"background_setting":"","lighting":{},"aesthetics":{},"photographic_characteristics":{},"style_medium":"","context":"","artistic_style":""})"; + } + if (caption.is_string()) { + return caption.get(); + } + return empty_value(caption) ? "{}" : caption.dump(); + } + + static sd::Tensor make_bundles(sd::Tensor states, const std::set& layers) { + constexpr int width = 2560; + constexpr int bundles[6][3] = {{9, 20, 31}, {5, 16, 27}, {7, 18, 29}, {9, 20, 31}, {11, 22, 33}, {13, 24, 35}}; + for (int64_t token = 0; token < states.shape()[1]; ++token) { + for (size_t layer = 0; layer < layers.size(); ++layer) { + float sum = 0.f; + int64_t offset = token * states.shape()[0] + layer * width; + for (int i = 0; i < width; ++i) { + sum += states[offset + i] * states[offset + i]; + } + float scale = 1.f / std::sqrt(sum / width + 1e-6f); + for (int i = 0; i < width; ++i) { + states[offset + i] *= scale; + } + } + } + sd::Tensor result({6 * 3 * width, states.shape()[1]}); + for (int b = 0; b < 6; ++b) { + for (int j = 0; j < 3; ++j) { + int index = static_cast(std::distance(layers.begin(), layers.find(bundles[b][j]))); + for (int64_t token = 0; token < states.shape()[1]; ++token) { + std::copy_n(states.data() + token * states.shape()[0] + index * width, width, + result.data() + token * result.shape()[0] + (b * 3 + j) * width); + } + } + } + return result; + } + + SDCondition get_learned_condition(int n_threads, const ConditionerParams& p) override { + int image_count = p.ref_images ? static_cast(p.ref_images->size()) : 0; + if (image_count && !llm->enable_vision) { + LOG_ERROR("Fibo 2.0 editing requires vision weights"); + return {}; + } + std::string prompt = "<|im_start|>user\n" + format_prompt(p.text, image_count) + "<|im_end|>\n<|im_start|>assistant\n"; + std::vector>> images; + LLM::DeepStackImageEmbeds deepstack; + std::vector grids; + size_t marker_start = 0; + for (int i = 0; i < image_count; ++i) { + const auto& image = (*p.ref_images)[i]; + int factor = llm->config.vision.patch_size * llm->config.vision.spatial_merge_size; + int height = static_cast(image.shape()[1]), width = static_cast(image.shape()[0]); + int h = std::max(factor, static_cast(std::round(static_cast(height) / factor)) * factor); + int w = std::max(factor, static_cast(std::round(static_cast(width) / factor)) * factor); + resize_image_dims(height, width, h, w, factor, + p.ref_image_params.vlm_min_size, + p.ref_image_params.vlm_max_size, + p.ref_image_params.vlm_resize_mode); + auto pixels = sd::ops::interpolate(image, {w, h, image.shape()[2], image.shape()[3]}, sd::ops::InterpolateMode::Bicubic, false, true); + auto rgb = sd::Tensor({w, h, 3, 1}); + for (int64_t pixel = 0; pixel < int64_t(w) * h; ++pixel) { + float alpha = pixels.shape()[2] == 4 ? pixels[pixel + 3 * w * h] : 1.f; + for (int channel = 0; channel < 3; ++channel) { + rgb[pixel + channel * w * h] = 2.f * std::clamp(pixels[pixel + channel * w * h] * alpha + 1.f - alpha, 0.f, 1.f) - 1.f; + } + } + auto outputs = llm->encode_image_outputs(n_threads, rgb, false); + if (!expand_vision_marker(prompt, marker_start, std::move(outputs), + h / llm->config.vision.patch_size, w / llm->config.vision.patch_size, + images, deepstack, grids, 4096)) { + return {}; + } + } + const std::set layers = {5, 7, 9, 11, 13, 16, 18, 20, 22, 24, 27, 29, 31, 33, 35}; + auto states = encode_prompt(n_threads, prompt, {0, 0}, 0, 0, images, layers, 0, false, 4096, deepstack, grids); + if (states.empty()) { + return {}; + } + SDCondition result; + result.c_crossattn = make_bundles(std::move(states), layers); + return result; + } +}; + +#endif // __SD_CONDITIONING_FIBO2_HPP__ diff --git a/src/core/ggml_extend.cpp b/src/core/ggml_extend.cpp index 80e2121c6d..46b8e7c54c 100644 --- a/src/core/ggml_extend.cpp +++ b/src/core/ggml_extend.cpp @@ -460,7 +460,8 @@ ggml_tensor* ggml_ext_conv_2d(ggml_context* ctx, p1 = 0; } - if (direct) { + // The direct kernels take F16 / F32 weights only (ggml-cpu asserts on anything else); other types keep im2col. + if (direct && (w->type == GGML_TYPE_F16 || w->type == GGML_TYPE_F32)) { x = ggml_conv_2d_direct(ctx, w, x, s0, s1, p0, p1, d0, d1); } else { // ggml_conv_2d, but skips the final [OC, N] -> [N, OC] relayout when N == 1 (no data moves). @@ -503,7 +504,7 @@ ggml_tensor* ggml_ext_conv_3d(ggml_context* ctx, if (scale != 1.f) { x = ggml_ext_scale(ctx, x, scale); } - if (direct) { + if (direct && (w->type == GGML_TYPE_F16 || w->type == GGML_TYPE_F32)) { int64_t OC = w->ne[3] / IC; int64_t N = x->ne[3] / IC; x = ggml_conv_3d_direct(ctx, w, x, s0, s1, s2, p0, p1, p2, d0, d1, d2, (int)IC, (int)N, (int)OC); diff --git a/src/core/ggml_extend_backend.cpp b/src/core/ggml_extend_backend.cpp index 9eafe6f1de..f1ff3795a1 100644 --- a/src/core/ggml_extend_backend.cpp +++ b/src/core/ggml_extend_backend.cpp @@ -296,6 +296,10 @@ bool add_rpc_devices(const std::string& servers) { return false; } ggml_backend_reg_t rpc_reg = ggml_backend_reg_by_name("RPC"); + if (!rpc_reg) { + ggml_backend_load_all(); + rpc_reg = ggml_backend_reg_by_name("RPC"); + } if (!rpc_reg) { LOG_ERROR("RPC backend not found, cannot add RPC servers"); return false; @@ -309,12 +313,19 @@ bool add_rpc_devices(const std::string& servers) { for (const auto& server : rpc_servers) { LOG_INFO("Adding RPC server: %s", server.c_str()); auto reg = ggml_backend_rpc_add_server_fn(server.c_str()); - // no return value to check for success but should print errors from the RPC backend if it fails to add the server + if (!reg) { + LOG_ERROR("failed to add RPC server: %s", server.c_str()); + return false; + } ggml_backend_register(reg); } return true; } +bool sd_add_rpc_servers(const char* servers) { + return add_rpc_devices(servers ? servers : ""); +} + static void ggml_backend_load_all_once() { // If the registry already has devices and the CPU backend is present, // assume either static registration or explicit host-side preloading has diff --git a/src/core/probe/conv_direct.cpp b/src/core/probe/conv_direct.cpp index 7f4666a41f..d6060856d7 100644 --- a/src/core/probe/conv_direct.cpp +++ b/src/core/probe/conv_direct.cpp @@ -8,6 +8,7 @@ #include "core/ggml_extend_backend.h" #include "core/probe/backend_probe.h" #include "core/util.h" +#include "ggml-cpu.h" static void build_conv_direct_probe(BackendProbeGraph& graph, bool conv3d, bool direct, int size = 32, int channels = 64) { if (conv3d) { @@ -83,6 +84,22 @@ bool sd_probe_conv_direct(ggml_backend_t backend, bool conv3d) { dims, name, result.candidate_faster ? "on" : "off", result.reference_ms, result.candidate_ms); return result.candidate_faster; } + if (sd_backend_is_cpu(backend)) { +#ifdef GGML_CPU_CONV_DIRECT_ROW_RUNS + // ggml-cpu's direct kernels work through a fixed 16 MB patch buffer instead of materialising the whole + // im2col matrix, which in a VAE decoder at 1024x1024 is 2.8 GB for one 128-channel conv and 5.6 GB for a + // 256-channel one. With F16 weights (every Conv2d / Conv3d block) they produce the same values as im2col + + // mul_mat, bit for bit; F32 weights (runtime LoRA tensors) and depthwise convs keep F32 activations where + // im2col rounds them to F16, so those come out closer to the exact result rather than identical. + LOG_INFO("conv%dd direct auto: %s -> on (CPU)", dims, name); + return true; +#else + // Without the row-run kernels (scripts/unsloth/ggml-patches/0014, applied by the prebuilt workflow) the + // per-element direct loops are slower than im2col on large convolutions. + LOG_INFO("conv%dd direct auto: %s -> off (CPU, ggml without the row-run direct kernels)", dims, name); + return false; +#endif + } if (sd_backend_supports_conv_implicit_gemm(backend, conv3d)) { LOG_INFO("conv%dd direct auto: %s -> on (implicit GEMM)", dims, name); return true; diff --git a/src/core/util.cpp b/src/core/util.cpp index c1d463e4da..080ca53414 100644 --- a/src/core/util.cpp +++ b/src/core/util.cpp @@ -29,6 +29,11 @@ #include #endif +#ifdef __linux__ +#include +#include +#endif + #include "ggml-backend.h" #include "ggml.h" #include "stable-diffusion.h" @@ -340,20 +345,47 @@ bool MmapWrapper::copy_data(void* buf, size_t n, size_t offset) const { // LICENSE: https://github.com/ggerganov/llama.cpp/blob/master/LICENSE int32_t sd_get_num_physical_cores() { #ifdef __linux__ - // enumerate the set of thread siblings, num entries is num cores - std::unordered_set siblings; - for (uint32_t cpu = 0; cpu < UINT32_MAX; ++cpu) { - std::ifstream thread_siblings("/sys/devices/system/cpu" + std::to_string(cpu) + "/topology/thread_siblings"); - if (!thread_siblings.is_open()) { - break; // no more cpus + // One thread per physical core this process may run on: the distinct thread-sibling sets among the CPUs in + // the affinity mask (a taskset / container cpuset leaves fewer CPUs than the machine has). Upstream's loop + // counts every CPU of the machine, so a 4-CPU cpuset on a 192-CPU host still gets 96 threads. + // The mask is sized dynamically: a kernel built for more than CPU_SETSIZE (1024) CPUs rejects a smaller one + // with EINVAL, so grow it until the call succeeds. + for (int ncpus = CPU_SETSIZE; ncpus <= (1 << 20); ncpus *= 2) { + cpu_set_t* mask = CPU_ALLOC(ncpus); + const size_t size = CPU_ALLOC_SIZE(ncpus); + if (mask == nullptr) { + break; } - std::string line; - if (std::getline(thread_siblings, line)) { - siblings.insert(line); + CPU_ZERO_S(size, mask); + if (sched_getaffinity(0, size, mask) != 0) { + CPU_FREE(mask); + if (errno == EINVAL) { + continue; + } + break; } - } - if (siblings.size() > 0) { - return static_cast(siblings.size()); + std::unordered_set siblings; + int allowed = 0; + for (int cpu = 0; cpu < ncpus; ++cpu) { + if (!CPU_ISSET_S(cpu, size, mask)) { + continue; + } + allowed++; + std::ifstream thread_siblings("/sys/devices/system/cpu/cpu" + std::to_string(cpu) + "/topology/thread_siblings"); + std::string line; + if (thread_siblings.is_open() && std::getline(thread_siblings, line) && !line.empty()) { + siblings.insert(line); + } + } + CPU_FREE(mask); + if (!siblings.empty()) { + return static_cast(siblings.size()); + } + if (allowed > 0) { + // No topology in sysfs (some containers): assume 2-way SMT above 4 CPUs, as below. + return allowed <= 4 ? allowed : allowed / 2; + } + break; } #elif defined(__APPLE__) && defined(__MACH__) int32_t num_physical_cores; @@ -568,50 +600,59 @@ bool parse_strict_bool(const std::string& text, bool& value) { } static std::string build_progress_bar(int step, int steps, char progress_char = '=', bool show_head = true) { - std::string progress = " |"; - int max_progress = 50; - int32_t current = 0; + std::string progress = " |"; + constexpr int max_progress = 50; + int32_t current = 0; if (steps > 0) { - current = (int32_t)(step * 1.f * max_progress / steps); - } - for (int i = 0; i < 50; i++) { - if (i > current) { - progress += " "; - } else if (show_head && i == current && i != max_progress - 1) { - progress += ">"; - } else { - progress += progress_char; - } + current = (int32_t)(std::clamp(step, 0, steps) * 1.f * max_progress / steps); } - progress += "|"; + int head_length = (current == max_progress) ? 0 : 1; + progress += std::string(current, progress_char) + std::string(head_length, show_head ? '>' : progress_char) + std::string(max_progress - current - head_length, ' ') + "|"; return progress; } -static void print_progress_line(int step, int steps, const std::string& speed_text, char progress_char = '=', bool show_head = true) { +static void print_progress_line(int step, int steps, const std::string& speed_text, char progress_char = '=', bool show_head = true, const std::string& eta_text = "") { if (step == 0) { return; } std::string progress = build_progress_bar(step, steps, progress_char, show_head); + int step_width = static_cast(std::to_string(steps).size()); const char* lf = (step == steps ? "\n" : ""); - printf("\r%s %i/%i - %s\033[K%s", progress.c_str(), step, steps, speed_text.c_str(), lf); + printf("\r%s %*i/%i - %10s%s%s\033[K%s", progress.c_str(), step_width, step, steps, speed_text.c_str(), + eta_text.empty() ? "" : " - ETA ", eta_text.c_str(), lf); fflush(stdout); // for linux } +static int show_eta = -1; // -1 - estimate duration, 0 - hide, 1 - show +static float passed_time = 0.f; + void pretty_progress(int step, int steps, float time) { if (sd_progress_cb) { sd_progress_cb(step, steps, time, sd_progress_cb_data); return; } if (step == 0) { + show_eta = -1; // reinitialize for batch generations + passed_time = 0.f; return; } + std::string remainder_text = ""; + if (show_eta && step < steps) { // if timer is not disabled and bar is not filled, + if (show_eta == -1 && step) // estimate the total time once, then + show_eta = (std::abs(time * steps / step) < 15.f) ? 0 : 1; // hide the timer if total is less than 15 seconds + if (show_eta) { + unsigned remainder = (passed_time += time) / step * (steps - step); // average speed; time*(steps-step) for current + remainder_text = remainder < 3600 ? sd_format("%02u:%02u", remainder / 60, remainder % 60) + : sd_format("%u:%02u:%02u", remainder / 3600, (remainder / 60) % 60, remainder % 60); + } + } const char* unit = "s/it"; float speed = time; if (speed < 1.0f && speed > 0.f) { speed = 1.0f / speed; unit = "it/s"; } - print_progress_line(step, steps, sd_format("%.2f%s", speed, unit)); + print_progress_line(step, steps, sd_format("%.2f%s", speed, unit), '=', true, remainder_text); } void pretty_bytes_progress(int step, int steps, uint64_t bytes_processed, float elapsed_seconds) { diff --git a/src/detailer.cpp b/src/detailer.cpp index 72d01dc6f0..ec4d6065a3 100644 --- a/src/detailer.cpp +++ b/src/detailer.cpp @@ -291,7 +291,9 @@ static void filter_and_sort_detections(std::vector* detectio std::partial_sort(detections->begin(), detections->begin() + params.mask_k_largest, detections->end(), - [](const auto& a, const auto& b) { return detection_area(a) > detection_area(b); }); + [&](const auto& a, const auto& b) { + return params.mask_filter_by_confidence ? a.confidence > b.confidence : detection_area(a) > detection_area(b); + }); detections->resize(params.mask_k_largest); } @@ -545,7 +547,16 @@ static sd_image_t resize_crop_mask(const Mask& source, const CropRegion& crop, i return output; } -static sd_image_t copy_as_rgb(const sd_image_t& source) { +static sd_image_t copy_for_detailing(const sd_image_t& source) { + if (source.channel == 4) { + sd_image_t output = source; + size_t bytes = static_cast(source.width) * source.height * source.channel; + output.data = static_cast(malloc(bytes)); + if (output.data != nullptr) { + memcpy(output.data, source.data, bytes); + } + return output; + } CropRegion full{0, 0, static_cast(source.width), static_cast(source.height)}; return resize_crop_image(source, full, source.width, source.height); } @@ -566,7 +577,7 @@ static void composite_crop(sd_image_t* destination, continue; } float source_x = (x - crop.x1 + 0.5f) * scale_x - 0.5f; - size_t destination_index = (static_cast(y) * destination->width + x) * 3; + size_t destination_index = (static_cast(y) * destination->width + x) * destination->channel; for (int c = 0; c < 3; ++c) { float generated_value = bilinear_channel(generated, source_x, source_y, c) * 255.f; float original_value = destination->data[destination_index + c]; @@ -787,6 +798,10 @@ static bool parse_adetailer_params(const sd_adetailer_params_t& public_params, ok = parse_int(value, ¶ms->max_detections); } else if (key == "mask_k_largest") { ok = parse_int(value, ¶ms->mask_k_largest); + } else if (key == "mask_filter_method") { + const std::string method = lower_copy(value); + ok = method == "area" || method == "confidence"; + params->mask_filter_by_confidence = method == "confidence"; } else if (key == "mask_min_ratio") { ok = parse_float(value, ¶ms->mask_min_ratio); } else if (key == "mask_max_ratio") { @@ -805,6 +820,8 @@ static bool parse_adetailer_params(const sd_adetailer_params_t& public_params, ok = parse_int(value, ¶ms->mask_blur); } else if (key == "inpaint_padding") { ok = parse_int(value, ¶ms->inpaint_padding); + } else if (key == "inpaint_only_masked") { + ok = parse_bool(value, ¶ms->inpaint_only_masked); } else if (key == "inpaint_width") { ok = parse_int(value, ¶ms->inpaint_width); } else if (key == "inpaint_height") { @@ -920,7 +937,7 @@ bool adetail_image(adetailer_ctx_t* context, filter_and_sort_detections(&detections, input_image.width, input_image.height, params); auto masks = make_masks(detections, input_image.width, input_image.height, params); - sd_image_t current = copy_as_rgb(input_image); + sd_image_t current = copy_for_detailing(input_image); if (current.data == nullptr) { return false; } @@ -939,12 +956,13 @@ bool adetail_image(adetailer_ctx_t* context, if (!mask_bbox(masks[i], &crop)) { continue; } - crop = expand_crop(crop, - current.width, - current.height, - params.inpaint_padding, - params.inpaint_width, - params.inpaint_height); + crop = params.inpaint_only_masked ? expand_crop(crop, + current.width, + current.height, + params.inpaint_padding, + params.inpaint_width, + params.inpaint_height) + : CropRegion{0, 0, static_cast(current.width), static_cast(current.height)}; sd_image_t local_image = resize_crop_image(current, crop, params.inpaint_width, params.inpaint_height); sd_image_t local_mask = resize_crop_mask(masks[i], crop, params.inpaint_width, params.inpaint_height); if (local_image.data == nullptr || local_mask.data == nullptr) { diff --git a/src/detailer.h b/src/detailer.h index 17a81a1dde..0e3ae50f5a 100644 --- a/src/detailer.h +++ b/src/detailer.h @@ -27,30 +27,32 @@ enum ADetailerSort { }; struct ADetailerParams { - const char* prompt = nullptr; - const char* negative_prompt = nullptr; - int input_size = 640; - float confidence = 0.3f; - float nms_threshold = 0.45f; - int max_detections = 100; - int mask_k_largest = 0; - float mask_min_ratio = 0.f; - float mask_max_ratio = 1.f; - int dilate_erode = 4; - int x_offset = 0; - int y_offset = 0; - bool merge_masks = false; - bool invert_mask = false; - int mask_blur = 4; - int inpaint_padding = 32; - int inpaint_width = 512; - int inpaint_height = 512; - float denoising_strength = 0.4f; - int steps = 0; - float cfg_scale = -1.f; - sample_method_t sample_method = SAMPLE_METHOD_COUNT; - scheduler_t scheduler = SCHEDULER_COUNT; - ADetailerSort sort_by = ADETAILER_SORT_NONE; + const char* prompt = nullptr; + const char* negative_prompt = nullptr; + int input_size = 640; + float confidence = 0.3f; + float nms_threshold = 0.45f; + int max_detections = 100; + int mask_k_largest = 0; + bool mask_filter_by_confidence = false; + float mask_min_ratio = 0.f; + float mask_max_ratio = 1.f; + int dilate_erode = 4; + int x_offset = 0; + int y_offset = 0; + bool merge_masks = false; + bool invert_mask = false; + int mask_blur = 4; + int inpaint_padding = 32; + bool inpaint_only_masked = true; + int inpaint_width = 512; + int inpaint_height = 512; + float denoising_strength = 0.4f; + int steps = 0; + float cfg_scale = -1.f; + sample_method_t sample_method = SAMPLE_METHOD_COUNT; + scheduler_t scheduler = SCHEDULER_COUNT; + ADetailerSort sort_by = ADETAILER_SORT_NONE; }; struct ADetailerGGML { diff --git a/src/model.h b/src/model.h index 0c1a37a66c..73a369301a 100644 --- a/src/model.h +++ b/src/model.h @@ -68,6 +68,7 @@ enum SDVersion { VERSION_IRIS, VERSION_FLUX3, VERSION_MOSS_IMAGE, + VERSION_FIBO2, VERSION_COUNT, }; @@ -293,7 +294,8 @@ static inline bool sd_version_uses_flux2_vae(SDVersion version) { sd_version_is_ideogram4(version) || sd_version_is_sefi_image(version) || sd_version_is_llada_image(version) || - version == VERSION_MOSS_IMAGE) { + version == VERSION_MOSS_IMAGE || + version == VERSION_FIBO2) { return true; } return false; @@ -351,7 +353,8 @@ static inline bool sd_version_is_dit(SDVersion version) { sd_version_is_sensenova_u1(version) || sd_version_is_pixart(version) || version == VERSION_IRIS || - version == VERSION_MOSS_IMAGE) { + version == VERSION_MOSS_IMAGE || + version == VERSION_FIBO2) { return true; } return false; @@ -385,7 +388,8 @@ static inline bool sd_version_supports_ref_image_conditioning(SDVersion version) sd_version_is_mage_flow(version) || sd_version_is_anima(version) || sd_version_is_sefi_image(version) || - sd_version_is_pid(version); + sd_version_is_pid(version) || + version == VERSION_FIBO2; } static bool sd_version_is_inpaint_or_unet_edit(SDVersion version) { diff --git a/src/model/diffusion/fibo2.hpp b/src/model/diffusion/fibo2.hpp new file mode 100644 index 0000000000..f53b6ecd1f --- /dev/null +++ b/src/model/diffusion/fibo2.hpp @@ -0,0 +1,395 @@ +#ifndef __SD_MODEL_DIFFUSION_FIBO2_HPP__ +#define __SD_MODEL_DIFFUSION_FIBO2_HPP__ + +#include "model/diffusion/z_image.hpp" + +// https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_bria_fibo2.py +namespace Fibo2 { + struct Fibo2Config { + int64_t dim = 3584; + int64_t channels = 128; + int64_t text_dim = 7680; + int64_t head_dim = 128; + int64_t depth = 44; + int64_t refiner_depth = 2; + int64_t perceiver_depth = 2; + int64_t max_gist = 512; + int64_t ffn_dim = 9536; + bool fused_qkv = false; + bool fused_kv = false; + bool fused_ffn = false; + std::vector injection_layers = {6, 14, 22, 30, 38}; + + static Fibo2Config detect_from_weights(const String2TensorStorage& weights, const std::string& prefix) { + Fibo2Config config; + int64_t depth = 0, refiner_depth = 0, perceiver_depth = 0; + std::set injection_layers; + for (const auto& [name, tensor] : weights) { + if (!starts_with(name, prefix + ".")) { + continue; + } + auto key = name.substr(prefix.size() + 1); + if (key == "x_embedder.weight") { + config.channels = tensor.ne[0]; + config.dim = tensor.ne[1]; + } else if (key == "perceiver.latent_queries") { + config.max_gist = tensor.ne[1]; + } else if (key == "perceiver.layers.0.cross_attention.to_k.weight" || key == "perceiver.layers.0.cross_attention.to_kv.weight") { + config.text_dim = tensor.ne[0]; + config.fused_kv = key == "perceiver.layers.0.cross_attention.to_kv.weight"; + } else if (key == "layers.0.attention.norm_q.weight") { + config.head_dim = tensor.ne[0]; + } else if (key == "layers.0.attention.qkv.weight") { + config.fused_qkv = true; + } else if (key == "layers.0.feed_forward.w1.weight") { + config.ffn_dim = tensor.ne[1]; + } else if (key == "layers.0.feed_forward.w13.weight") { + config.ffn_dim = tensor.ne[1] / 2; + config.fused_ffn = true; + } + if (starts_with(key, "layers.")) { + int index = std::stoi(key.substr(7)); + depth = std::max(depth, index + 1); + if (ends_with(key, ".cross_attn_gate")) { + injection_layers.insert(index); + } + } else if (starts_with(key, "noise_refiner.")) { + refiner_depth = std::max(refiner_depth, std::stoi(key.substr(14)) + 1); + } else if (starts_with(key, "perceiver.layers.")) { + perceiver_depth = std::max(perceiver_depth, std::stoi(key.substr(17)) + 1); + } + } + if (depth) { + config.depth = depth; + config.refiner_depth = refiner_depth; + config.perceiver_depth = perceiver_depth; + config.injection_layers.assign(injection_layers.begin(), injection_layers.end()); + LOG_VERBOSE("fibo2: dim = %" PRId64 ", depth = %" PRId64 ", head_dim = %" PRId64 ", text_dim = %" PRId64, + config.dim, config.depth, config.head_dim, config.text_dim); + } + return config; + } + + int gist_count(int text_length) const { + double ratio = std::clamp((text_length - 256.0) / (2048.0 - 256.0), 0.0, 1.0); + return std::min(static_cast(max_gist), static_cast(std::ceil((256 + ratio * (max_gist - 256)) / 64)) * 64); + } + }; + + struct Attention : public GGMLBlock { + int64_t head_dim; + int64_t heads; + bool fused_qkv; + bool fused_kv; + + Attention(const Fibo2Config& config, int64_t kv_dim, bool joint = false) + : head_dim(config.head_dim), heads(config.dim / config.head_dim), fused_qkv(joint && config.fused_qkv), fused_kv(!joint && config.fused_kv) { + if (fused_qkv) { + blocks["qkv"] = std::make_shared(config.dim, 3 * config.dim, false); + } else { + blocks["to_q"] = std::make_shared(config.dim, config.dim, false); + if (fused_kv) { + blocks["to_kv"] = std::make_shared(kv_dim, 2 * config.dim, false); + } else { + blocks["to_k"] = std::make_shared(kv_dim, config.dim, false); + blocks["to_v"] = std::make_shared(kv_dim, config.dim, false); + } + } + blocks["to_out.0"] = std::make_shared(config.dim, config.dim, false); + blocks["norm_q"] = std::make_shared(head_dim, 1e-5f); + blocks["norm_k"] = std::make_shared(head_dim, 1e-5f); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* kv, ggml_tensor* pe = nullptr) { + auto c = ctx->ggml_ctx; + auto project = [&](const char* name, ggml_tensor* h) { + return std::dynamic_pointer_cast(blocks[name])->forward(ctx, h); + }; + ggml_tensor *q, *k, *v; + if (fused_qkv) { + auto parts = ggml_ext_chunk(c, project("qkv", x), 3, 0); + q = parts[0]; + k = parts[1]; + v = parts[2]; + } else { + q = project("to_q", x); + if (fused_kv) { + auto parts = ggml_ext_chunk(c, project("to_kv", kv), 2, 0); + k = parts[0]; + v = parts[1]; + } else { + k = project("to_k", kv); + v = project("to_v", kv); + } + } + auto reshape = [&](ggml_tensor* h) { + return ggml_reshape_4d(c, h, head_dim, heads, h->ne[1], h->ne[2]); + }; + q = std::dynamic_pointer_cast(blocks["norm_q"])->forward(ctx, reshape(q)); + k = std::dynamic_pointer_cast(blocks["norm_k"])->forward(ctx, reshape(k)); + v = reshape(v); + auto prepare = [&](ggml_tensor* h) { + if (pe) { + return Rope::apply_rope(ctx, h, pe); + } + h = ggml_cont(c, ggml_permute(c, h, 0, 2, 1, 3)); + return ggml_reshape_3d(c, h, head_dim, h->ne[1], heads * h->ne[3]); + }; + auto y = ggml_ext_attention_ext(ctx, prepare(q), prepare(k), v, heads, nullptr, true, ctx->flash_attn_enabled); + return std::dynamic_pointer_cast(blocks["to_out.0"])->forward(ctx, y); + } + }; + + struct FeedForward : public ZImage::FeedForward { + bool fused; + + FeedForward(const Fibo2Config& config) + : ZImage::FeedForward(config.dim, config.ffn_dim, 1), fused(config.fused_ffn) { + if (fused) { + blocks.erase("w1"); + blocks.erase("w3"); + blocks["w13"] = std::make_shared(config.dim, 2 * config.ffn_dim, false); + } + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + if (!fused) { + return ZImage::FeedForward::forward(ctx, x); + } + auto w2 = std::dynamic_pointer_cast(blocks["w2"]); + if (sd_backend_is(ctx->backend, "Vulkan") || sd_backend_is(ctx->backend, "ROCm")) { + w2->set_force_prec_f32(true); + } + x = std::dynamic_pointer_cast(blocks["w13"])->forward(ctx, x); + return w2->forward(ctx, ggml_swiglu(ctx->ggml_ctx, x)); + } + }; + + struct PerceiverLayer : public GGMLBlock { + PerceiverLayer(const Fibo2Config& config) { + for (const char* name : {"cross_attn_norm", "self_attn_norm", "ffn_norm"}) { + blocks[name] = std::make_shared(config.dim, 1e-5f); + } + blocks["cross_attention"] = std::make_shared(config, config.text_dim); + blocks["self_attention"] = std::make_shared(config, config.dim); + blocks["ffn"] = std::make_shared(config); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* text) { + auto norm = [&](const char* name, ggml_tensor* h) { + return std::dynamic_pointer_cast(blocks[name])->forward(ctx, h); + }; + auto h = std::dynamic_pointer_cast(blocks["cross_attention"])->forward(ctx, norm("cross_attn_norm", x), text); + x = ggml_add(ctx->ggml_ctx, x, h); + h = norm("self_attn_norm", x); + h = std::dynamic_pointer_cast(blocks["self_attention"])->forward(ctx, h, h); + x = ggml_add(ctx->ggml_ctx, x, h); + h = std::dynamic_pointer_cast(blocks["ffn"])->forward(ctx, norm("ffn_norm", x)); + return ggml_add(ctx->ggml_ctx, x, h); + } + }; + + struct Perceiver : public GGMLBlock { + Fibo2Config config; + + Perceiver(const Fibo2Config& config) + : config(config) { + for (int i = 0; i < config.perceiver_depth; ++i) { + blocks["layers." + std::to_string(i)] = std::make_shared(config); + } + blocks["norm_out"] = std::make_shared(config.dim, 1e-5f); + } + + void init_params(ggml_context* ctx, const String2TensorStorage& weights = {}, const std::string prefix = "") override { + params["latent_queries"] = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, config.dim, config.max_gist); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* text, int gist_count) { + auto x = ggml_ext_slice(ctx->ggml_ctx, params["latent_queries"], 1, 0, gist_count); + for (int i = 0; i < config.perceiver_depth; ++i) { + x = std::dynamic_pointer_cast(blocks["layers." + std::to_string(i)])->forward(ctx, x, text); + } + return std::dynamic_pointer_cast(blocks["norm_out"])->forward(ctx, x); + } + }; + + struct TransformerBlock : public GGMLBlock { + bool injection; + + TransformerBlock(const Fibo2Config& config, bool injection) + : injection(injection) { + blocks["attention"] = std::make_shared(config, config.dim, true); + blocks["feed_forward"] = std::make_shared(config); + for (const char* name : {"attention_norm1", "attention_norm2", "ffn_norm1", "ffn_norm2"}) { + blocks[name] = std::make_shared(config.dim, 1e-5f); + } + blocks["adaLN_modulation.0"] = std::make_shared(std::min(config.dim, 256), 4 * config.dim); + if (injection) { + blocks["cross_attention"] = std::make_shared(config, config.text_dim); + blocks["cross_attn_norm1"] = std::make_shared(config.dim, 1e-5f); + blocks["cross_attn_norm2"] = std::make_shared(config.dim, 1e-5f); + } + } + + void init_params(ggml_context* ctx, const String2TensorStorage& weights = {}, const std::string prefix = "") override { + if (injection) { + params["cross_attn_gate"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 1); + } + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* pe, ggml_tensor* noisy, ggml_tensor* clean, int image_start, int image_count, ggml_tensor* text = nullptr) { + auto c = ctx->ggml_ctx; + auto norm = [&](const char* name, ggml_tensor* h) { + return std::dynamic_pointer_cast(blocks[name])->forward(ctx, h); + }; + auto modulation = std::dynamic_pointer_cast(blocks["adaLN_modulation.0"]); + auto mn = ggml_ext_chunk(c, modulation->forward(ctx, noisy), 4, 0); + auto mc = ggml_ext_chunk(c, modulation->forward(ctx, clean), 4, 0); + auto select = [&](int index, bool gate) { + auto n = gate ? ggml_tanh(c, mn[index]) : mn[index]; + auto k = gate ? ggml_tanh(c, mc[index]) : mc[index]; + auto repeat = [&](ggml_tensor* h, int64_t count) { + return ggml_repeat(c, h, ggml_new_tensor_2d(c, GGML_TYPE_F32, h->ne[0], count)); + }; + auto m = repeat(n, image_count); + if (image_start) { + m = ggml_concat(c, repeat(k, image_start), m, 1); + } + int64_t tail = x->ne[1] - image_start - image_count; + if (tail) { + m = ggml_concat(c, m, repeat(k, tail), 1); + } + return m; + }; + auto h = ZImage::modulate(c, norm("attention_norm1", x), select(0, false)); + h = std::dynamic_pointer_cast(blocks["attention"])->forward(ctx, h, h, pe); + x = ggml_add(c, x, ggml_mul(c, norm("attention_norm2", h), select(1, true))); + if (injection) { + h = std::dynamic_pointer_cast(blocks["cross_attention"])->forward(ctx, norm("cross_attn_norm1", x), text); + h = ggml_mul(c, norm("cross_attn_norm2", h), ggml_tanh(c, params["cross_attn_gate"])); + x = ggml_add(c, x, h); + } + h = ZImage::modulate(c, norm("ffn_norm1", x), select(2, false)); + h = std::dynamic_pointer_cast(blocks["feed_forward"])->forward(ctx, h); + return ggml_add(c, x, ggml_mul(c, norm("ffn_norm2", h), select(3, true))); + } + }; + + struct Fibo2Model : public GGMLBlock { + Fibo2Config config; + + Fibo2Model(const Fibo2Config& config) + : config(config) { + blocks["x_embedder"] = std::make_shared(config.channels, config.dim); + blocks["t_embedder"] = std::make_shared(1024, 256, std::min(config.dim, 256)); + blocks["perceiver"] = std::make_shared(config); + for (int i = 0; i < config.refiner_depth; ++i) { + blocks["noise_refiner." + std::to_string(i)] = std::make_shared(config, false); + } + for (int i = 0; i < config.depth; ++i) { + bool injection = std::find(config.injection_layers.begin(), config.injection_layers.end(), i) != config.injection_layers.end(); + blocks["layers." + std::to_string(i)] = std::make_shared(config, injection); + } + blocks["final_layer"] = std::make_shared(config.dim, 1, config.channels); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* t, ggml_tensor* text, ggml_tensor* pe, int gist_count, const std::vector& refs) { + auto c = ctx->ggml_ctx; + int64_t width = x->ne[0], height = x->ne[1]; + int image_count = static_cast(width * height); + // The shared FLUX.2 VAE already packs (channel, dy, dx) into 128 channels. + auto tokens = [&](ggml_tensor* h) { + h = ggml_reshape_3d(c, h, h->ne[0] * h->ne[1], config.channels, 1); + return ggml_cont(c, ggml_permute(c, h, 1, 0, 2, 3)); + }; + x = tokens(x); + for (auto ref : refs) { + x = ggml_concat(c, x, tokens(ref), 1); + } + x = std::dynamic_pointer_cast(blocks["x_embedder"])->forward(ctx, x); + auto embedder = std::dynamic_pointer_cast(blocks["t_embedder"]); + auto noisy = embedder->forward(ctx, t); + auto clean = embedder->forward(ctx, ggml_scale(c, t, 0.f)); + auto image_pe = ggml_ext_slice(c, pe, 3, gist_count, pe->ne[3]); + for (int i = 0; i < config.refiner_depth; ++i) { + x = std::dynamic_pointer_cast(blocks["noise_refiner." + std::to_string(i)])->forward(ctx, x, image_pe, noisy, clean, 0, image_count); + sd::ggml_graph_cut::mark_graph_cut(x, "fibo2.noise_refiner." + std::to_string(i), "x"); + } + auto bundle = [&](int i) { return ggml_ext_slice(c, text, 0, i * config.text_dim, (i + 1) * config.text_dim); }; + auto gist = std::dynamic_pointer_cast(blocks["perceiver"])->forward(ctx, bundle(0), gist_count); + x = ggml_concat(c, gist, x, 1); + int bundle_index = 1; + for (int i = 0; i < config.depth; ++i) { + auto block = std::dynamic_pointer_cast(blocks["layers." + std::to_string(i)]); + auto context = block->injection ? bundle(bundle_index++) : nullptr; + x = block->forward(ctx, x, pe, noisy, clean, gist_count, image_count, context); + sd::ggml_graph_cut::mark_graph_cut(x, "fibo2.layers." + std::to_string(i), "x"); + } + x = ggml_ext_slice(c, x, 1, gist_count, gist_count + image_count); + x = std::dynamic_pointer_cast(blocks["final_layer"])->forward(ctx, x, noisy); + x = ggml_cont(c, ggml_permute(c, x, 1, 0, 2, 3)); + return ggml_reshape_4d(c, x, width, height, config.channels, 1); + } + }; + + struct Fibo2Runner : public DiffusionModelRunner { + Fibo2Config config; + Fibo2Model model; + std::vector pe_values; + + Fibo2Runner(ggml_backend_t backend, const String2TensorStorage& weights, const std::string& prefix, std::shared_ptr weight_manager = nullptr) + : DiffusionModelRunner(backend, prefix, weight_manager), + config(Fibo2Config::detect_from_weights(weights, prefix)), + model(config) { + model.init(params_ctx, weights, prefix); + } + + std::string get_desc() override { return "fibo2"; } + + void get_param_tensors(std::map& tensors, const std::string& prefix) override { + model.get_param_tensors(tensors, prefix); + } + + sd::Tensor compute(int n_threads, const DiffusionParams& p) override { + if (!p.x || !p.timesteps || !p.context || p.context->shape()[0] != config.text_dim * 6) { + LOG_ERROR("Fibo 2.0 requires six text feature bundles"); + return {}; + } + auto build = [&]() { + auto graph = new_graph_custom(65536); + auto x = make_input(*p.x); + GGML_ASSERT(x->ne[3] == 1 && x->ne[2] == config.channels); + int count = config.gist_count(static_cast(p.context->shape()[1])); + Rope::Embedding embedding; + for (int i = 0; i < count; ++i) { + embedding.ids.push_back({static_cast(i + 1), 0.f, 0.f}); + } + embedding.positions.append_tokens(count); + auto append_image = [&](ggml_tensor* image, int plane) { + auto ids = Rope::gen_flux_img_ids(static_cast(image->ne[1]), static_cast(image->ne[0]), + 1, 1, 3, plane, 0, 0, false, &embedding.positions); + embedding.ids.insert(embedding.ids.end(), ids.begin(), ids.end()); + }; + append_image(x, static_cast(config.max_gist + 1)); + std::vector refs; + if (p.ref_latents && p.ref_image_params.pass_to_dit) { + for (const auto& ref : *p.ref_latents) { + refs.push_back(make_input(ref)); + append_image(refs.back(), static_cast(config.max_gist + 1 + refs.size())); + } + } + embedding.values = Rope::embed_nd(embedding.ids, 1, 256.f, {32, 48, 48}, embedding.layout, &embedding.frequencies); + pe_values = finish_rope_pe(std::move(embedding)); + auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, 64, static_cast(pe_values.size() / 256)); + set_backend_tensor_data(pe, pe_values.data()); + auto ctx = get_context(); + auto out = model.forward(&ctx, x, make_input(*p.timesteps), make_input(*p.context), pe, count, refs); + ggml_build_forward_expand(graph, out); + return graph; + }; + return restore_trailing_singleton_dims(GGMLRunner::compute(build, n_threads, false), p.x->dim()); + } + }; +} + +#endif // __SD_MODEL_DIFFUSION_FIBO2_HPP__ diff --git a/src/model/diffusion/hidream_o1.hpp b/src/model/diffusion/hidream_o1.hpp index 2befb723a3..ca1559bcb7 100644 --- a/src/model/diffusion/hidream_o1.hpp +++ b/src/model/diffusion/hidream_o1.hpp @@ -56,15 +56,6 @@ namespace HiDreamO1 { } }; - static inline std::string repeat_special_token(const std::string& token, int64_t count) { - std::string out; - out.reserve(static_cast(count) * token.size()); - for (int64_t i = 0; i < count; ++i) { - out += token; - } - return out; - } - static inline std::pair calculate_dimensions(int max_size, double ratio) { int width = static_cast(std::sqrt(max_size * max_size * ratio)); int height = static_cast(width / ratio); @@ -405,7 +396,7 @@ namespace HiDreamO1 { txt = LLM::splice_image_embeds(&runner_ctx, txt, image_embeds); auto t_emb = t_embedder1->forward(&runner_ctx, timestep); - int64_t txt_seq_len = input_ids->ne[0]; + int64_t txt_seq_len = txt->ne[1]; if (txt_seq_len > 1) { auto prefix = ggml_ext_slice(compute_ctx, txt, 1, 0, txt_seq_len - 1); txt = ggml_concat(compute_ctx, prefix, ggml_reshape_3d(compute_ctx, t_emb, t_emb->ne[0], 1, 1), 1); @@ -616,8 +607,7 @@ namespace HiDreamO1 { sd::Tensor vlm_image = sd::ops::interpolate( resized_ref, {dims.first, dims.second, resized_ref.shape()[2], resized_ref.shape()[3]}); - vlm_image = vlm_image * 2.0f - 1.0f; - int64_t image_tokens = static_cast(dims.first / PATCH_SIZE) * static_cast(dims.second / PATCH_SIZE); + vlm_image = vlm_image * 2.0f - 1.0f; auto patch_img = resized_ref * 2.0f - 1.0f; result.c_ref_images.push_back(std::move(patch_img)); @@ -626,8 +616,11 @@ namespace HiDreamO1 { return {}; } int64_t prompt_start = static_cast(prefix_tokens.size()); + for (const auto& grid : image_grids) { + prompt_start += static_cast(grid[1]) * grid[2] - 1; + } prompt += "<|vision_start|>"; - prompt += repeat_special_token("<|image_pad|>", image_tokens); + prompt += "<|image_pad|>"; prompt += "<|vision_end|>"; vlm_images.emplace_back(static_cast(prompt_start), std::move(vlm_image)); image_grids.push_back({1, dims.second / PATCH_SIZE, dims.first / PATCH_SIZE}); @@ -642,6 +635,13 @@ namespace HiDreamO1 { } std::vector input_ids_pad = input_ids; + for (size_t i = 0; i < vlm_images.size(); ++i) { + const auto& grid = image_grids[i]; + const int64_t count = static_cast(grid[1]) * grid[2]; + input_ids_pad.insert(input_ids_pad.begin() + vlm_images[i].first, count - 1, IMAGE_TOKEN_ID); + } + + int txt_seq_len = static_cast(input_ids_pad.size()); input_ids_pad.push_back(VISION_START_TOKEN_ID); input_ids_pad.insert(input_ids_pad.end(), target_image_len - 1, IMAGE_TOKEN_ID); image_grids.push_back({1, static_cast(height / PATCH_SIZE), static_cast(width / PATCH_SIZE)}); @@ -656,8 +656,7 @@ namespace HiDreamO1 { } std::vector token_types(input_ids_pad.size(), 0); - int txt_seq_len = static_cast(input_ids.size()); - int bgn = txt_seq_len - TIMESTEP_TOKEN_NUM; + int bgn = txt_seq_len - TIMESTEP_TOKEN_NUM; for (int i = bgn; i < static_cast(token_types.size()); ++i) { token_types[i] = 1; } diff --git a/src/model/diffusion/model.hpp b/src/model/diffusion/model.hpp index 955c5277a7..7c14d9d1fd 100644 --- a/src/model/diffusion/model.hpp +++ b/src/model/diffusion/model.hpp @@ -27,9 +27,25 @@ struct RefImageParams { int vlm_min_size = -1; int vlm_max_size = -1; bool resize_vae_to_target = false; + bool resize_single_ref_to_target = false; + bool vae_size_buckets = false; + bool vae_center_crop = false; + bool vlm_use_vae_input = false; }; const std::unordered_map REF_IMAGE_PRESETS = { + {"fibo2", [] { + RefImageParams params; + params.pass_to_vlm = true; + params.ref_index_mode = Rope::RefIndexMode::INCREASE; + params.vlm_resize_mode = RefImageResizeMode::LONGEST_SIDE; + params.vlm_max_size = 768; + params.resize_single_ref_to_target = true; + params.vae_size_buckets = true; + params.vae_center_crop = true; + params.vlm_use_vae_input = true; + return params; + }()}, {"flux_kontext", {false, true, Rope::RefIndexMode::FIXED, false, true, -1, RefImageResizeMode::NONE, -1, -1}}, {"longcat", {true, true, Rope::RefIndexMode::FIXED, false, true, -1, RefImageResizeMode::AREA, -1, -1}}, {"flux2", {false, true, Rope::RefIndexMode::INCREASE, false, true, -1, RefImageResizeMode::NONE, -1, -1}}, diff --git a/src/model/te/llm.hpp b/src/model/te/llm.hpp index 8ea918025e..85f7511693 100644 --- a/src/model/te/llm.hpp +++ b/src/model/te/llm.hpp @@ -924,6 +924,8 @@ namespace LLM { int64_t txt_token_start = 0; int64_t txt_token_end = 0; ggml_tensor* input_embed = nullptr; + // Image and DeepStack indices remain in the expanded sequence coordinates. + int64_t removed_tokens = 0; for (int i = 0; i < image_embeds.size(); i++) { if (i == 0) { @@ -931,7 +933,10 @@ namespace LLM { } else { txt_token_start = image_embeds[i - 1].first + image_embeds[i - 1].second->ne[1]; } - txt_token_end = image_embeds[i].first; + txt_token_start -= removed_tokens; + txt_token_end = image_embeds[i].first - removed_tokens; + removed_tokens += image_embeds[i].second->ne[1] - 1; + GGML_ASSERT(txt_token_start >= 0 && txt_token_end >= txt_token_start && txt_token_end < raw_x->ne[1]); // An embed can sit flush against the previous one or at the very start/end of the // sequence, leaving no text tokens to splice around it. @@ -951,14 +956,14 @@ namespace LLM { } } - txt_token_start = image_embeds[image_embeds.size() - 1].first + image_embeds[image_embeds.size() - 1].second->ne[1]; + txt_token_start = image_embeds[image_embeds.size() - 1].first + image_embeds[image_embeds.size() - 1].second->ne[1] - removed_tokens; txt_token_end = raw_x->ne[1]; if (txt_token_end > txt_token_start) { auto final_txt_embed = ggml_ext_slice(ctx->ggml_ctx, raw_x, 1, txt_token_start, txt_token_end); input_embed = ggml_concat(ctx->ggml_ctx, input_embed, final_txt_embed, 1); } - GGML_ASSERT(raw_x->ne[1] == input_embed->ne[1]); + GGML_ASSERT(raw_x->ne[1] + removed_tokens == input_embed->ne[1]); return input_embed; } @@ -2506,6 +2511,9 @@ namespace LLM { } int64_t n_tokens = input_ids->ne[0]; + for (const auto& image_embed : image_embeds) { + n_tokens += image_embed.second->ne[1] - 1; + } if (config.arch == LLMArch::MISTRAL_SMALL_3_2 || config.arch == LLMArch::MINISTRAL_3_3B || config.arch == LLMArch::QWEN3 || @@ -3014,20 +3022,17 @@ namespace LLM { LOG_VERBOSE("llm encode_image test done in %lldms", t1 - t0); } - std::string placeholder = "<|image_pad|>"; - std::string img_prompt = "Picture 1: <|vision_start|>"; // [24669, 220, 16, 25, 220, 151652] - int64_t num_image_tokens = image_embed.shape()[1]; - img_prompt.reserve(num_image_tokens * placeholder.size()); - for (int i = 0; i < num_image_tokens; i++) { - img_prompt += placeholder; - } - img_prompt += "<|vision_end|>"; - - std::vector>> image_embeds; - image_embeds.emplace_back(64, image_embed); + const std::string image_prefix = "Picture 1: <|vision_start|>"; + const std::string img_prompt = image_prefix + "<|image_pad|><|vision_end|>"; std::pair prompt_attn_range; std::string text = "<|im_start|>system\nDescribe the key features of the input image (color, shape, size, texture, objects, background), then explain how the user's text instruction should alter or modify the image. Generate a new image that meets the user's requirements while maintaining consistency with the original input where appropriate.<|im_end|>\n<|im_start|>user\n"; + std::vector prefix_tokens; + if (!tokenizer->encode(text + image_prefix, prefix_tokens, nullptr)) { + return; + } + std::vector>> image_embeds; + image_embeds.emplace_back(static_cast(prefix_tokens.size()), image_embed); text += img_prompt; prompt_attn_range.first = static_cast(text.size()); text += "change 'flux.cpp' to 'edit.cpp'"; diff --git a/src/model_loader.cpp b/src/model_loader.cpp index 90ab65eed5..03bcd74a94 100644 --- a/src/model_loader.cpp +++ b/src/model_loader.cpp @@ -439,6 +439,10 @@ bool ModelLoader::init_from_diffusers_file(const std::string& file_path, const s SDVersion ModelLoader::get_sd_version() const { TensorStorage token_embedding_weight, input_block_weight, context_ebedding_weight; + if (tensor_storage_map.count("model.diffusion_model.perceiver.latent_queries")) { + return VERSION_FIBO2; + } + if (tensor_storage_map.find("model.diffusion_model.emb_in.image.weight") != tensor_storage_map.end() && tensor_storage_map.find("model.diffusion_model.txt_mode_blocks.0.q_proj.weight") != tensor_storage_map.end()) { return VERSION_FLUX3; diff --git a/src/model_manager.cpp b/src/model_manager.cpp index 05ee55f676..4a3103aa18 100644 --- a/src/model_manager.cpp +++ b/src/model_manager.cpp @@ -2004,12 +2004,12 @@ bool ModelManager::ensure_compute_backend_capacity( }; size_t last_eviction_count = 0; auto release_eviction_states = [&]() { - if (eviction_states.size() == last_eviction_count) { - return false; + if (eviction_states.size() != last_eviction_count) { + last_eviction_count = eviction_states.size(); + release_compute_staging_blocks(false, &eviction_states); + release_params_storage_blocks(false, &eviction_states); } - last_eviction_count = eviction_states.size(); - release_compute_staging_blocks(false, &eviction_states); - release_params_storage_blocks(false, &eviction_states); + // Device memory reports can lag behind releases even when no new candidates are added. return fits(); }; diff --git a/src/name_conversion.cpp b/src/name_conversion.cpp index 6a519ba21f..f37c4ddbe3 100644 --- a/src/name_conversion.cpp +++ b/src/name_conversion.cpp @@ -197,7 +197,11 @@ std::string convert_cond_stage_model_name(std::string name, std::string prefix, replace_with_name_map(name, llm_vision_name_map); } else { replace_with_name_map(name, llm_name_map); - name = convert_llada2_moe_te_name(name); + name = convert_llada2_moe_te_name(name); + const std::string language_prefix = "llm.model.language_model."; + if (starts_with(name, language_prefix)) { + name = "llm.model." + name.substr(language_prefix.size()); + } } } else { name = convert_open_clip_to_hf_clip_name(name); @@ -1034,6 +1038,18 @@ static std::string convert_ming_image_dit_name(std::string name) { return name; } +static std::string convert_fibo2_dit_name(std::string name) { + static const std::vector> name_map = { + {".attention.q_norm.", ".attention.norm_q."}, + {".attention.k_norm.", ".attention.norm_k."}, + {".attention.out.", ".attention.to_out.0."}, + {".cross_attention.to_out.weight", ".cross_attention.to_out.0.weight"}, + {".self_attention.to_out.weight", ".self_attention.to_out.0.weight"}, + }; + replace_with_name_map(name, name_map); + return name; +} + std::string convert_diffusers_dit_to_original_minimax_h3(std::string name) { static const std::vector> prefix_map = { {"transformer_blocks.", "blocks."}, @@ -1067,6 +1083,8 @@ std::string convert_diffusion_model_name(std::string name, std::string prefix, S name = convert_moss_image_dit_name(name); } else if (version == VERSION_MING_IMAGE) { name = convert_ming_image_dit_name(name); + } else if (version == VERSION_FIBO2) { + name = convert_fibo2_dit_name(name); } else if (sd_version_is_z_image(version) || sd_version_is_z_image_l2p(version)) { name = convert_diffusers_dit_to_original_lumina2(name); } else if (sd_version_is_llada_image(version)) { @@ -1715,7 +1733,7 @@ std::string convert_tensor_name(std::string name, SDVersion version) { } } - if (version == VERSION_QWEN_IMAGE_2_1 || sd_version_is_boogu_image(version) || sd_version_is_krea2(version) || sd_version_is_mage_flow(version) || sd_version_is_minimax_h3(version)) { + if (version == VERSION_FIBO2 || version == VERSION_QWEN_IMAGE_2_1 || sd_version_is_boogu_image(version) || sd_version_is_krea2(version) || sd_version_is_mage_flow(version) || sd_version_is_minimax_h3(version)) { const std::string hf_vision_prefix = "text_encoders.llm.model.visual."; if (starts_with(name, hf_vision_prefix)) { name = "text_encoders.llm.visual." + name.substr(hf_vision_prefix.size()); diff --git a/src/pipeline/diffusion_engine.cpp b/src/pipeline/diffusion_engine.cpp index 82b40473a5..ed200d45e6 100644 --- a/src/pipeline/diffusion_engine.cpp +++ b/src/pipeline/diffusion_engine.cpp @@ -114,6 +114,7 @@ const char* model_version_to_str[] = { "Iris-3B", "Flux.3 Image Klein", "MOSS-Image", + "Fibo-2.0", }; static_assert(VERSION_COUNT == sizeof(model_version_to_str) / sizeof(model_version_to_str[0]), @@ -1520,7 +1521,8 @@ bool StableDiffusionGGML::build_denoiser() { sd_version_is_pid(version) || sd_version_is_ideogram4(version) || version == VERSION_IRIS || - version == VERSION_MOSS_IMAGE) { + version == VERSION_MOSS_IMAGE || + version == VERSION_FIBO2) { pred_type = FLOW_PRED; if (sd_version_is_wan(version) || version == VERSION_MOSS_IMAGE) { default_flow_shift = 5.f; @@ -3359,7 +3361,9 @@ bool StableDiffusionGGML::is_flow_denoiser() { } std::string StableDiffusionGGML::get_default_ref_image_preset(SDVersion version) const { - if (sd_version_is_longcat(version)) { + if (version == VERSION_FIBO2) { + return "fibo2"; + } else if (sd_version_is_longcat(version)) { return "longcat"; } else if (sd_version_is_flux(version)) { return "flux_kontext"; @@ -3373,7 +3377,7 @@ std::string StableDiffusionGGML::get_default_ref_image_preset(SDVersion version) return "qwen"; } else if (sd_version_is_mage_flow(version)) { return "mage_flow"; - } else if (sd_version_is_z_image(version) || sd_version_is_boogu_image(version)) { + } else if (sd_version_is_boogu_image(version)) { return "z_image_omni"; } else if (sd_version_is_llada_image(version)) { return "llada_image"; @@ -3411,6 +3415,9 @@ RefImageParams StableDiffusionGGML::resolve_ref_image_params(const char* ref_ima LOG_INFO("Using '%s' preset for reference images", preset_name.c_str()); params = REF_IMAGE_PRESETS.at(preset_name); } + if (version == VERSION_QWEN_IMAGE_2_1) { + params.vlm_use_vae_input = true; + } for (const auto& [key, value] : parse_key_value_args(ref_image_args, "reference image args")) { if (key == "pass_to_vlm") { @@ -3439,6 +3446,26 @@ RefImageParams StableDiffusionGGML::resolve_ref_image_params(const char* ref_ima if (!parse_strict_bool(value, params.resize_before_vae)) { LOG_WARN("ignoring invalid reference image arg '%s=%s'", key.c_str(), value.c_str()); } + } else if (key == "resize_vae_to_target") { + if (!parse_strict_bool(value, params.resize_vae_to_target)) { + LOG_WARN("ignoring invalid reference image arg '%s=%s'", key.c_str(), value.c_str()); + } + } else if (key == "resize_single_ref_to_target") { + if (!parse_strict_bool(value, params.resize_single_ref_to_target)) { + LOG_WARN("ignoring invalid reference image arg '%s=%s'", key.c_str(), value.c_str()); + } + } else if (key == "vae_size_buckets") { + if (!parse_strict_bool(value, params.vae_size_buckets)) { + LOG_WARN("ignoring invalid reference image arg '%s=%s'", key.c_str(), value.c_str()); + } + } else if (key == "vae_center_crop") { + if (!parse_strict_bool(value, params.vae_center_crop)) { + LOG_WARN("ignoring invalid reference image arg '%s=%s'", key.c_str(), value.c_str()); + } + } else if (key == "vlm_use_vae_input") { + if (!parse_strict_bool(value, params.vlm_use_vae_input)) { + LOG_WARN("ignoring invalid reference image arg '%s=%s'", key.c_str(), value.c_str()); + } } else if (key == "vae_input_max_pixels") { if (!parse_strict_int(value, params.vae_input_max_pixels)) { LOG_WARN("ignoring invalid reference image arg '%s=%s'", key.c_str(), value.c_str()); diff --git a/src/pipeline/image.cpp b/src/pipeline/image.cpp index 7de85b3595..d20c02e6b8 100644 --- a/src/pipeline/image.cpp +++ b/src/pipeline/image.cpp @@ -286,26 +286,34 @@ namespace sd::pipeline { LOG_VERBOSE("auto resize ref images"); double vae_width; double vae_height; - if (ref_image_params.resize_vae_to_target) { + int factor = sd_version_is_qwen_image(sd->version) ? 32 : 16; + if (ref_image_params.resize_vae_to_target || (ref_image_params.resize_single_ref_to_target && ref_images.size() == 1)) { vae_width = request->width; vae_height = request->height; } else { int default_pixels = sd->version == VERSION_QWEN_IMAGE_2_1 ? request->width * request->height : 1024 * 1024; int target_pixels = ref_image_params.vae_input_max_pixels > 0 ? ref_image_params.vae_input_max_pixels : default_pixels; - int vae_image_size = std::min(target_pixels, request->width * request->height); - vae_width = sqrt(vae_image_size * ref_images[i].shape()[0] / ref_images[i].shape()[1]); - vae_height = vae_width * ref_images[i].shape()[1] / ref_images[i].shape()[0]; + if (ref_image_params.vae_size_buckets) { + auto size = sd::select_image_size_bucket(static_cast(ref_images[i].shape()[0]), static_cast(ref_images[i].shape()[1]), target_pixels, factor, request->width, request->height); + vae_width = size.first; + vae_height = size.second; + } else { + int vae_image_size = std::min(target_pixels, request->width * request->height); + vae_width = sqrt(vae_image_size * ref_images[i].shape()[0] / ref_images[i].shape()[1]); + vae_height = vae_width * ref_images[i].shape()[1] / ref_images[i].shape()[0]; + } } - int factor = sd_version_is_qwen_image(sd->version) ? 32 : 16; vae_height = round(vae_height / factor) * factor; vae_width = round(vae_width / factor) * factor; - auto resized_ref_img = sd::ops::interpolate(ref_images[i], - {static_cast(vae_width), - static_cast(vae_height), - ref_images[i].shape()[2], - ref_images[i].shape()[3]}); + auto resized_ref_img = ref_image_params.vae_center_crop + ? sd::resize_and_center_crop(ref_images[i], static_cast(vae_width), static_cast(vae_height)) + : sd::ops::interpolate(ref_images[i], + {static_cast(vae_width), + static_cast(vae_height), + ref_images[i].shape()[2], + ref_images[i].shape()[3]}); LOG_VERBOSE("resize vae ref image %d from %" PRId64 "x%" PRId64 " to %" PRId64 "x%" PRId64, static_cast(i), @@ -315,7 +323,7 @@ namespace sd::pipeline { resized_ref_img.shape()[0]); ref_latent = sd->encode_first_stage(resized_ref_img); - if (sd->version == VERSION_QWEN_IMAGE_2_1) { + if (ref_image_params.vlm_use_vae_input) { ref_images[i] = std::move(resized_ref_img); } } else { @@ -484,8 +492,8 @@ namespace sd::pipeline { } condition_params.text = request->negative_prompt; condition_params.zero_out_masked = zero_out_masked; - if (sd_version_is_llada_image(sd->version)) { - // LLaDA-Image CFG keeps the source latent but drops its SigVQ features. + if (sd_version_is_llada_image(sd->version) || sd->version == VERSION_FIBO2) { + // Keep source latents in CFG, but omit source images from the text encoder. condition_params.ref_images = nullptr; } uncond = sd->get_learned_condition(condition_params); diff --git a/src/pipeline/model_builders.cpp b/src/pipeline/model_builders.cpp index 4b435170f7..ca4d1998b1 100644 --- a/src/pipeline/model_builders.cpp +++ b/src/pipeline/model_builders.cpp @@ -5,6 +5,7 @@ #include #include "conditioning/conditioner.hpp" +#include "conditioning/fibo2.hpp" #include "core/ggml_extend_backend.h" #include "core/util.h" #include "extensions/generation_extension.h" @@ -14,6 +15,7 @@ #include "model/diffusion/boogu.hpp" #include "model/diffusion/control.hpp" #include "model/diffusion/ernie_image.hpp" +#include "model/diffusion/fibo2.hpp" #include "model/diffusion/flux.hpp" #include "model/diffusion/flux3.hpp" #include "model/diffusion/hidream_o1.hpp" @@ -94,7 +96,12 @@ namespace sd::model_builders { return false; } - if (sd_version_is_sd3(version)) { + if (version == VERSION_FIBO2) { + result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), + tensor_storage_map, weight_manager, tokenizers); + result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), + tensor_storage_map, "model.diffusion_model", weight_manager); + } else if (sd_version_is_sd3(version)) { result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), tensor_storage_map, weight_manager, diff --git a/src/pipeline/request.cpp b/src/pipeline/request.cpp index 4256fee4e2..830a5fb990 100644 --- a/src/pipeline/request.cpp +++ b/src/pipeline/request.cpp @@ -65,7 +65,7 @@ namespace sd::pipeline { return LCM_SCHEDULER; } else if (sample_method == DDIM_TRAILING_SAMPLE_METHOD) { return SIMPLE_SCHEDULER; - } else if (sd != nullptr && (sd_version_is_flux(sd->version) || sd->version == VERSION_QWEN_IMAGE_2_1)) { + } else if (sd != nullptr && (sd_version_is_flux(sd->version) || sd->version == VERSION_QWEN_IMAGE_2_1 || sd->version == VERSION_FIBO2)) { return FLUX_SCHEDULER; } else if (sd != nullptr && sd_version_is_flux2(sd->version)) { return FLUX2_SCHEDULER; @@ -329,6 +329,11 @@ namespace sd::pipeline { } void GenerationRequest::resolve(StableDiffusionGGML* sd) { + if (sd_version_is_minit2i(sd->version) && (width != 512 || height != 512)) { + width = 512; + height = 512; + LOG_WARN("align to fixed size 512x512 for MiniT2I"); + } align_generation_request_size(); resolve_hires(); seed = resolve_seed(seed); diff --git a/src/runtime/image_preprocess.cpp b/src/runtime/image_preprocess.cpp index a8df7c3b56..ec649390be 100644 --- a/src/runtime/image_preprocess.cpp +++ b/src/runtime/image_preprocess.cpp @@ -6,6 +6,42 @@ namespace sd { + std::pair select_image_size_bucket(int source_width, int source_height, int area, int factor, int fallback_width, int fallback_height) { + int width = fallback_width; + int height = fallback_height; + double ratio = static_cast(source_width) / source_height; + double best_crop = INFINITY; + int64_t best_area_error = INT64_MAX; + int min_size = 4 * factor; + for (int w = min_size; w < 2LL * area / min_size; w += factor) { + for (int h = min_size; h < 2LL * area / w; h += factor) { + int64_t pixels = int64_t(w) * h; + if (pixels < 0.9 * area || pixels > 1.1 * area) { + continue; + } + double candidate_ratio = static_cast(w) / h; + double crop = std::round((1.0 - std::min(candidate_ratio / ratio, ratio / candidate_ratio)) * 1e9); + int64_t area_error = std::abs(pixels - area); + if (crop < best_crop || (crop == best_crop && (area_error < best_area_error || (area_error == best_area_error && w > width)))) { + width = w; + height = h; + best_crop = crop; + best_area_error = area_error; + } + } + } + return {width, height}; + } + + Tensor resize_and_center_crop(const Tensor& image, int width, int height, ops::InterpolateMode filter) { + double scale = std::max(static_cast(width) / image.shape()[0], static_cast(height) / image.shape()[1]); + int w = static_cast(std::ceil(image.shape()[0] * scale)); + int h = static_cast(std::ceil(image.shape()[1] * scale)); + auto resized = ops::interpolate(image, {w, h, image.shape()[2], image.shape()[3]}, filter); + resized = ops::slice(resized, 0, (w - width) / 2, (w - width) / 2 + width); + return ops::slice(resized, 1, (h - height) / 2, (h - height) / 2 + height); + } + static constexpr std::pair image_targets[] = { {"init", ImageTarget::Init}, {"end", ImageTarget::End}, diff --git a/src/runtime/image_preprocess.h b/src/runtime/image_preprocess.h index b75c4e0e50..7dfc458992 100644 --- a/src/runtime/image_preprocess.h +++ b/src/runtime/image_preprocess.h @@ -3,6 +3,7 @@ #include #include +#include #include #include "core/tensor.hpp" @@ -10,6 +11,9 @@ namespace sd { + std::pair select_image_size_bucket(int source_width, int source_height, int area, int factor, int fallback_width, int fallback_height); + Tensor resize_and_center_crop(const Tensor& image, int width, int height, ops::InterpolateMode filter = ops::InterpolateMode::Lanczos); + enum class ImageTarget { Init, End,