commit f4e68ed970c5f3b6ecfc0ff7cd0be6d19695f9c0 Author: wehub-resource-sync Date: Mon Jul 13 12:34:46 2026 +0800 chore: import upstream snapshot with attribution diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..a12c095 --- /dev/null +++ b/.gitignore @@ -0,0 +1,27 @@ +# Build +build/ +dist/ +*.egg-info/ +*.bak +*.tmp +# Python +__pycache__/ +*.pyc +*.pyo +*.a +deprecated + +# C++ / CMake +CMakeCache.txt +CMakeFiles/ +cmake_install.cmake +Makefile +*.o +*.so +*.dylib + +# IDE +.idea/ +.vscode/ +*.swp +.DS_Store diff --git a/CMakeLists.txt b/CMakeLists.txt new file mode 100644 index 0000000..efb0c95 --- /dev/null +++ b/CMakeLists.txt @@ -0,0 +1,67 @@ +cmake_minimum_required(VERSION 3.27) + +project(cider LANGUAGES CXX OBJCXX) + +set(CMAKE_CXX_STANDARD 17) +set(CMAKE_CXX_STANDARD_REQUIRED ON) +set(CMAKE_POSITION_INDEPENDENT_CODE ON) + +option(BUILD_SHARED_LIBS "Build extensions as a shared library" ON) + +# ── Python + nanobind + MLX ────────────────────────────────────── +find_package( + Python 3.12 + COMPONENTS Interpreter Development.Module + REQUIRED) + +execute_process( + COMMAND "${Python_EXECUTABLE}" -m nanobind --cmake_dir + OUTPUT_STRIP_TRAILING_WHITESPACE + OUTPUT_VARIABLE nanobind_ROOT) +find_package(nanobind CONFIG REQUIRED) + +execute_process( + COMMAND "${Python_EXECUTABLE}" -m mlx --cmake-dir + OUTPUT_STRIP_TRAILING_WHITESPACE + OUTPUT_VARIABLE MLX_ROOT) +find_package(MLX CONFIG REQUIRED) + +# ── Shared library: all primitives ─────────────────────────────── +add_library(cider_prim_lib) + +target_sources(cider_prim_lib PUBLIC + ${CMAKE_CURRENT_LIST_DIR}/csrc/src/w8a8_primitive.mm + ${CMAKE_CURRENT_LIST_DIR}/csrc/src/w4a8_primitive.mm + ${CMAKE_CURRENT_LIST_DIR}/csrc/src/pergroup_primitive.mm + ${CMAKE_CURRENT_LIST_DIR}/csrc/src/sdpa_primitive.mm +) + +target_include_directories(cider_prim_lib PUBLIC + ${CMAKE_CURRENT_LIST_DIR}/csrc/include +) + +target_link_libraries(cider_prim_lib PUBLIC + mlx + "-framework Metal" + "-framework Foundation" +) + +# ── nanobind Python module ─────────────────────────────────────── +nanobind_add_module( + _cider_prim + NB_STATIC + STABLE_ABI + LTO + NOMINSIZE + NB_DOMAIN mlx + ${CMAKE_CURRENT_LIST_DIR}/csrc/src/prim_bindings.cpp +) +target_link_libraries(_cider_prim PRIVATE cider_prim_lib) + +if(BUILD_SHARED_LIBS) + target_link_options(_cider_prim PRIVATE -Wl,-rpath,@loader_path) +endif() + +# ── Install ────────────────────────────────────────────────────── +install(TARGETS cider_prim_lib LIBRARY DESTINATION cider/lib) +install(TARGETS _cider_prim LIBRARY DESTINATION cider/lib) diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..bbbe85d --- /dev/null +++ b/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2026 Mininglamp contributors + +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. diff --git a/README.md b/README.md new file mode 100644 index 0000000..9211e35 --- /dev/null +++ b/README.md @@ -0,0 +1,605 @@ +# cider + +Cider is developed on top of MLX for macOS. It provides: + +1. **Online activation quantization** (W8A8/W4A8) absent in MLX, with custom INT8 TensorOps kernels built as MLX custom primitives supporting full lazy evaluation. +2. **Optimized SDPA decode kernel** inspired by [FlashInfer](https://github.com/flashinfer-ai/flashinfer), featuring contiguous chunk GQA scheduling and register tiling for faster autoregressive decode. + +It also includes service-side extensions and non-intrusive compatibility patches for mlx_vlm (validated on mlx_vlm 0.4.3), including fixes for Qwen3-VL multi-image inference issues related to RoPE position handling and chunked prefill. + +## Conditional Compilation (M4 / M5) + +Cider uses **conditional compilation**: the INT8 TensorOps C++ extension is only built on Apple M5+. + +| Chip | `pip install -e .` behavior | `import cider` behavior | +|------|----------------------------|------------------------| +| **M5+** | Full build (CMake + Metal kernels) | All features available | +| **M4 and below** | Skips C++ build, installs pure-Python package | `is_available()` → False, `convert_model()` is a warning no-op | + +**Override via environment variable:** +```bash +CIDER_FORCE_BUILD=1 pip install -e . # Force build (e.g., CI) +CIDER_FORCE_BUILD=0 pip install -e . # Force skip +``` + +## Modes + +| Mode | Weights | Activations | Compute Path | Status | +|------|---------|-------------|--------------|--------| +| **W8A8** | INT8 symmetric | INT8 per-token | TensorOps matmul2d | ✅ Implemented | +| **W4A8** | INT4 packed (uint8) | INT8 per-token | Unpack→TensorOps | ✅ Implemented | +| W4A16 | — | — | MLX built-in | Baseline | +| W8A16 | — | — | MLX built-in | Baseline | + +**W4A16 and W8A16 are already supported by MLX natively** — this SDK provides the missing W8A8 and W4A8 modes that MLX does not implement. + +## Optimized SDPA (Decode Attention) + +Cider includes an optimized **Scaled Dot-Product Attention** kernel for autoregressive decode (`Q_seq=1`), inspired by [FlashInfer](https://github.com/flashinfer-ai/flashinfer). Key techniques: + +- **Contiguous Chunk GQA Scheduling**: Instead of dispatching each Q head independently (reading KV `H_q` times), a single threadgroup processes all Q heads sharing the same KV group. This reduces DRAM bandwidth by up to `gqa_factor`×. +- **Register Tiling (TILE=4)**: Each thread holds partial results for 4 Q heads in registers — no shared memory, no barrier synchronization — maximizing ALU throughput for grouped-query decode. +- **2-Pass FlashDecoding**: For long sequences (≥1024 tokens), a 2-pass reduction partitions the KV sequence across threadgroups to improve occupancy. +- **AutoTune**: Automatic sweep of block sizes per (architecture, GQA ratio, head count, sequence length), cached to `~/.cider_sdpa_tune.json`. + +### Usage + +```python +import cider + +# One-line patch — all downstream frameworks (mlx_lm, mlx_vlm) benefit automatically +cider.patch_sdpa() + +# Optional: run AutoTune once to find optimal block sizes for your chip +cider.autotune_sdpa() + +# Undo patch if needed +cider.unpatch_sdpa() +``` + +The patch is **fully transparent**: prefill (`Q_seq > 1`) and masked attention automatically fall back to MLX's native implementation. Only decode steps (`Q_seq=1`, no mask) route through Cider's optimized kernel. **Zero code changes** to mlx_lm or mlx_vlm. + +### SDPA Micro-Benchmark (Apple M5 Pro, D=128) + +Speedup vs MLX `mx.fast.scaled_dot_product_attention`: + +| Config | N=1K | N=2K | N=4K | N=8K | N=16K | N=32K | +|--------|------|------|------|------|-------|-------| +| MHA (H_q=32, H_kv=32) | 1.03x | 1.01x | 1.02x | 1.00x | 1.00x | 0.99x | +| GQA4 (H_q=32, H_kv=8) | 1.04x | 1.06x | 1.06x | 1.12x | 1.16x | 1.21x | +| GQA8 (H_q=32, H_kv=4) | 1.04x | 1.07x | 1.13x | 1.21x | 1.34x | **1.57x** | + +Speedup range: **0.99×–1.57×**. Correctness: 18/18 PASS (max diff ≤ 0.000122 vs MLX, fp16). 14 WIN / 4 TIE / 0 LOSS — zero regressions. + +### End-to-End Decode Speed with SDPA (Qwen3-VL-4B, Apple M5 Pro) + +Tested with `benchmarks/bench_e2e_wxa16.py` on real multi-image VLM trajectories (1–3 screenshots, 1334–3455 prompt tokens): + +| Config | Decode tok/s | vs Baseline | +|--------|:-----------:|:-----------:| +| FP16 | 32.0 | — | +| FP16 + Cider SDPA | 32.6 | +1.9% | +| W8A16 | 54.0 | — | +| W8A16 + Cider SDPA | **58.0** | **+7.4%** | +| W8A8 per-group | 51.0 | — | +| W8A8 per-group + Cider SDPA | 52.3 | +2.5% | +| W8A8 per-channel | 49.1 | — | +| W8A8 per-channel + Cider SDPA | **52.6** | **+7.1%** | + +Cider SDPA stacks on top of W8A8 prefill gains. The benefit scales with GQA ratio and KV cache length — models with higher GQA ratios (e.g., Qwen3-32B with GQA8) and longer contexts see larger gains. + +MLX's quantization is **weight-only**: QuantizedLinear dequantizes weights to FP16 and uses FP16 GEMM. While MLX's Steel NAX templates are generic enough to be instantiated with INT8 types (and would achieve identical raw matmul throughput — [see our transparent benchmark](benchmarks/mlx_native/cider_vs_mlx_int8.md)), MLX does not provide the quantization/dequantization pipeline needed for actual W8A8 inference. Cider fills this gap with fused quantize-matmul-dequant primitives. + +This SDK implements online INT8 activation quantization and INT8 TensorOps-based compute for the supported inference paths. + +### W8A8 Quantization Granularity + +| Granularity | Description | Speed | Precision | +|-------------|-------------|-------|-----------| +| **Per-channel** | One scale per output channel | Fastest (1.8x prefill) | Slightly lower | +| **Per-group (gs=128)** | One scale per 128 elements | Fast (1.5x prefill) | Moderate precision retention | +| **Per-group (gs=64)** | One scale per 64 elements | Moderate (1.3x prefill) | Higher precision | + +## Performance (Apple M5 Pro) + +### Individual Operator Latency + +Shape [N=10240, K=2560] +| M | PC(ms) | PG(ms) | w8a16 | w4a16 | PC/w8 | PC/w4 | PG/w8 | PG/w4| +|-----|------|------|-----|-----|-----|------|-----|----| +| 1 | 0.27ms | 0.26ms | 0.26ms | 0.18ms | 0.96x | 0.67x | 0.99x | 0.69x | +|128 | 0.34ms | 0.39ms | 0.49ms | 0.44ms | 1.43x | 1.28x | 1.26x | 1.13x | +|1024 | 1.23ms | 1.52ms | 2.24ms | 2.04ms | 1.82x | 1.66x | 1.47x | 1.34x| +|4096 | 4.41ms | 5.65ms | 8.12ms | 7.72ms | 1.84x | 1.75x | 1.44x | 1.37x | +|8192 | 8.71ms | 11.40ms | 16.23ms | 15.09ms | 1.86x | 1.73x | 1.42x | 1.32x| + + +Shape [N=2560, K=10240] +| M | PC(ms) | PG(ms) | w8a16 | w4a16 | PC/w8 | PC/w4 | PG/w8 | PG/w4 | +|--------|------|--------|-------| ---|--------|------|-------------|------------------| +| 1 | 0.25ms | 0.26ms | 0.26ms | 0.20ms | 1.03x | 0.78x | 0.98x | 0.75x | +|128 | 0.39ms | 0.41ms | 0.55ms | 0.46ms | 1.43x | 1.19x | 1.35x | 1.12x | +| 1024 | 1.31ms | 1.65ms | 2.35ms | 2.14ms | 1.80x | 1.64x | 1.43x | 1.30x | +| 4096 | 5.37ms | 6.79ms | 8.54ms | 8.04ms | 1.59x | 1.50x | 1.26x | 1.18x | +| 8192 | 10.97ms | 12.94ms | 17.28ms | 16.23ms | 1.58x | 1.48x | 1.34x | 1.25x | + +### End-to-End VLM + +**Qwen3-VL-2B** + +| Prompt Tokens | FP16 Prefill (tok/s) | W8A16 Prefill (tok/s) | **W8A8 PC Prefill (tok/s)** | FP16 Decode (tok/s) | W8A16 Decode (tok/s) | **W8A8 PC Decode (tok/s)** | +|:---:|:---:|:---:|:---:|:---:|:---:|:---:| +| 1334 | 3010 | 2065 | **3242** | 70 | 107 | **104** | +| 2393 | 2868 | 1847 | **2983** | 69 | 97 | **100** | +| 3455 | 2777 | 1741 | **2796** | 66 | 90 | **95** | + +**Qwen3-VL-4B** + + Prompt Tokens | FP16 Prefill (tok/s) | W8A16 Prefill (tok/s) | **W8A8 PC Prefill (tok/s)** | FP16 Decode (tok/s) | W8A16 Decode (tok/s) | **W8A8 PC Decode (tok/s)** | +|:---:|:---:|:---:|:---:|:---:|:---:|:---:| +| 1334 | 1884 | 1786 | **2186** | 32 | **56** | 54 | +| 2393 | 1815 | 1700 | **2028** | 31 | **55** | 52 | +| 3455 | 1755 | 1603 | **1881** | 30 | **52** | 49 | + + +### LLM Quantization: Precision vs. Speed Comparison + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +
ModelsQuantization Configurationwikitext2 PPL（↓）Prefill Time (s)（↓）Peak Memory (GB)（↓）
Qwen3-8BFP169.726179.918.93
W8A16 (mlx RTN)9.707221.312.07
W8A8 (per-channel)9.756123.511.32
W8A8 (per-group gs=64)9.744179.111.83
W8A8 (per-group gs=128)9.727165.811.61
Llama3-8BFP166.138175.818.32
W8A16 (mlx RTN)6.147236.911.46
W8A8 (per-channel)6.271123.310.69
W8A8 (per-group, gs=64)6.269178.711.19
W8A8 (per-group, gs=128)6.270155.710.98
+ +## NOTE +It's important to note that cider assumes the weights are **quantization-friendly** . This means you need to ensure your model has already undergone some post-training quantization calibration methods, such as GTPQ, SmoothQuant, or Quarot etc, to handle outliers or the model itself was trained using QAT. When cider transforms the model (e.g., with w8a8) and calculates the quantization scale, it defaults to the simplest **min-max** method, although it also supports the 99th percentile method (clip_percentile=99.9). If you encounter garbled output, it's highly likely that your model is affected by **outliers** . Therefore, some preprocessing is **necessary**. + + +## Requirements + + +- **Apple M5+** for INT8 TensorOps (M4 and below: installs as pure-Python, `is_available()` returns False) +- Python 3.12+ +- MLX >= 0.31 +- nanobind >= 2.12 (only needed on M5+ for C++ build) +- CMake >= 3.27 (only needed on M5+ for C++ build) + +## Install + +```bash +pip install -e . +``` + +On M5+, this runs CMake to compile the C++ extension, then installs the Python package. +On M4 and below, only the Python package is installed (no compilation errors). + +## Quick Start + +### One-line Model Conversion (Recommended) + +```python +from cider import convert_model, is_available + +model, proc = load("path/to/model") + +if is_available(): + convert_model(model) + # CiderLinear auto-detects: + # seq_len > 1 → W8A8 INT8 TensorOps (faster prefill) + # seq_len == 1 → INT8 MV kernel (near-native decode speed) +else: + pass # Falls back to standard MLX inference on M4 +``` + +**Important** +When quantizing Vision-Language Models (VLMs), the vision transformer (ViT) is generally not replaced. Directly using convert_model will quantize the vision model's linear layers as well, which typically causes accuracy drop. For VLMs, we recommend calling convert_model(model.language_model) to apply existing quantization methods like GPTQ, SmoothQuant, and AWQ to the language model only. + +Tested on selected MLX transformer models, including Qwen3, Qwen3-VL and Llama3 families. Other architectures may require adaptation. + + +### Layer-level API + +```python +import numpy as np +import mlx.core as mx +from cider import W8A8Linear, W4A8Linear, is_available + +assert is_available(), "Requires Apple M5+" + +# Prepare weight +W = np.random.randn(4096, 4096).astype(np.float16) + +# W8A8 linear (per-channel) +from cider.ops import quantize_weight_int8 +w_int8, scale = quantize_weight_int8(W) +layer = W8A8Linear( + w_int8=mx.array(w_int8), scale_w=mx.array(scale), + group_size=0, in_features=4096, out_features=4096 +) +x = mx.random.normal((32, 4096)).astype(mx.float16) +y = layer(x) # lazy — builds MLX graph +mx.eval(y) # GPU executes + +# W4A8 linear (half the weight memory) +layer4 = W4A8Linear.from_weights(W) +y4 = layer4(x) +mx.eval(y4) +``` + +## Low-Level API + +```python +from cider import perchannel_linear, w4a8_linear, quantize_weight_int8, pack_weight_int4 + +# Quantize weights (numpy, offline) +w_int8, scale = quantize_weight_int8(W_np) +packed_w4, scale4 = pack_weight_int4(W_np) + +# Primitive calls (return lazy mx.array) +y = perchannel_linear(x, mx.array(w_int8), mx.array(scale)) +y4 = w4a8_linear(x, mx.array(packed_w4), mx.array(scale4)) +``` + +## Project Structure + +``` +cider/ +├── cider/ # Python package +│ ├── __init__.py # Public API (conditional on is_available) +│ ├── ops.py # Primitive wrappers + quantize helpers +│ ├── nn.py # CiderLinear, W4A8Linear (nn.Module) +│ ├── convert.py # convert_model() high-level API +│ ├── attention/ # Optimized SDPA +│ │ ├── __init__.py +│ │ └── sdpa.py # patch_sdpa / unpatch_sdpa / autotune_sdpa +│ └── kernels/ # Metal shaders (bundled) +│ ├── w8a8_matmul.metal # W8A8 GEMM (prefill, M>1) +│ ├── w8a8_int8_mv.metal # W8A8 per-channel MV (decode, M=1) +│ ├── w8a8_quantize.metal # Per-token activation quantization +│ ├── w4a8_matmul.metal # W4A8 GEMM (prefill) +│ ├── pergroup_int8_gemm.metal # Per-group GEMM (prefill) +│ ├── pergroup_int8_mv.metal # Per-group MV (decode) +│ ├── cider_sdpa_vector.h # SDPA kernel header (1-pass & 2-pass) +│ └── cider_sdpa_vector.metal # SDPA kernel template instantiations +├── csrc/ # C++ MLX primitives (nanobind, M5+ only) +│ ├── include/ +│ │ ├── w8a8_primitive.h +│ │ ├── w4a8_primitive.h +│ │ ├── pergroup_primitive.h +│ │ └── sdpa_primitive.h +│ └── src/ +│ ├── w8a8_primitive.mm +│ ├── w4a8_primitive.mm +│ ├── pergroup_primitive.mm +│ ├── sdpa_primitive.mm # SDPA 1-pass & 2-pass primitives +│ └── prim_bindings.cpp +├── benchmarks/ +│ ├── bench_e2e_wxa16.py # End-to-end VLM benchmark (W8A8 + SDPA) +│ ├── bench_sdpa.py # SDPA micro-benchmark (correctness + speed vs MLX) +│ ├── bench_full.py # Isolated kernel latency (per-channel/per-group vs MLX) +│ ├── test_bitexact.py # Numerical correctness verification +│ └── mlx_native/ # MLX native INT8 comparison +├── tutorial/ +│ ├── how_to_write_efficient_int_gemm_m5_en.md +│ └── how_to_write_efficient_int_gemm_m5_zh.md +├── tools/ +│ ├── eval_ppl_all.py # Unified PPL eval (FP16/W8A16/per-channel/per-group) +│ ├── convert_compressed_tensors_to_mlx.py +│ └── smoothquant.py # SmoothQuant calibration +├── examples/ +│ └── basic_usage.py +├── vlm_service/ # OpenAI-style VLM inference server +│ ├── server.py # FastAPI server (streaming + non-streaming) +│ ├── core_infer.py # HMInference engine (singleton) +│ ├── custom_qwen3vl.py # Custom Qwen3-VL generation loop +│ ├── config.py # Config loader +│ ├── bench_client.py # Server benchmark client +│ └── client.py # API client example +├── config/ +│ └── config.yaml # Server & model configuration +├── experimental/ # ANE+GPU hybrid tensor parallelism (M4) +│ ├── split_linear.py # SplitLinear + ANEBridge + patch_model() +│ ├── bench.py # End-to-end benchmark +│ ├── libane_bridge_v6.m # ANE private API bridge (Obj-C source) +│ └── README.md +├── CMakeLists.txt +├── pyproject.toml +├── setup.py # Conditional build (M5+: full, M4: pure-Python) +└── README.md +``` + +## VLM Inference Service + +`vlm_service/` provides a ready-to-use **OpenAI-style** VLM inference server with W8A8 acceleration. + +### Quick Start + +1. **Configure** `config/config.yaml`: + +```yaml +model_name_or_path: /path/to/your/model # MLX VLM model (e.g., Qwen3-VL-2B W8A16) +sampling: + max_new_tokens: 1024 + temperature: 1.0 + top_p: 1.0 +server: + host: 0.0.0.0 + port: 8341 + ttl: 1800 +w8a8: + mode: 'off' # 'auto' | 'on' | 'off' +``` + +- `auto`: Enable W8A8 if hardware supports it, fallback to default otherwise +- `on`: Force W8A8 (error if unsupported). "When 'on' is selected, it means your model needs to perform online activation quantization. In this case, Cider itself does **not** guarantee quantization accuracy, and you need to apply some quantization algorithms yourself, such as SmoothQuant, QuaRot, GPTQ, or even QAT, to ensure that the accuracy does not degrade significantly after activation quantization. This option simply provides a way for you to leverage the hardware's computational advantages when your model applies W8A8, rather than just simulating quantization." +- `off`: Disable W8A8, use standard MLX inference + +2. **Start the server**: + +```bash +cd vlm_service +python server.py --config ../config/config.yaml +``` + +3. **Send requests** (OpenAI-style API): + +```bash +# Text-only +curl http://localhost:8341/v1/chat/completions \ + -H "Content-Type: application/json" \ + -d '{ + "model": "vlm", + "messages": [{"role": "user", "content": "Hello"}], + "stream": false + }' + +# With image (base64) +curl http://localhost:8341/v1/chat/completions \ + -H "Content-Type: application/json" \ + -d '{ + "model": "vlm", + "messages": [{"role": "user", "content": [ + {"type": "image_url", "image_url": {"url": "data:image/png;base64,..."} }, + {"type": "text", "text": "What is in this image?"} + ]}], + "stream": true + }' +``` + +### API Endpoints + +| Endpoint | Method | Description | +|----------|--------|-------------| +| `/v1/chat/completions` | POST | Chat completion (stream / non-stream) | +| `/v1/models` | GET | List available models | +| `/health` | GET | Health check | +| `/v1/queue` | GET | Request queue status | + +### How W8A8 Works in the Service + +When `w8a8.mode` is `auto` or `on`, the server calls `cider.convert_model()` at startup to replace all Linear layers with `CiderLinear`. During inference: + +- **Prefill** (processing input tokens, seq_len > 1): Uses W8A8 INT8 GEMM for faster computation +- **Decode** (generating tokens one by one, seq_len == 1): Uses INT8 MV kernel (near-native speed) + +No code changes needed — the switching is automatic based on input sequence length. + +## Architecture + +### MLX Custom Primitives + +Both W8A8Linear and W4A8Linear are implemented as `mlx::core::Primitive` subclasses. This means: + +1. **Lazy evaluation**: `y = layer(x)` builds a graph node, not immediate computation +2. **Graph composition**: Multiple primitive calls compose into a single MLX graph +3. **Stream scheduling**: MLX's scheduler handles GPU dispatch order + +### Metal Kernel Pipeline + +Each primitive dispatches Metal compute kernels: + +**Prefill (M > 1):** +1. **quantize_per_token**: FP16 activations → INT8 + per-token scales +2. **matmul_fused_dequant**: INT8 × INT8 → INT32 → FP16 (with fused scale dequantization) + +**Decode (M = 1):** +- **int8_mv**: Direct INT8 matrix-vector product with on-the-fly weight dequantization (no activation quantization needed) + +For W4A8, the GEMM step includes inline INT4→INT8 unpacking in the fragment load. + +### TensorOps matmul2d + +The INT8 GEMM uses Apple's `mpp::tensor_ops::matmul2d(16, 32, 16)` — hardware-accelerated INT8×INT8→INT32 matrix multiply available on M5+ via Metal 4's `cooperative_tensor` API. This is the same hardware instruction available to MLX's NAX templates. Cider's kernel adds fused dequantization (INT32 × scales → FP16) in the store phase, avoiding an extra device memory round-trip. See [kernel comparison](benchmarks/mlx_native/cider_vs_mlx_int8.md) for details. + +### Tile Configurations + +| Config | BM | BN | BK | SK | Threads | Use When | +|--------|----|----|----|----|---------|----------| +| Large | 128 | 128 | 512 | 32 | 512 | M > 64 | +| Small | 32 | 128 | 512 | 32 | 128 | M ≤ 64 | + +Auto-selected based on M. L2 cache swizzle dispatch included. + +## ANE+GPU Heterogeneous Tensor Parallelism (experimental) + +We found that during inference on Mac, only two hardware computing units—GPU and CPU—were utilized, while the ANE (Apple Neural Engine) computing unit on Mac remained idle. We identified this as a potential optimization opportunity. Inspired by [maderix/ANE](https://github.com/maderix/ANE), we conducted experimental work on a hybrid ANE+GPU inference mode. Currently, we apply this approach to tensor parallel computing. On the M4 chip, during synchronous-only forward inference (MLX natively uses a technique called lazy evaluation, which reduces synchronization overhead; in end-to-end testing, the hybrid inference currently shows no advantage, mainly because we have not yet implemented this using MLX's lazy evaluation—this remains future work), we observed approximately **3%~16%** performance improvement compared to pure GPU inference under synchronize pipeline. This remains exploratory work, and end-to-end gains are currently limited by the lack of a lazy-evaluation-compatible implementation. + + +During LLM prefill, the GPU's matrix units are fully occupied — but the **Apple Neural Engine sits completely idle**. ANE Split exploits this by splitting each linear layer's GEMM along output channels: + +- **ANE** computes ~65% of output channels (FP32, via reverse-engineered private `_ANEClient` API) +- **GPU** computes the remaining ~35% (FP16, standard MLX matmul) +- Both run **concurrently**, and results are concatenated + +This is a form of **heterogeneous tensor parallelism** — not data parallelism, not pipeline parallelism — exploiting two distinct compute units on the same SoC. + +### Performance (Apple M4, Qwen3-VL-2B Prefill) + +| seq | W8A16 GPU | SplitLinear | Speedup vs W8A16 | +|-----|----------|-----------|-------------| +| 512 | 639.9 ms | **615.9 ms** | **1.039×** | +| 1024 | 1348.6 ms | **1156.9 ms** | **1.17×** | + +In the tested benchmark cases, cosine similarity was close to 1.0 and top-1 token agreement was 100%. + + +### Key Design Choices + +- **Prefill only**: Decode falls back to original GPU linear (zero overhead) +- **Shared input preparation**: Q/K/V and Gate/Up projections share a single input transpose+numpy copy via `_InputGroup` +- **Auto-routing**: Down projections (IC > 2×OC) stay GPU-only where ANE is inefficient +- **Short-seq bypass**: Sequences < 192 tokens skip splitting (overhead > benefit) + +See [`experimental/README.md`](experimental/README.md) for full documentation, usage, and build instructions and limitations. + +> **Note:** ANE Split is tested on M4. M5 introduced ANE architecture changes that may break the private API bridge — not yet validated on M5. + +## Quantization + +| Component | Scheme | Granularity | +|-----------|--------|-------------| +| W8A8 weights | Symmetric INT8 | Per-channel or per-group (gs=64/128) | +| W4A8 weights | Symmetric INT4 (zp=8) | Per-column | +| Activations | Symmetric INT8 | Per-token | +| Accumulation | INT32 | — | +| Output dequant | `C_fp16 = C_int32 * s_act * s_weight` | Per-element | + +## Limitations + +- **M=1 individual operator**: Per-channel MV kernel is slower than MLX W4A16 for isolated decode calls. The per-group MV kernel is within 5% of MLX W8A16 decode speed in end-to-end benchmarks. +- **Apple M5+ only** for INT8 TensorOps: M4 and below installs but `is_available()` returns False. +- **W4A8 slower than W8A8**: INT4→INT8 unpack ALU overhead (Metal 4 matmul2d has no native INT4 operand). + +## Tools + +### Unified PPL Evaluation + +```bash +# Run all 5 configurations in one script +python tools/eval_ppl_all.py --num-samples 50 + +# Evaluates: FP16, W8A16 (MLX native), W8A8 per-channel, per-group(gs=64), per-group(gs=128) +# Outputs comparison table at the end +``` + +## Roadmap + +- [x] One-line model conversion API (`convert_model`, auto prefill/decode) +- [x] Automatic dtype handling (float16 / bfloat16) +- [x] Per-channel and per-group W8A8 quantization +- [x] Dedicated decode MV kernel (matches native MLX speed) +- [x] Conditional compilation (M4 graceful fallback) +- [x] mlx_vlm and mlx_lm integration examples +- [x] Optimized SDPA decode kernel (FlashInfer-inspired, GQA-aware) +- [x] One-line `patch_sdpa()` for transparent decode acceleration +- [x] SDPA AutoTune (per-architecture block size sweep) +- [ ] ANE primitives lazy evaluation +- [ ] Integrated Pruning Feature +- [ ] KVCache quantization + +## Authors + +Multimodal Team, Mininglamp Technology + +For bug reports, feature requests, and usage questions, please open an issue in this repository. + + + +## Citation + +If you find this work useful, please cite: + +```bibtex +@software{wang2026cider, + author = {Multimodal Team, Mininglamp Technology}, + title = {Cider: Exploiting Unused INT8 TensorOps for Faster LLM Prefill on Apple Silicon}, + year = {2026}, + howpublished = {https://github.com/Mininglamp-AI/cider} +} +``` + +## License + +MIT + +## Acknowledgments + +- [MLX](https://github.com/ml-explore/mlx) by Apple — primitive API, NAXFrag kernel architecture +- [FlashInfer](https://github.com/flashinfer-ai/flashinfer) — inspired contiguous chunk GQA scheduling and register tiling for SDPA decode kernel +- Metal 4 MetalPerformancePrimitives for INT8 TensorOps +- [maderix/ANE](https://github.com/maderix/ANE) — inspired and informed our ANE+GPU tensor-parallel implementation diff --git a/README.wehub.md b/README.wehub.md new file mode 100644 index 0000000..051591e --- /dev/null +++ b/README.wehub.md @@ -0,0 +1,7 @@ +# WeHub 来源说明 + +- 原始项目：`Mininglamp-AI/cider` +- 原始仓库：https://github.com/Mininglamp-AI/cider +- 导入方式：上游默认分支的最新快照 +- 原作者、版权和许可证信息以原始仓库及本仓库 LICENSE 为准 +- 本文件仅用于记录来源，不代表 WeHub 是原项目作者 diff --git a/benchmarks/bench_e2e_wxa16.py b/benchmarks/bench_e2e_wxa16.py new file mode 100644 index 0000000..05fa3f5 --- /dev/null +++ b/benchmarks/bench_e2e_wxa16.py @@ -0,0 +1,347 @@ +#!/usr/bin/env python3 +""" +End-to-end benchmark on REAL trajectory data: + 1. GPU FP16 baseline + 2. GPU W8A16 (native quantized_matmul) + 3. GPU W8A8_pg (INT8 TensorOps via cider, pergroup) + 4. GPU W8A8_pc (INT8 TensorOps via cider, perchannel) + 5-8. Above 4 configs + Cider SDPA patch (decode acceleration) + +Uses replay_prompt.build_prompt_at_step() to build real prompts with real screenshots. +Compares both accuracy (action output) and speed (prefill + decode). +""" +import sys, os, time, gc, re, json, io, base64 +import numpy as np +import mlx.core as mx +import mlx.nn as nn +from PIL import Image +from pathlib import Path +ROOT_DIR = str(Path(__file__).parent.parent) +sys.path.insert(0, ROOT_DIR) +sys.path.insert(0, os.path.join(ROOT_DIR, "vlm_service")) + +from session_data.replay_prompt import build_prompt_at_step +from vlm_service.custom_qwen3vl import custom_stream_generate +from mlx_vlm.utils import load as vlm_load +home_dir = os.path.expanduser("~") +FP16_MODEL = os.path.join(home_dir, 'Downloads/checkpoint-50349') +W8A16_MODEL = os.path.join(home_dir, 'Downloads/checkpoint-50349-mlx_w8a16') +SESSION_DIR = os.path.join(ROOT_DIR, "session_data") +STEPS = [0, 1, 2] # step 0: 1img, step 1: 2imgs, step 2: 3imgs +MAX_TOKENS = 200 +PREFILL_STEP = 8192 +N_SPEED_RUNS = 2 # per-step speed runs + + +def load_step(session_dir, step): + """Load prompt + images for a given step.""" + data = build_prompt_at_step(session_dir, step) + pil_images = [Image.open(io.BytesIO(base64.b64decode(b))) for b in data['images']] + return data, pil_images + + +def build_messages(data): + """Build chat messages from replay_prompt data, with placeholders.""" + prompt_text = data['prompt'] + n_images = prompt_text.count('') + parts = prompt_text.split('') + content = [] + for i, part in enumerate(parts): + if part: + content.append({"type": "text", "text": part}) + if i < n_images: + content.append({"type": "image"}) + + messages = [ + {"role": "system", "content": data['system_prompt']}, + {"role": "user", "content": content}, + ] + return messages + + +def run_generate(model, processor, messages, pil_images): + """Run one generation, return (text, timing_dict).""" + prompt = processor.tokenizer.apply_chat_template( + messages, tokenize=False, add_generation_prompt=True + ) + + gc.collect() + mx.clear_cache() + mx.reset_peak_memory() + + text = "" + prefill_time = 0 + decode_tps = 0 + prompt_tokens = 0 + gen_tokens = 0 + + try: + for resp in custom_stream_generate( + model, processor, + prompt=prompt, + image=pil_images, + max_tokens=MAX_TOKENS, + temperature=0.0, + prefill_step_size=PREFILL_STEP, + verbose=False, + ): + text += resp.text + prefill_time = resp.prompt_tokens / resp.prompt_tps if resp.prompt_tps > 0 else 0 + decode_tps = resp.generation_tps + prompt_tokens = resp.prompt_tokens + gen_tokens = resp.generation_tokens + except UnicodeDecodeError: + # mlx_vlm detokenizer sometimes fails on partial utf-8 sequences + pass + + total_time = prefill_time + (gen_tokens / decode_tps if decode_tps > 0 else 0) + + return text, { + 'prefill_ms': prefill_time * 1000, + 'prefill_tps': prompt_tokens / prefill_time if prefill_time > 0 else 0, + 'decode_tps': decode_tps, + 'prompt_tokens': prompt_tokens, + 'gen_tokens': gen_tokens, + 'total_ms': total_time * 1000, + } + + +def extract_action(text): + """Extract action from model output.""" + m = re.search(r'(.*?)', text, re.DOTALL) + return m.group(1).strip() if m else None + + +def extract_coords(action_str): + """Extract (x, y) from action string.""" + if not action_str: + return None + nums = re.findall(r'\d+', action_str) + if len(nums) >= 2: + return (int(nums[0]), int(nums[1])) + return None + + +def main(): + print("=" * 80) + print(" E2E Benchmark: FP16 / W8A16 / W8A8 × MLX SDPA / Cider SDPA") + print("=" * 80) + print(f" Session: {SESSION_DIR}") + print(f" Steps: {STEPS}") + print(f" Max tokens: {MAX_TOKENS}") + print(f" Speed runs per step: {N_SPEED_RUNS}") + + # Load all steps + print("\nLoading trajectory data...") + step_data = {} + for s in STEPS: + data, images = load_step(SESSION_DIR, s) + msgs = build_messages(data) + step_data[s] = (data, images, msgs) + print(f" Step {s}: {len(images)} images") + + # Import cider for SDPA patching + import cider + + # configs: (key, model_path, use_cider_quant, use_cider_sdpa, label) + configs = [ + # Without Cider SDPA (MLX default attention) + ('fp16', FP16_MODEL, False, False, "FP16"), + ('fp16+sdpa', FP16_MODEL, False, True, "FP16+CiderSDPA"), + ('w8a16', W8A16_MODEL, False, False, "W8A16"), + ('w8a16+sdpa', W8A16_MODEL, False, True, "W8A16+CiderSDPA"), + ('w8a8_pg', W8A16_MODEL, True, False, "W8A8_pg"), + ('w8a8_pg+sdpa', W8A16_MODEL, True, True, "W8A8_pg+CiderSDPA"), + ('w8a8_pc', FP16_MODEL, True, False, "W8A8_pc"), + ('w8a8_pc+sdpa', FP16_MODEL, True, True, "W8A8_pc+CiderSDPA"), + ] + + all_results = {} + + for cfg_key, model_path, use_cider_quant, use_cider_sdpa, label in configs: + print(f"\n{'='*80}") + print(f" {label}") + print(f"{'='*80}") + + # Load model + print(f" Loading {model_path}...") + model, proc = vlm_load(model_path) + + # Apply cider quantization if needed + if use_cider_quant: + from cider import convert_model + stats = convert_model(model) + print(f" [cider quant] {stats}") + + # Apply cider SDPA if needed + if use_cider_sdpa: + cider.patch_sdpa(verbose=True) + else: + cider.unpatch_sdpa(verbose=False) + + cfg_results = {} + + for s in STEPS: + data, images, msgs = step_data[s] + print(f"\n Step {s} ({len(images)} imgs):") + + # Warmup + text, timing = run_generate(model, proc, msgs, images) + print(f" Warmup: {timing['prompt_tokens']} prompt tok, " + f"prefill {timing['prefill_ms']:.0f}ms " + f"({timing['prefill_tps']:.0f} tok/s), " + f"decode {timing['decode_tps']:.1f} tok/s") + + # Accuracy run (use warmup result) + action = extract_action(text) + coords = extract_coords(action) + + # Speed runs + timings = [] + for r in range(N_SPEED_RUNS): + _, t = run_generate(model, proc, msgs, images) + timings.append(t) + print(f" Run {r+1}: prefill {t['prefill_ms']:.0f}ms " + f"({t['prefill_tps']:.0f} tok/s) | " + f"decode {t['decode_tps']:.1f} tok/s | " + f"total {t['total_ms']:.0f}ms") + + cfg_results[s] = { + 'text': text, + 'action': action, + 'coords': coords, + 'timings': timings, + 'prompt_tokens': timings[0]['prompt_tokens'], + } + + print(f" Action: {action}") + + all_results[cfg_key] = cfg_results + + # Free model before loading next + del model, proc + gc.collect() + mx.clear_cache() + + # Ensure SDPA is unpatched after benchmark + cider.unpatch_sdpa(verbose=False) + + # ── Speed Summary ── + print(f"\n{'='*80}") + print(f" SPEED SUMMARY (median of {N_SPEED_RUNS} runs)") + print(f"{'='*80}") + + cfg_keys = [c[0] for c in configs] + cfg_labels = {c[0]: c[4] for c in configs} + + # Table: per-step results + print(f"\n {'Config':<22s}", end="") + for s in STEPS: + print(f" | Step{s} prefill decode total", end="") + print() + print(f" {'─'*22}", end="") + for _ in STEPS: + print(f" | {'─'*30}", end="") + print() + + for k in cfg_keys: + line = f" {cfg_labels[k]:<22s}" + for s in STEPS: + ts = all_results[k][s]['timings'] + pf = float(np.median([t['prefill_tps'] for t in ts])) + dc = float(np.median([t['decode_tps'] for t in ts])) + tt = float(np.median([t['total_ms'] for t in ts])) + line += f" | {pf:>7.0f}t/s {dc:>5.1f}t/s {tt:>5.0f}ms" + print(line) + + # ── Decode speed comparison (SDPA impact) ── + print(f"\n{'='*80}") + print(f" DECODE SPEED: MLX SDPA vs Cider SDPA") + print(f"{'='*80}") + + # Group by base config + base_configs = ['fp16', 'w8a16', 'w8a8_pg', 'w8a8_pc'] + for base in base_configs: + sdpa_key = base + '+sdpa' + if base not in all_results or sdpa_key not in all_results: + continue + + print(f"\n {cfg_labels[base]}:") + for s in STEPS: + ts_base = all_results[base][s]['timings'] + ts_sdpa = all_results[sdpa_key][s]['timings'] + dc_base = float(np.median([t['decode_tps'] for t in ts_base])) + dc_sdpa = float(np.median([t['decode_tps'] for t in ts_sdpa])) + speedup = dc_sdpa / dc_base if dc_base > 0 else 0 + pf_base = float(np.median([t['prefill_tps'] for t in ts_base])) + pf_sdpa = float(np.median([t['prefill_tps'] for t in ts_sdpa])) + print(f" Step {s}: decode {dc_base:.1f} → {dc_sdpa:.1f} tok/s " + f"({speedup:.2f}x) | prefill {pf_base:.0f} → {pf_sdpa:.0f} tok/s") + + # ── Accuracy comparison ── + print(f"\n{'='*80}") + print(f" ACCURACY: SDPA patch should not change outputs") + print(f"{'='*80}") + + for base in base_configs: + sdpa_key = base + '+sdpa' + if base not in all_results or sdpa_key not in all_results: + continue + mismatches = 0 + for s in STEPS: + a1 = all_results[base][s]['action'] + a2 = all_results[sdpa_key][s]['action'] + if a1 != a2: + mismatches += 1 + c1 = all_results[base][s]['coords'] + c2 = all_results[sdpa_key][s]['coords'] + if c1 and c2: + diff = (abs(c1[0]-c2[0]), abs(c1[1]-c2[1])) + close = all(d <= 5 for d in diff) + tag = f"{'≈' if close else '≠'} diff=({diff[0]},{diff[1]})" + else: + tag = "≠" + print(f" {cfg_labels[base]} Step {s}: {tag}") + print(f" base: {a1}") + print(f" sdpa: {a2}") + if mismatches == 0: + print(f" {cfg_labels[base]}: all {len(STEPS)} steps IDENTICAL ✅") + + # ── Overall Summary ── + print(f"\n{'='*80}") + print(f" OVERALL DECODE tok/s (median across all steps)") + print(f"{'='*80}") + for k in cfg_keys: + all_decode = [] + for s in STEPS: + ts = all_results[k][s]['timings'] + all_decode.append(float(np.median([t['decode_tps'] for t in ts]))) + med = np.median(all_decode) + print(f" {cfg_labels[k]:<22s}: {med:.1f} tok/s") + + # Save + out_path = '/tmp/e2e_wxa16_sdpa_results.json' + save_data = {} + for k in cfg_keys: + save_data[k] = {} + for s in STEPS: + r = all_results[k][s] + save_data[k][str(s)] = { + 'action': r['action'], + 'coords': r['coords'], + 'prompt_tokens': r['prompt_tokens'], + 'timings_median': { + 'prefill_ms': float(np.median([t['prefill_ms'] for t in r['timings']])), + 'prefill_tps': float(np.median([t['prefill_tps'] for t in r['timings']])), + 'decode_tps': float(np.median([t['decode_tps'] for t in r['timings']])), + 'total_ms': float(np.median([t['total_ms'] for t in r['timings']])), + }, + } + with open(out_path, 'w') as f: + json.dump(save_data, f, indent=2) + print(f"\n Results saved to {out_path}") + print(f"{'='*80}") + + +if __name__ == '__main__': + main() diff --git a/benchmarks/bench_full.py b/benchmarks/bench_full.py new file mode 100644 index 0000000..8b5557b --- /dev/null +++ b/benchmarks/bench_full.py @@ -0,0 +1,106 @@ +#!/usr/bin/env python3 +"""Full kernel benchmark: Cider per-channel & per-group vs MLX w8a16 & w4a16. + +Tests M = {1, 128, 1024, 4096, 8192} across Qwen3-2B/7B shapes. +""" +import sys, os, time +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +import numpy as np +import mlx.core as mx +from cider import perchannel_linear, pergroup_linear, quantize_weight_int8 + +WARMUP = 5 +REPEAT = 20 + +# Qwen3-VL shapes: [N, K] +SHAPES_NK = [ + (3584, 3584), # qkv square + (18944, 3584), # up-proj + (3584, 18944), # down-proj + (2560, 2560), # 2B square + (10240, 2560), # 2B up + (2560, 10240), # 2B down +] + +M_VALUES = [1, 128, 1024, 4096, 8192] +GROUP_SIZE = 128 + + +def timed(fn, warmup=WARMUP, repeat=REPEAT): + for _ in range(warmup): + y = fn(); mx.eval(y) + times = [] + for _ in range(repeat): + t0 = time.perf_counter() + y = fn(); mx.eval(y) + times.append(time.perf_counter() - t0) + return np.median(times) * 1000 + + +def main(): + print(f"Full Kernel Benchmark (warmup={WARMUP}, repeat={REPEAT}, gs={GROUP_SIZE})") + print() + + for N, K in SHAPES_NK: + print(f"=== Shape [N={N}, K={K}] ===") + print(f"{'M':>6s} | {'PC(ms)':>8s} {'PG(ms)':>8s} {'w8a16':>8s} {'w4a16':>8s} | {'PC/w8':>7s} {'PC/w4':>7s} {'PG/w8':>7s} {'PG/w4':>7s}") + print("-" * 95) + + # Prepare per-channel weights + np.random.seed(42) + w_fp = np.random.randn(N, K).astype(np.float32) + # Per-channel: scale per row + absmax = np.abs(w_fp).max(axis=1, keepdims=True).clip(min=1e-8) + scale_pc = (absmax / 127.0).astype(np.float32) + w_int8_pc = np.clip(np.round(w_fp / scale_pc), -127, 127).astype(np.int8) + w_int8_pc_mx = mx.array(w_int8_pc) + scale_pc_mx = mx.array(scale_pc.squeeze()) + + # Per-group weights (symmetric, gs=GROUP_SIZE) + ng = (K + GROUP_SIZE - 1) // GROUP_SIZE + w_pg = w_fp.copy() + scale_pg = np.zeros((N, ng), dtype=np.float32) + w_int8_pg = np.zeros((N, K), dtype=np.int8) + for g in range(ng): + k0 = g * GROUP_SIZE + k1 = min(k0 + GROUP_SIZE, K) + blk = w_pg[:, k0:k1] + amax = np.abs(blk).max(axis=1, keepdims=True).clip(min=1e-8) + s = (amax / 127.0).astype(np.float32) + scale_pg[:, g] = s.squeeze() + w_int8_pg[:, k0:k1] = np.clip(np.round(blk / s), -127, 127).astype(np.int8) + w_int8_pg_mx = mx.array(w_int8_pg) + scale_pg_mx = mx.array(scale_pg) + + # MLX w8a16 (bits=8, gs=64) + w_mx = mx.array(w_fp.astype(np.float16)) + w8_q, w8_s, w8_b = mx.quantize(w_mx, bits=8, group_size=64) + + # MLX w4a16 (bits=4, gs=64) + w4_q, w4_s, w4_b = mx.quantize(w_mx, bits=4, group_size=64) + + mx.eval(w_int8_pc_mx, scale_pc_mx, w_int8_pg_mx, scale_pg_mx, + w8_q, w8_s, w8_b, w4_q, w4_s, w4_b) + + for M in M_VALUES: + x = mx.random.normal((M, K)).astype(mx.float16) + mx.eval(x) + + t_pc = timed(lambda: perchannel_linear(x, w_int8_pc_mx, scale_pc_mx)) + t_pg = timed(lambda: pergroup_linear(x, w_int8_pg_mx, scale_pg_mx, GROUP_SIZE)) + t_w8 = timed(lambda: mx.quantized_matmul(x, w8_q, w8_s, w8_b, bits=8, group_size=64)) + t_w4 = timed(lambda: mx.quantized_matmul(x, w4_q, w4_s, w4_b, bits=4, group_size=64)) + + r_pc_w8 = t_w8 / t_pc if t_pc > 0 else 0 + r_pc_w4 = t_w4 / t_pc if t_pc > 0 else 0 + r_pg_w8 = t_w8 / t_pg if t_pg > 0 else 0 + r_pg_w4 = t_w4 / t_pg if t_pg > 0 else 0 + + print(f"{M:>6d} | {t_pc:>6.2f}ms {t_pg:>6.2f}ms {t_w8:>6.2f}ms {t_w4:>6.2f}ms | {r_pc_w8:>5.2f}x {r_pc_w4:>5.2f}x {r_pg_w8:>5.2f}x {r_pg_w4:>5.2f}x") + + print() + + +if __name__ == "__main__": + main() diff --git a/benchmarks/bench_sdpa.py b/benchmarks/bench_sdpa.py new file mode 100644 index 0000000..380462d --- /dev/null +++ b/benchmarks/bench_sdpa.py @@ -0,0 +1,155 @@ +"""benchmarks/bench_sdpa.py — Cider SDPA vs MLX native: correctness + performance. + +Usage: + python benchmarks/bench_sdpa.py # Default configs + python benchmarks/bench_sdpa.py --autotune # Run AutoTune first + python benchmarks/bench_sdpa.py --iters 500 # More iterations for stable results +""" +import argparse +import math +import sys +import time + +sys.path.insert(0, ".") + +import mlx.core as mx + + +def bench_one(fn, q, k, v, scale, warmup, iters): + """Benchmark a single SDPA function. Returns p10 latency in microseconds.""" + for _ in range(warmup): + r = fn(q, k, v, scale=scale) + mx.eval(r) + + times = [] + for _ in range(iters): + t0 = time.perf_counter() + r = fn(q, k, v, scale=scale) + mx.eval(r) + times.append((time.perf_counter() - t0) * 1e6) + + times.sort() + p10 = times[max(0, len(times) // 10)] + median = times[len(times) // 2] + return p10, median + + +def main(): + parser = argparse.ArgumentParser(description="Cider SDPA benchmark") + parser.add_argument("--warmup", type=int, default=50) + parser.add_argument("--iters", type=int, default=200) + parser.add_argument("--autotune", action="store_true", help="Run AutoTune before benchmark") + parser.add_argument("--D", type=int, default=128, help="Head dimension") + args = parser.parse_args() + + from cider.attention.sdpa import scaled_dot_product_attention, _mlx_sdpa + + # GPU info + info = mx.metal.device_info() if hasattr(mx.metal, 'device_info') else mx.device_info() + arch = info.get("architecture", "unknown") + print(f"GPU: {arch}") + print(f"MLX: {mx.__version__}") + print(f"Config: warmup={args.warmup}, iters={args.iters}, D={args.D}") + + if args.autotune: + import cider + print("\n[AutoTune]") + cider.autotune_sdpa(D=args.D, iters=args.iters // 4) + print() + + configs = [ + # (B, H_q, H_kv, N, label) + (1, 32, 32, 1024, "MHA N=1K"), + (1, 32, 32, 2048, "MHA N=2K"), + (1, 32, 32, 4096, "MHA N=4K"), + (1, 32, 32, 8192, "MHA N=8K"), + (1, 32, 32, 16384, "MHA N=16K"), + (1, 32, 32, 32768, "MHA N=32K"), + ] + for H_q in [32, 64, 96, 128]: + for gqa_factor in [2, 4, 8, 16]: + H_kv = H_q // gqa_factor + for N in [1024, 2048, 4096, 8192, 16384, 32768]: + label = f"GQA{H_q}_{gqa_factor} N={N // 1024}K" + configs.append((1, H_q, H_kv, N, label)) + + D = args.D + + # ── Correctness ───────────────────────────────────────────── + print("\n" + "=" * 72) + print("CORRECTNESS (max |cider - mlx|)") + print("=" * 72) + print(f"{'Config':<16} {'MaxDiff':>12} {'Status':>8}") + print("-" * 40) + + mx.random.seed(42) + all_correct = True + for B, Hq, Hkv, N, label in configs: + q = mx.random.normal((B, Hq, 1, D)).astype(mx.float16) + k = mx.random.normal((B, Hkv, N, D)).astype(mx.float16) + v = mx.random.normal((B, Hkv, N, D)).astype(mx.float16) + scale = 1.0 / math.sqrt(D) + + ref = _mlx_sdpa(q, k, v, scale=scale) + out = scaled_dot_product_attention(q, k, v, scale=scale) + mx.eval(ref, out) + + diff = mx.abs(ref - out).max().item() + ok = diff < 0.01 + if not ok: + all_correct = False + status = "PASS" if ok else "FAIL" + print(f"{label:<16} {diff:>12.6f} {status:>8}") + + print() + if all_correct: + print(">>> ALL CORRECT <<<") + else: + print(">>> SOME FAILED <<<") + sys.exit(1) + + # ── Performance ───────────────────────────────────────────── + print("\n" + "=" * 72) + print(f"PERFORMANCE (p10 latency, {args.iters} iters)") + print("=" * 72) + print(f"{'Config':<16} {'MLX(us)':>10} {'Cider(us)':>10} {'Speedup':>10} {'Status':>8}") + print("-" * 58) + + wins = ties = losses = 0 + + mx.random.seed(42) + for B, Hq, Hkv, N, label in configs: + q = mx.random.normal((B, Hq, 1, D)).astype(mx.float16) + k = mx.random.normal((B, Hkv, N, D)).astype(mx.float16) + v = mx.random.normal((B, Hkv, N, D)).astype(mx.float16) + scale = 1.0 / math.sqrt(D) + mx.eval(q, k, v) + + mlx_p10, _ = bench_one(_mlx_sdpa, q, k, v, scale, args.warmup, args.iters) + cider_p10, _ = bench_one(scaled_dot_product_attention, q, k, v, scale, args.warmup, args.iters) + + ratio = mlx_p10 / cider_p10 + if ratio >= 1.02: + status = "WIN" + wins += 1 + elif ratio >= 0.98: + status = "TIE" + ties += 1 + else: + status = "LOSS" + losses += 1 + + print(f"{label:<16} {mlx_p10:>10.1f} {cider_p10:>10.1f} {ratio:>9.2f}x {status:>8}") + + print("-" * 58) + print(f"Summary: {wins} wins / {ties} ties / {losses} losses") + print() + + if losses == 0: + print(">>> NO REGRESSIONS <<<") + else: + print(f">>> {losses} REGRESSION(S) DETECTED <<<") + + +if __name__ == "__main__": + main() diff --git a/benchmarks/mlx_native/bench_cider_vs_mlx_int8.py b/benchmarks/mlx_native/bench_cider_vs_mlx_int8.py new file mode 100644 index 0000000..0e0faca --- /dev/null +++ b/benchmarks/mlx_native/bench_cider_vs_mlx_int8.py @@ -0,0 +1,353 @@ +#!/usr/bin/env python3 +""" +Benchmark: Cider W8A8 INT8 matmul vs MLX-style NAX INT8 matmul. + +Both use the same mpp::tensor_ops::matmul2d hardware instruction. +The difference is in data loading/storing patterns. + +We write a standalone INT8 GEMM kernel in MLX Steel NAX style, +compiled via mx.fast.metal_kernel, and compare against cider's kernel. +""" +import sys, os, time +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +import numpy as np +import mlx.core as mx +from cider import perchannel_linear, quantize_weight_int8 +from cider.ops import int8_matmul_int32 + +# ============================================================= +# MLX-style INT8 NAX GEMM kernel (standalone) +# Same tile config as cider: BM=128, BN=128, BK=512, WM=4, WN=4 +# Uses matmul2d - identical HW instruction to cider +# ============================================================= + +MLX_INT8_HEADER = r""" +#include +#include + +using namespace metal; + +constant constexpr int kElemsPerFrag = 8; +constant constexpr int kElemCols = 4; +constant constexpr int kElemRowsJump = 8; + +METAL_FUNC short2 nax_get_coord(ushort lid) { + short qid = short(lid >> 2); + short fm = ((qid & 4) | ((short(lid) >> 1) & 3)); + short fn = ((qid & 2) | (short(lid) & 1)) * 4; + return short2(fn, fm); +} + +METAL_FUNC void load_frag_int8(thread int8_t *dst, + const device int8_t *base, + int stride, short2 sc, + short row_off, short col_off) { + const device int8_t *ptr = base + (sc.y + row_off) * stride + sc.x + col_off; + for (short i = 0; i < 2; i++) { + for (short j = 0; j < kElemCols; j++) { + dst[i * kElemCols + j] = ptr[i * kElemRowsJump * stride + j]; + } + } +} + +METAL_FUNC void store_frag_int32(const thread int32_t *src, + device int32_t *base, + int stride, short2 sc, + short row_off, short col_off, + uint M, uint N, + uint m_base, uint n_base) { + for (short i = 0; i < 2; i++) { + uint r = m_base + row_off + sc.y + i * kElemRowsJump; + for (short j = 0; j < kElemCols; j++) { + uint c = n_base + col_off + sc.x + j; + if (r < M && c < N) { + base[r * N + c] = src[i * kElemCols + j]; + } + } + } +} +""" + +MLX_INT8_SOURCE = r""" + uint M_val = M[0], N_val = N[0], K_val = K[0]; + uint tiles_m_val = tilesm[0], tiles_n_val = tilesn[0]; + + constexpr int BM = 128, BN = 128, BK = 512; + constexpr int WM = 4, WN = 4; + constexpr int SM = BM / WM; + constexpr int SN = BN / WN; + constexpr short TM = SM / 16; + constexpr short TN = SN / 16; + constexpr short TK = 32 / 16; + constexpr short SK = 32; + + uint2 tgid = uint2(threadgroup_position_in_grid.x, threadgroup_position_in_grid.y); + uint sgid = simdgroup_index_in_threadgroup; + uint lid = thread_index_in_simdgroup; + + if (tgid.x >= tiles_n_val || tgid.y >= tiles_m_val) return; + + short2 sc = nax_get_coord(ushort(lid)); + uint sg_row = sgid / WN; + uint sg_col = sgid % WN; + uint m_base = tgid.y * BM + sg_row * SM; + uint n_base = tgid.x * BN + sg_col * SN; + + const device int8_t *sg_A = A + m_base * K_val; + const device int8_t *sg_B = B + n_base; + + constexpr auto desc = mpp::tensor_ops::matmul2d_descriptor( + 16, 32, 16, false, false, true, + mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate); + mpp::tensor_ops::matmul2d gemm_op; + + auto ct_a = gemm_op.get_left_input_cooperative_tensor(); + auto ct_b = gemm_op.get_right_input_cooperative_tensor(); + auto ct_c = gemm_op.get_destination_cooperative_tensor(); + + int32_t c_frags[TM * TN][kElemsPerFrag]; + for (int f = 0; f < TM * TN; f++) + for (int i = 0; i < kElemsPerFrag; i++) + c_frags[f][i] = 0; + + int gemm_k_iters = int(K_val) / BK; + for (int kk0 = 0; kk0 < gemm_k_iters; kk0++) { + threadgroup_barrier(mem_flags::mem_none); + for (int kk1 = 0; kk1 < BK; kk1 += SK) { + int8_t a_frags[TM][TK][kElemsPerFrag]; + int8_t b_frags[TK][TN][kElemsPerFrag]; + volatile int compiler_barrier; + + for (short mm = 0; mm < TM; mm++) + for (short kk = 0; kk < TK; kk++) + load_frag_int8(a_frags[mm][kk], sg_A + kk1, int(K_val), sc, short(mm*16), short(kk*16)); + + for (short kk = 0; kk < TK; kk++) + for (short nn = 0; nn < TN; nn++) + load_frag_int8(b_frags[kk][nn], sg_B + kk1 * N_val, int(N_val), sc, short(kk*16), short(nn*16)); + + for (short mm = 0; mm < TM; mm++) { + for (short nn = 0; nn < TN; nn += 2) { + for (short kk = 0; kk < TK; kk++) { + for (short i = 0; i < kElemsPerFrag; i++) ct_a[i] = a_frags[mm][kk][i]; + for (short i = 0; i < kElemsPerFrag; i++) { + ct_b[i] = b_frags[kk][nn][i]; + ct_b[kElemsPerFrag + i] = b_frags[kk][nn+1][i]; + } + short c0 = mm*TN+nn, c1 = c0+1; + for (short i = 0; i < kElemsPerFrag; i++) { + ct_c[i] = c_frags[c0][i]; + ct_c[kElemsPerFrag+i] = c_frags[c1][i]; + } + gemm_op.run(ct_a, ct_b, ct_c); + for (short i = 0; i < kElemsPerFrag; i++) { + c_frags[c0][i] = ct_c[i]; + c_frags[c1][i] = ct_c[kElemsPerFrag+i]; + } + } + } + } + (void)compiler_barrier; + } + sg_A += BK; + sg_B += BK * N_val; + } + + for (short mm = 0; mm < TM; mm++) + for (short nn = 0; nn < TN; nn++) + store_frag_int32(c_frags[mm*TN+nn], out, int(N_val), sc, short(mm*16), short(nn*16), M_val, N_val, m_base, n_base); +""" + +WARMUP = 20 +REPEAT = 50 + +SHAPES = [ + (16, 4096, 4096), + (64, 4096, 4096), + (128, 4096, 4096), + (256, 4096, 4096), + (512, 4096, 4096), + (1024, 4096, 4096), + (2048, 4096, 4096), + # MLP shapes + (128, 3584, 18944), + (512, 3584, 18944), + (128, 18944, 3584), + (512, 18944, 3584), +] + + +def bench_cider_w8a8(x, w_int8, scale_w): + """Cider full pipeline: quantize_act + int8_matmul + fused_dequant -> FP16""" + for _ in range(WARMUP): + y = perchannel_linear(x, w_int8, scale_w) + mx.eval(y) + times = [] + for _ in range(REPEAT): + t0 = time.perf_counter() + y = perchannel_linear(x, w_int8, scale_w) + mx.eval(y) + times.append(time.perf_counter() - t0) + return np.median(times) * 1000 + + +def bench_cider_int8_raw(A_int8, B_int8): + """Cider raw INT8 matmul: int8 x int8 -> int32 (no quantize, no dequant)""" + for _ in range(WARMUP): + y = int8_matmul_int32(A_int8, B_int8) + mx.eval(y) + times = [] + for _ in range(REPEAT): + t0 = time.perf_counter() + y = int8_matmul_int32(A_int8, B_int8) + mx.eval(y) + times.append(time.perf_counter() - t0) + return np.median(times) * 1000 + + +def bench_mlx_int8(A_int8, B_int8, M, N, K): + """MLX-style standalone INT8 matmul: int8 x int8 -> int32 (raw GEMM only)""" + BM, BN = 128, 128 + tiles_m = (M + BM - 1) // BM + tiles_n = (N + BN - 1) // BN + + M_buf = mx.array([M], dtype=mx.uint32) + N_buf = mx.array([N], dtype=mx.uint32) + K_buf = mx.array([K], dtype=mx.uint32) + tm_buf = mx.array([tiles_m], dtype=mx.uint32) + tn_buf = mx.array([tiles_n], dtype=mx.uint32) + + kernel = mx.fast.metal_kernel( + name='mlx_style_int8_gemm', + input_names=['A', 'B', 'M', 'N', 'K', 'tilesm', 'tilesn'], + output_names=['out'], + source=MLX_INT8_SOURCE, + header=MLX_INT8_HEADER, + ) + + # grid = total threads, NOT threadgroup count + # threadgroup=(512,1,1), so grid_x = tiles_n * 512, grid_y = tiles_m * 1 + grid_x = tiles_n * 512 + grid_y = tiles_m + + for _ in range(WARMUP): + result = kernel( + inputs=[A_int8, B_int8, M_buf, N_buf, K_buf, tm_buf, tn_buf], + output_shapes=[(M, N)], + output_dtypes=[mx.int32], + grid=(grid_x, grid_y, 1), + threadgroup=(512, 1, 1), + ) + mx.eval(result[0]) + + times = [] + for _ in range(REPEAT): + t0 = time.perf_counter() + result = kernel( + inputs=[A_int8, B_int8, M_buf, N_buf, K_buf, tm_buf, tn_buf], + output_shapes=[(M, N)], + output_dtypes=[mx.int32], + grid=(grid_x, grid_y, 1), + threadgroup=(512, 1, 1), + ) + mx.eval(result[0]) + times.append(time.perf_counter() - t0) + return np.median(times) * 1000 + + +def verify_correctness(K, N): + """Quick correctness check: MLX-style int8 matmul vs numpy.""" + M = 16 + np.random.seed(123) + a_np = np.random.randint(-5, 5, (M, K), dtype=np.int8) + b_np = np.random.randint(-5, 5, (K, N), dtype=np.int8) + ref = a_np.astype(np.int32) @ b_np.astype(np.int32) + + A_int8 = mx.array(a_np) + B_int8 = mx.array(b_np) + + BM, BN = 128, 128 + tiles_m = (M + BM - 1) // BM + tiles_n = (N + BN - 1) // BN + + kernel = mx.fast.metal_kernel( + name='mlx_style_int8_gemm', + input_names=['A', 'B', 'M', 'N', 'K', 'tilesm', 'tilesn'], + output_names=['out'], + source=MLX_INT8_SOURCE, + header=MLX_INT8_HEADER, + ) + + result = kernel( + inputs=[A_int8, B_int8, + mx.array([M], dtype=mx.uint32), + mx.array([N], dtype=mx.uint32), + mx.array([K], dtype=mx.uint32), + mx.array([tiles_m], dtype=mx.uint32), + mx.array([tiles_n], dtype=mx.uint32)], + output_shapes=[(M, N)], + output_dtypes=[mx.int32], + grid=(tiles_n * 512, tiles_m, 1), + threadgroup=(512, 1, 1), + ) + mx.eval(result[0]) + got = np.array(result[0]) + + max_diff = np.max(np.abs(ref - got)) + match = np.array_equal(ref, got) + print(f"Correctness check ({M}x{K} @ {K}x{N}): max_diff={max_diff}, exact_match={match}") + if not match: + print(f" ref[0,:8] = {ref[0,:8]}") + print(f" got[0,:8] = {got[0,:8]}") + return match + + +def main(): + print("=" * 70) + print("Cider W8A8 (full pipeline) vs MLX-style INT8 GEMM (raw matmul)") + print("=" * 70) + print(f"Warmup={WARMUP}, Repeat={REPEAT}") + print() + print("NOTE: Cider includes quantize_act + INT8 matmul + fused dequant") + print(" MLX-style is raw INT8 matmul only (no quantize/dequant)") + print(" So MLX-style SHOULD be faster (does less work)") + print() + + # Verify correctness first + print("--- Correctness Check ---") + ok = verify_correctness(4096, 4096) + if not ok: + print("WARNING: correctness check failed! Results may be invalid.") + print() + + print(f"{'M':>5s} {'K':>6s} {'N':>6s} | {'Cider-full':>10s} | {'Cider-raw':>10s} | {'MLX-INT8':>9s} | {'full/raw':>8s} | {'raw/MLX':>7s}") + print("-" * 80) + + prev_kn = None + for M, K, N in SHAPES: + if (K, N) != prev_kn: + np.random.seed(42) + w_fp32 = np.random.randn(K, N).astype(np.float32) + w_int8_np, scale_w_np = quantize_weight_int8(w_fp32) + w_int8 = mx.array(w_int8_np) + scale_w = mx.array(scale_w_np) + B_int8 = mx.array(w_int8_np) + mx.eval(w_int8, scale_w, B_int8) + prev_kn = (K, N) + + x_fp16 = mx.random.normal((M, K)).astype(mx.float16) + A_int8 = mx.array(np.random.randint(-127, 127, (M, K), dtype=np.int8)) + mx.eval(x_fp16, A_int8) + + t_full = bench_cider_w8a8(x_fp16, w_int8, scale_w) + t_raw = bench_cider_int8_raw(A_int8, B_int8) + t_mlx = bench_mlx_int8(A_int8, B_int8, M, N, K) + + r_full_raw = t_full / t_raw if t_raw > 0 else 0 + r_raw_mlx = t_raw / t_mlx if t_mlx > 0 else 0 + print(f"{M:>5d} {K:>6d} {N:>6d} | {t_full:>8.2f}ms | {t_raw:>8.2f}ms | {t_mlx:>7.2f}ms | {r_full_raw:>6.2f}x | {r_raw_mlx:>5.2f}x") + + +if __name__ == "__main__": + main() diff --git a/benchmarks/mlx_native/cider_vs_mlx_int8.md b/benchmarks/mlx_native/cider_vs_mlx_int8.md new file mode 100644 index 0000000..7622218 --- /dev/null +++ b/benchmarks/mlx_native/cider_vs_mlx_int8.md @@ -0,0 +1,117 @@ +# Cider W8A8 Kernel vs MLX INT8 GEMM: A Transparent Comparison + + +MLX's Steel NAX GEMM templates appear generic enough at the shader level to support `` instantiations. In our standalone test, an MLX-style INT8×INT8→INT32 kernel can be built on the same `mpp::tensor_ops::matmul2d` TensorOps instruction used by Cider. At the raw matmul level, both kernels achieve **comparable throughput** (within 2–5% across all tested sizes). + +Cider's contribution is **not** a novel matmul kernel. It is an end-to-end W8A8 inference pipeline for Apple Silicon, combining online activation quantization, INT8 TensorOps matmul, and dequantization in a usable MLX-integrated path. + +**Apple M5 introduced cooperative_tensor INT8 TensorOps (matmul2d 16×32×16), delivering 2× the TOPS of FP16 — but no ML framework on macOS uses them.** In the MLX quantization path we examined (v0.31), quantized weights are dequantized for computation and the corresponding matmul path remains FP16-based rather than end-to-end INT8. This means prefill (the compute-bound phase of LLM inference) gets zero benefit from INT8 hardware. W8A8 is the only quantization scheme that accelerates both compute and memory: activations and weights are both INT8, so the matmul runs on the faster INT8 datapath. Cider provides the missing pipeline — online activation quantization, INT8 TensorOps matmul, and fused dequantization — to unlock this dormant hardware capability. In our current experiments on Qwen3-VL-2B (M5 Pro), Cider achieves 1.15–1.21× prefill speedup over a W8A16 baseline, while preserving comparable language-model perplexity in our Llama-3-8B check (PPL Δ < 0.01). + + +## Background: What MLX Does Today + +In the MLX v0.31 code paths we examined, we did not find a public end-to-end INT8 TensorOps inference path. + +More specifically: + +- **Standard matmul** (`steel_gemm_fused_nax.metal`): Instantiated for `float16`, `bfloat16`, `float32` only. No `int8` instantiation exists. +- **Quantized matmul** (`quantized_nax.h`): Dequantizes packed int{2,3,4,5,6,8} weights → FP16, then uses **FP16 TensorOps**. This is weight-only quantization — activations stay in FP16. + +This does **not** mean INT8 TensorOps are impossible in MLX. At the shader level, the underlying GEMM template structure appears flexible enough to support INT8 instantiations. The missing piece is a complete and exposed W8A8 execution pipeline. + + +## The Benchmark + +We wrote a standalone INT8 GEMM kernel in MLX's NAX style, compiled via `mx.fast.metal_kernel`, using the same: +- `matmul2d` TensorOps instruction +- BM=128, BN=128, BK=512 tile configuration +- NAXFrag cooperative tensor layout +- Device memory direct load (no threadgroup staging) + +This is **not** a hypothetical comparison — the MLX-style kernel compiles, runs, and produces bit-exact results (verified: `max_diff=0`). + +### Three-Way Comparison (Apple M5 Pro, Warmup=30, Repeat=100, 3 runs best-of) + +**Square shapes (K=4096, N=4096)**: + +| M | Cider Full Pipeline | Cider Raw INT8 | MLX-style INT8 | Full/Raw | Raw/MLX | +|---|-------------------|---------------|----------------|----------|---------| +| 16 | 0.24ms | 0.21ms | 0.25ms | 1.12x | 0.84x | +| 64 | 0.25ms | 0.24ms | 0.26ms | 1.04x | 0.95x | +| 128 | 0.28ms | 0.27ms | 0.26ms | 1.03x | 1.02x | +| 256 | 0.40ms | 0.40ms | 0.37ms | 1.01x | 1.08x | +| 512 | 0.59ms | 0.54ms | 0.52ms | 1.09x | 1.04x | +| 1024 | 0.94ms | 0.90ms | 0.86ms | 1.05x | 1.05x | +| 2048 | 1.68ms | 1.55ms | 1.49ms | 1.08x | 1.04x | + +**MLP shapes (Llama-3 8B dimensions)**: + +| M | K | N | Cider Full | Cider Raw | MLX-style | Full/Raw | Raw/MLX | +|---|---|---|-----------|-----------|-----------|----------|---------| +| 128 | 3584 | 18944 | 0.73ms | 0.74ms | 0.70ms | 0.99x | 1.06x | +| 512 | 3584 | 18944 | 1.66ms | 1.65ms | 1.55ms | 1.00x | 1.07x | +| 128 | 18944 | 3584 | 0.83ms | 0.78ms | 0.75ms | 1.06x | 1.05x | +| 512 | 18944 | 3584 | 1.87ms | 1.81ms | 1.68ms | 1.04x | 1.07x | + +Where: +- **Cider Full Pipeline** = `perchannel_linear()`: FP16→INT8 quantize + INT8 matmul + fused INT32→FP16 dequant +- **Cider Raw INT8** = `int8_matmul_int32()`: Pure INT8×INT8→INT32 (no quantize, no dequant) +- **MLX-style INT8** = Standalone NAX kernel via `mx.fast.metal_kernel`: Same pure INT8→INT32 + +## What the Numbers Suggest + +### 1. Raw INT8 TensorOps throughput is comparable + +Across the tested shapes, Cider Raw INT8 and the standalone MLX-style INT8 kernel deliver similar throughput. This is consistent with the fact that both rely on the same underlying `matmul2d` TensorOps instruction with similar tile configurations. + + + +### 2. The main cost of a usable W8A8 path is pipeline integration + +Compared with raw INT8 matmul, the full Cider path additionally includes: +- activation quantization, +- scale handling, +- and dequantization back to the desired output dtype. + +These steps introduce overhead, but they are also what make W8A8 usable in practice. + +### 3. Cider's main value is exposing an end-to-end W8A8 path +Our results suggest that the key contribution of Cider is not a claim of unique raw GEMM throughput, but the availability of an integrated W8A8 execution path on Apple Silicon. + + +## Why Cider Uses a Custom Kernel Path + +If a standalone MLX-style INT8 kernel can achieve similar raw throughput, why does Cider still implement its own kernel path? + + +### 1. To support dequantization as part of the usable W8A8 path + +A practical W8A8 inference path needs more than raw INT32 accumulation. It also needs output scaling and dtype conversion in a form that fits the surrounding runtime. + + +### 2. To keep quantization and INT8 matmul in one integrated execution path + +Cider packages activation quantization and INT8 matmul into a custom MLX primitive, which is more practical than treating them as disconnected experimental kernels. + + +### 3. To bypass the current lack of a public end-to-end INT8 route in standard MLX dispatch + +Even if INT8 kernels are possible at the shader level, they are not currently exposed as a standard end-to-end W8A8 inference path in the MLX stack we examined. + + +## Reproducing + +```bash +cd /path/to/cider +python benchmarks/mlx_native/bench_cider_vs_mlx_int8.py +``` + +Requires: Apple M5+, MLX ≥ 0.31, Cider installed. + +The MLX-style kernel is compiled at runtime via `mx.fast.metal_kernel` with Metal 4 + MPP headers. No Xcode installation required. + +## Notes + +- All statements in this document are scoped to the MLX version and code paths we examined during this comparison. +- The standalone MLX-style INT8 kernel is a raw-kernel benchmark, not a full MLX-native end-to-end W8A8 inference pipeline. +- Cider Full Pipeline results include activation quantization and output dequantization overhead, while raw INT8 kernel results do not. \ No newline at end of file diff --git a/benchmarks/test_bitexact.py b/benchmarks/test_bitexact.py new file mode 100644 index 0000000..6c2a69c --- /dev/null +++ b/benchmarks/test_bitexact.py @@ -0,0 +1,83 @@ +#!/usr/bin/env python3 +"""Bit-exact correctness test for INT8×INT8→INT32 TensorOps kernel. + +Pure integer matmul must produce EXACT results — no floating-point +tolerance, no cosine similarity. Every element must match bit-for-bit. + +Usage: + cd /path/to/cider + python tests/test_bitexact.py +""" +import sys, os +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +import numpy as np +import mlx.core as mx +from cider import int8_matmul_int32 + +# ── Test shapes ───────────────────────────────────────────────── +SHAPES = [ + # (M, K, N) + (1, 2, 2), # minimal + (2, 4, 2), # tiny + (1, 128, 128), # single row + (1, 4096, 4096), # decode-size + (4, 2048, 2048), # small batch + (16, 4096, 4096), # medium + (32, 4096, 4096), # tile boundary (small→large) + (64, 4096, 4096), # large tile + (128, 4096, 4096), # full large tile + (16, 4096, 8192), # wide N + (16, 8192, 4096), # wide K + (64, 2048, 4096), # mixed + (7, 513, 1025), # non-aligned (prime-ish) + (33, 4097, 4095), # non-aligned near 4096 +] + + +def test_bitexact(): + print("=" * 70) + print("INT8×INT8→INT32 Bit-Exact Tests") + print("=" * 70) + passed, failed = 0, 0 + + for M, K, N in SHAPES: + np.random.seed(42) + a_np = np.random.randint(-128, 128, (M, K), dtype=np.int8) + b_np = np.random.randint(-128, 128, (N, K), dtype=np.int8) + + # GPU kernel + c_gpu = int8_matmul_int32(mx.array(a_np), mx.array(b_np)) + mx.eval(c_gpu) + c_gpu_np = np.array(c_gpu) + + # Numpy reference (int64 to avoid overflow in accumulation) + c_ref = (a_np.astype(np.int64) @ b_np.astype(np.int64).T).astype(np.int32) + + # Bit-exact check + match = np.array_equal(c_gpu_np, c_ref) + if match: + passed += 1 + tag = "PASS" + else: + failed += 1 + diff = (c_gpu_np != c_ref) + n_diff = diff.sum() + # Show first mismatch + idx = np.argwhere(diff)[0] + tag = f"FAIL ({n_diff}/{M*N} mismatches, first@{tuple(idx)}: gpu={c_gpu_np[tuple(idx)]} ref={c_ref[tuple(idx)]})" + + print(f" M={M:>4d} K={K:>4d} N={N:>4d} [{tag}]") + + print(f"\nBit-exact: {passed}/{passed+failed} passed") + return failed == 0 + + +if __name__ == "__main__": + ok = test_bitexact() + print("\n" + "=" * 70) + if ok: + print("ALL BIT-EXACT TESTS PASSED ✓") + else: + print("SOME TESTS FAILED") + sys.exit(1) diff --git a/cider/__init__.py b/cider/__init__.py new file mode 100644 index 0000000..e741889 --- /dev/null +++ b/cider/__init__.py @@ -0,0 +1,56 @@ +"""cider — INT8 TensorOps quantized matmul + optimized SDPA for Apple Silicon. + +Quick start (attention acceleration): + import cider + cider.patch_sdpa() # One line. mlx_lm/mlx_vlm auto-accelerated. + cider.autotune_sdpa() # Optional: sweep blocks for best perf on your GPU. + +Quick start (quantization): + from cider import convert_model + model, proc = load("model_path") + convert_model(model) +""" + +__version__ = "0.8.0" + +# ── Attention acceleration (always available) ─────────────────── +from .attention import patch_sdpa, unpatch_sdpa, autotune_sdpa + +# ── Quantization (M5+ only) ──────────────────────────────────── +from .ops import is_available + +if is_available(): + from .convert import convert_model + from .nn import CiderLinear, set_mode, get_mode, W4A8Linear, W8A8Linear + from .ops import ( + perchannel_linear, w4a8_linear, pergroup_linear, + int8_matmul_int32, quantize_weight_int8, pack_weight_int4, kernel_dir, + ) + + __all__ = [ + "patch_sdpa", "unpatch_sdpa", "autotune_sdpa", + "convert_model", "set_mode", "get_mode", + "CiderLinear", "W8A8Linear", "W4A8Linear", + "perchannel_linear", "w4a8_linear", "pergroup_linear", + "int8_matmul_int32", "quantize_weight_int8", "pack_weight_int4", + "is_available", "kernel_dir", + ] +else: + def convert_model(*args, **kwargs): + import warnings + warnings.warn( + "cider.convert_model() is a no-op: INT8 TensorOps require Apple M5+. " + "Model will use standard MLX inference.", + RuntimeWarning, stacklevel=2, + ) + + def set_mode(*args, **kwargs): + pass + + def get_mode(): + return "unavailable" + + __all__ = [ + "patch_sdpa", "unpatch_sdpa", "autotune_sdpa", + "convert_model", "set_mode", "get_mode", "is_available", + ] diff --git a/cider/attention/__init__.py b/cider/attention/__init__.py new file mode 100644 index 0000000..2d5130f --- /dev/null +++ b/cider/attention/__init__.py @@ -0,0 +1,23 @@ +"""cider.attention — Optimized SDPA kernels for Apple Silicon. + +Usage: + import cider + cider.patch_sdpa() # monkey-patch mx.fast.scaled_dot_product_attention + cider.autotune_sdpa() # optional: sweep blocks for best perf +""" + +from .sdpa import ( + scaled_dot_product_attention, + load_autotune, + autotune_sdpa, + patch_sdpa, + unpatch_sdpa, +) + +__all__ = [ + "scaled_dot_product_attention", + "load_autotune", + "autotune_sdpa", + "patch_sdpa", + "unpatch_sdpa", +] diff --git a/cider/attention/sdpa.py b/cider/attention/sdpa.py new file mode 100644 index 0000000..3f95f8f --- /dev/null +++ b/cider/attention/sdpa.py @@ -0,0 +1,350 @@ +"""cider.attention.sdpa — Compiled SDPA with monkey-patch support. + +Uses precompiled Metal kernels via C++ Custom Primitive (zero JIT overhead). +Falls back to MLX default for unsupported cases. +""" + +import math +import json +import os +from typing import Optional + +import mlx.core as mx + + +# ── Kernel directory (where .metal/.h files live) ─────────────── +def _get_kernel_dir() -> str: + return os.path.join(os.path.dirname(os.path.dirname(__file__)), "kernels") + + +# ── C++ primitive import ──────────────────────────────────────── +_prim = None + +def _get_prim(): + global _prim + if _prim is None: + try: + from cider.lib import _cider_prim + _prim = _cider_prim + except ImportError: + raise RuntimeError( + "cider C++ extension not built. " + "Run: pip install -e . (from cider root)" + ) + return _prim + + +# ── GPU arch ──────────────────────────────────────────────────── +_gpu_arch = None + +def _get_gpu_arch(): + global _gpu_arch + if _gpu_arch is None: + info = mx.device_info() if hasattr(mx, 'device_info') else mx.metal.device_info() + _gpu_arch = info.get("architecture", "unknown") + return _gpu_arch + + +# ── Routing ───────────────────────────────────────────────────── +def _use_cider(N: int, gqa_factor: int) -> bool: + """Decide whether Cider kernel beats MLX default. + + Compiled primitive has near-zero dispatch overhead (unlike JIT). + Our kernel uses contiguous chunk + TILE=4 which benefits all configs. + Always route through Cider for decode (Q_seq=1) — the calling code + already handles mask/prefill fallback before reaching here. + """ + return True + + +def _use_2pass(N: int, gqa_factor: int) -> bool: + return N >= 1024 + + +def _select_blocks(N: int, gqa_factor: int, H_q: int, arch: str) -> int: + cached = _autotune_lookup(N, gqa_factor, H_q, arch) + if cached is not None: + return cached + if gqa_factor >= 8: + if N >= 8192: + return 128 + elif N >= 4096: + return 64 + else: + return 32 + elif gqa_factor >= 4: + if N >= 8192: + return 64 + else: + return 32 + else: + if N >= 16384: + return 128 + elif N >= 4096: + return 64 + else: + return 32 + + +# ── Core SDPA ─────────────────────────────────────────────────── +def scaled_dot_product_attention( + q: mx.array, + k: mx.array, + v: mx.array, + *, + scale: Optional[float] = None, + mask: Optional[mx.array] = None, + **kwargs, +) -> mx.array: + """Optimized SDPA for decode (q_seq_len=1). + + Drop-in replacement for mx.fast.scaled_dot_product_attention. + Falls back to MLX for unsupported cases (mask, non-decode, prefill). + Accepts **kwargs for forward-compat with new MLX args (e.g. sinks). + """ + if mask is not None: + return _mlx_sdpa(q, k, v, scale=scale, mask=mask, **kwargs) + + B, H_q, Q_seq, D = q.shape + _, H_kv, N, _ = k.shape + + if Q_seq != 1: + return _mlx_sdpa(q, k, v, scale=scale, **kwargs) + + if scale is None: + scale = 1.0 / math.sqrt(D) + + gqa_factor = H_q // H_kv + + # Fall back to MLX when Cider doesn't help + if not _use_cider(N, gqa_factor): + return _mlx_sdpa(q, k, v, scale=scale, **kwargs) + + prim = _get_prim() + kernel_dir = _get_kernel_dir() + + if not _use_2pass(N, gqa_factor): + # ── 1-pass ── + # Flatten to [B*H_q, D] for queries, [B*H_kv, N, D] for K/V + q_flat = q.reshape(B * H_q, D) + k_cont = k.reshape(B * H_kv, N, D) + v_cont = v.reshape(B * H_kv, N, D) + + out_flat = prim.cider_sdpa_1pass( + q_flat, k_cont, v_cont, + gqa_factor, scale, kernel_dir, + ) + return out_flat.reshape(B, H_q, 1, D) + + else: + # ── 2-pass ── + arch = _get_gpu_arch() + blocks = _select_blocks(N, gqa_factor, H_q, arch) + + out = prim.cider_sdpa_2pass( + q, k, v, + gqa_factor, blocks, scale, kernel_dir, + ) + return out + + +# ── Monkey-patch ──────────────────────────────────────────────── +_mlx_sdpa = mx.fast.scaled_dot_product_attention # save original +_patched = False + + +def patch_sdpa(verbose: bool = True): + """Monkey-patch mx.fast.scaled_dot_product_attention with Cider's optimized version. + + Usage: + import cider + cider.patch_sdpa() + # Now all calls to mx.fast.scaled_dot_product_attention use Cider kernels. + # mlx_lm, mlx_vlm etc. benefit automatically. + """ + global _patched + if _patched: + if verbose: + print("[cider] SDPA already patched.") + return + + # Load AutoTune cache if available + n = load_autotune() + + mx.fast.scaled_dot_product_attention = scaled_dot_product_attention + _patched = True + + if verbose: + arch = _get_gpu_arch() + tune_msg = f", {n} tuned configs loaded" if n > 0 else "" + print(f"[cider] SDPA patched — compiled kernels, {arch}{tune_msg}") + + +def unpatch_sdpa(verbose: bool = True): + """Restore original mx.fast.scaled_dot_product_attention.""" + global _patched + if not _patched: + return + mx.fast.scaled_dot_product_attention = _mlx_sdpa + _patched = False + if verbose: + print("[cider] SDPA unpatched — restored MLX default.") + + +# ── AutoTune ──────────────────────────────────────────────────── +_autotune_cache = {} +_AUTOTUNE_FILE = os.path.expanduser("~/.cider_sdpa_tune.json") +_BLOCKS_CANDIDATES = [32, 64, 128] + + +def _autotune_lookup(N: int, gqa_factor: int, H_q: int, arch: str) -> Optional[int]: + """Look up cached blocks value. Tries exact H_q match first, then generic.""" + key = f"{arch}_gqa{gqa_factor}_H{H_q}_N{N}" + val = _autotune_cache.get(key) + if val is not None: + return val + # Fallback: try without H_q (legacy cache or generic entry) + key_generic = f"{arch}_gqa{gqa_factor}_N{N}" + return _autotune_cache.get(key_generic) + + +def load_autotune() -> int: + """Load cached AutoTune results from disk.""" + if os.path.exists(_AUTOTUNE_FILE): + with open(_AUTOTUNE_FILE) as f: + data = json.load(f) + _autotune_cache.update(data) + return len(data) + return 0 + + +def save_autotune(): + """Save AutoTune cache to disk.""" + with open(_AUTOTUNE_FILE, 'w') as f: + json.dump(_autotune_cache, f, indent=2) + + +def _bench_one(q, k, v, scale, blocks, warmup=10, iters=50): + """Time a single (blocks) config. Returns median microseconds.""" + import time + prim = _get_prim() + kernel_dir = _get_kernel_dir() + gqa_factor = q.shape[1] // k.shape[1] + + # Warmup + for _ in range(warmup): + out = prim.cider_sdpa_2pass( + q, k, v, gqa_factor, blocks, scale, kernel_dir) + mx.eval(out) + + # Timed runs + times = [] + for _ in range(iters): + t0 = time.perf_counter() + out = prim.cider_sdpa_2pass( + q, k, v, gqa_factor, blocks, scale, kernel_dir) + mx.eval(out) + times.append((time.perf_counter() - t0) * 1e6) + + times.sort() + # Return p10 (robust against outliers) + idx = max(0, len(times) // 10) + return times[idx] + + +def autotune_sdpa( + gqa_factors: Optional[list] = None, + seq_lens: Optional[list] = None, + D: int = 128, + warmup: int = 10, + iters: int = 50, + verbose: bool = True, +) -> dict: + """Run AutoTune sweep for 2-pass SDPA blocks parameter. + + Sweeps blocks in [32, 64, 128] for each (gqa_factor, N) config. + Results cached to ~/.cider_sdpa_tune.json. + + Args: + gqa_factors: list of GQA ratios to tune (default: [1,2,4,8]) + seq_lens: list of sequence lengths (default: [1024,2048,4096,8192,16384,32768]) + D: head dimension (default: 128) + warmup: warmup iterations per config + iters: benchmark iterations per config + verbose: print results + + Returns: + dict mapping config keys to best blocks values + """ + if gqa_factors is None: + gqa_factors = [1, 2, 4, 8] + if seq_lens is None: + seq_lens = [1024, 2048, 4096, 8192, 16384, 32768] + + # Typical H_q values across popular models: + # 32: Qwen2.5-7B, LLaMA-3-8B, Gemma-2-9B, SmolLM + # 96: Qwen3-32B, Qwen2.5-72B + # 112: LLaMA-3.1-70B + _TYPICAL_HQ = [32, 96, 112] + + arch = _get_gpu_arch() + results = {} + + if verbose: + print(f"[cider] AutoTune starting — arch={arch}, D={D}") + print(f" GQA factors: {gqa_factors}") + print(f" H_q values: {_TYPICAL_HQ}") + print(f" Seq lengths: {seq_lens}") + print(f" Blocks candidates: {_BLOCKS_CANDIDATES}") + print() + + mx.random.seed(42) + + for gqa in gqa_factors: + for H_q in _TYPICAL_HQ: + H_kv = H_q // gqa if gqa > 1 else H_q + if H_kv < 1: + continue # Skip invalid combos (e.g., H_q=32, GQA=96) + if H_q % gqa != 0 and gqa > 1: + continue # H_q must be divisible by gqa_factor + B = 1 + scale = 1.0 / math.sqrt(D) + + for N in seq_lens: + if N < 1024: + continue # 1-pass, no blocks parameter + + q = mx.random.normal((B, H_q, 1, D)).astype(mx.float16) + k = mx.random.normal((B, H_kv, N, D)).astype(mx.float16) + v = mx.random.normal((B, H_kv, N, D)).astype(mx.float16) + mx.eval(q, k, v) + + best_blocks = _BLOCKS_CANDIDATES[0] + best_time = float('inf') + timings = {} + + for blocks in _BLOCKS_CANDIDATES: + t = _bench_one(q, k, v, scale, blocks, warmup, iters) + timings[blocks] = t + if t < best_time: + best_time = t + best_blocks = blocks + + key = f"{arch}_gqa{gqa}_H{H_q}_N{N}" + _autotune_cache[key] = best_blocks + results[key] = best_blocks + + if verbose: + details = " | ".join( + f"b{b}={timings[b]:.0f}us{'*' if b == best_blocks else ''}" + for b in _BLOCKS_CANDIDATES + ) + print(f" GQA{gqa} H{H_q:<3} N={N:>6}: {details} -> blocks={best_blocks}") + + # Save to disk + save_autotune() + + if verbose: + print(f"\n[cider] AutoTune done — {len(results)} configs saved to {_AUTOTUNE_FILE}") + + return results diff --git a/cider/convert.py b/cider/convert.py new file mode 100644 index 0000000..7596020 --- /dev/null +++ b/cider/convert.py @@ -0,0 +1,100 @@ +"""cider.convert — One-line model conversion for W8A8 prefill acceleration. + +Usage: + from cider import convert_model, set_mode + + model, proc = load("model_path") + convert_model(model) # Patch all Linear layers in-place + +works with standard MLX LLM/VLM architectures — tested on Qwen, Llama, Qwen3-VL +Supports float16 and bfloat16 models automatically. +""" + +import gc +import time + +import mlx.nn as nn + +from . import ops +from .nn import CiderLinear + +# Re-export mode control from nn (single source of truth) + + +# ── Model conversion ─────────────────────────────────────────── +_TARGET_TYPES = (nn.Linear, nn.QuantizedLinear) + + +def _convert_children(module, counter, clip_percentile=None): + """Walk module.children(), replace Linear/QuantizedLinear in-place.""" + for name, child in module.children().items(): + if isinstance(child, _TARGET_TYPES): + setattr(module, name, CiderLinear.from_float(child, clip_percentile=clip_percentile)) + counter[0] += 1 + if counter[0] % 28 == 0: + gc.collect() + elif isinstance(child, list): + for i, item in enumerate(child): + if isinstance(item, _TARGET_TYPES): + child[i] = CiderLinear.from_float(item, clip_percentile=clip_percentile) + counter[0] += 1 + if counter[0] % 28 == 0: + gc.collect() + elif isinstance(item, nn.Module): + _convert_children(item, counter, clip_percentile=clip_percentile) + elif isinstance(child, nn.Module): + _convert_children(child, counter, clip_percentile=clip_percentile) + # Skip dict/other non-Module children (e.g. rope_scaling) + + +def convert_model( + model: nn.Module, + *, + clip_percentile: float = None, + verbose: bool = True, +) -> dict: + """Convert all Linear/QuantizedLinear layers to CiderLinear (in-place). + + After conversion: + - set_mode("prefill") → all linears use W8A8 INT8 TensorOps + - set_mode("decode") → all linears use original weights (no overhead) + + Args: + model: Any MLX nn.Module (Qwen3-VL, Llama, Mistral, etc.). + verbose: Print conversion summary. + + Returns: + dict with stats: n_converted, elapsed_s. + + Example: + from cider import convert_model, set_mode + + model, proc = load("model_path") + stats = convert_model(model) + + set_mode("prefill") + # ... run prefill (W8A8, ~15-19% faster) ... + + set_mode("decode") + # ... run decode (original weights, optimal for single token) ... + """ + if not ops.is_available(): + raise RuntimeError( + "W8A8 INT8 TensorOps not available. Requires Apple M5+ with Metal 4." + ) + + t0 = time.perf_counter() + counter = [0] + _convert_children(model, counter, clip_percentile=clip_percentile) + + elapsed = time.perf_counter() - t0 + n = counter[0] + stats = {"n_converted": n, "elapsed_s": elapsed} + + if verbose: + print( + f"[cider] Converted {n} layers to CiderLinear " + f"in {elapsed:.1f}s" + ) + + return stats diff --git a/cider/fused_hybrid.py b/cider/fused_hybrid.py new file mode 100644 index 0000000..ac7b09f --- /dev/null +++ b/cider/fused_hybrid.py @@ -0,0 +1,296 @@ +""" +mlx_w8a8.fused_hybrid — Fused W8A8 Hybrid for Qwen3-VL prefill acceleration. + +QKV fuse + gate/up fuse → fewer kernel dispatches → better GPU occupancy. +Prefill: W8A8 INT8 TensorOps with fused projections. +Decode: original QuantizedLinear (W8A16) unchanged. + +Usage: + from mlx_w8a8.fused_hybrid import convert_model_fused, set_mode + + model, proc = load(model_path) # Load W8A16 model + convert_model_fused(model) # Patch in-place + set_mode("prefill") # Before prefill chunks + # ... run prefill ... + set_mode("decode") # Before decode loop + # ... run decode (uses original W8A16) ... +""" + +import time, gc +import numpy as np +import mlx.core as mx +import mlx.nn as nn + +# ── Chip capability detection ─────────────────────────────────── +def is_w8a8_available() -> bool: + """Check if the current hardware supports W8A8 INT8 TensorOps. + + Requires Apple M5 or later (Metal 4.0 MPP cooperative_tensor + matmul2d). + M4 and earlier do NOT support INT8 TensorOps — returns False. + + Returns: + True if W8A8 acceleration is available and the extension loads correctly. + """ + import subprocess + try: + chip = subprocess.run( + ["sysctl", "-n", "machdep.cpu.brand_string"], + capture_output=True, text=True, timeout=5 + ).stdout.strip() + except Exception: + return False + + # M5, M5 Pro, M5 Max, M5 Ultra, and future M6+ all support Metal 4 + # M4 and below do NOT have INT8 TensorOps + import re + m = re.match(r"Apple M(\d+)", chip) + if not m or int(m.group(1)) < 5: + return False + + # Verify extension actually loads + try: + _ensure_ext() + return True + except Exception: + return False + + +# Lazy imports for extension +_w8a8_prim = None +_KERNEL_DIR = None + + +def _quantize_per_channel_np(w_f32): + """Per-channel INT8 quantization. w: [K, N] float32 → (int8 [K,N], scale [N]).""" + col_max = np.max(np.abs(w_f32), axis=0) # [N] + scale = col_max / 127.0 + scale = np.where(scale == 0, 1.0, scale) # avoid div-by-zero + w_int8 = np.clip(np.round(w_f32 / scale[np.newaxis, :]), -128, 127).astype(np.int8) + return w_int8, scale.astype(np.float32) + + +def _ensure_ext(): + global _w8a8_prim, _KERNEL_DIR + if _w8a8_prim is None: + import os, sys + ext_lib = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname( + os.path.abspath(__file__)))), "ext", "lib") + if ext_lib not in sys.path: + sys.path.insert(0, ext_lib) + import _w8a8_prim as _prim + _w8a8_prim = _prim + _KERNEL_DIR = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname( + os.path.abspath(__file__)))), "kernels") + + +# ── Global mode switch ────────────────────────────────────────── +_hybrid_mode = "decode" + + +def set_mode(mode: str): + """Set the hybrid mode: 'prefill' (W8A8 fused) or 'decode' (W8A16 original).""" + global _hybrid_mode + assert mode in ("prefill", "decode"), f"Invalid mode: {mode}" + _hybrid_mode = mode + + +def get_mode() -> str: + return _hybrid_mode + + +# ── Helpers ───────────────────────────────────────────────────── +def _dequant_ql(ql): + """Dequantize a QuantizedLinear → FP16 weight [out_features, in_features].""" + w = mx.dequantize( + ql.weight, ql.scales, + ql.biases if hasattr(ql, 'biases') and ql.biases is not None else None, + ql.group_size, ql.bits, + ) + mx.eval(w) + return w + + +def _make_w8a8(w_fp16_out_in): + """Convert [out, in] FP16 weight → (INT8 [in, out], FP32 scales [out]).""" + _ensure_ext() + w_t = mx.transpose(w_fp16_out_in) + mx.eval(w_t) + w_np = np.array(w_t, dtype=np.float32) + w_int8, scales = _quantize_per_channel_np(w_np) + result = (mx.array(w_int8), mx.array(scales)) + mx.eval(*result) + return result + + +def _w8a8_gemm(x, weight_int8, scales): + """Dispatch W8A8 GEMM via mlx primitive. x: [..., K], weight: [K, N].""" + _ensure_ext() + orig_shape = x.shape + K = orig_shape[-1] + N = weight_int8.shape[1] + if x.ndim > 2: + x_2d = mx.reshape(x, (-1, K)) + else: + x_2d = x + y_2d = _w8a8_prim.perchannel_linear(x_2d, weight_int8, scales, _KERNEL_DIR) + if x.ndim > 2: + return mx.reshape(y_2d, orig_shape[:-1] + (N,)) + return y_2d + + +# ── Fused Attention forward ──────────────────────────────────── +def _fused_attention_call(self, x, mask=None, cache=None, position_ids=None): + """Attention forward: fused QKV GEMM in prefill, original in decode.""" + B, L, D = x.shape + + if _hybrid_mode == "prefill" and hasattr(self, '_fused_qkv_w'): + # Single fused QKV GEMM: [B,L,dim] → [B,L, n_q+n_k+n_v] + qkv = _w8a8_gemm(x, self._fused_qkv_w, self._fused_qkv_s) + n_q, n_k, n_v = self._qkv_split + queries = qkv[..., :n_q] + keys = qkv[..., n_q:n_q + n_k] + values = qkv[..., n_q + n_k:] + else: + queries = self.q_proj(x) + keys = self.k_proj(x) + values = self.v_proj(x) + + queries = self.q_norm( + queries.reshape(B, L, self.n_heads, self.head_dim) + ).transpose(0, 2, 1, 3) + keys = self.k_norm( + keys.reshape(B, L, self.n_kv_heads, self.head_dim) + ).transpose(0, 2, 1, 3) + values = values.reshape(B, L, self.n_kv_heads, self.head_dim).transpose( + 0, 2, 1, 3 + ) + + kv_seq_len = keys.shape[-2] + if position_ids is None: + kv_seq_len += cache.offset + 1 + position_ids = mx.arange(cache.offset, cache.offset + L) + position_ids = mx.expand_dims(position_ids, axis=0) + position_ids = mx.tile(position_ids, (3, 1, 1)) + else: + kv_seq_len += cache.offset + 1 if cache is not None else 0 + + from mlx_vlm.models.qwen3_vl.language import ( + apply_multimodal_rotary_pos_emb, + scaled_dot_product_attention, + ) + + cos, sin = self.rotary_emb(values, position_ids) + if mask is not None and isinstance(mask, mx.array): + if isinstance(kv_seq_len, mx.array): + kv_seq_len = kv_seq_len.max().item() + mask = mask[..., :int(kv_seq_len)] + + queries, keys = apply_multimodal_rotary_pos_emb(queries, keys, cos, sin) + if cache is not None: + keys, values = cache.update_and_fetch(keys, values) + + output = scaled_dot_product_attention( + queries, keys, values, cache, scale=self.scale, mask=mask + ) + output = output.transpose(0, 2, 1, 3).reshape(B, L, -1) + + # o_proj + if _hybrid_mode == "prefill" and hasattr(self.o_proj, '_w8a8_w'): + return _w8a8_gemm(output, self.o_proj._w8a8_w, self.o_proj._w8a8_s) + return self.o_proj(output) + + +# ── Fused MLP forward ────────────────────────────────────────── +def _fused_mlp_call(self, x): + """MLP forward: fused gate+up GEMM in prefill, original in decode.""" + if _hybrid_mode == "prefill" and hasattr(self, '_fused_gu_w'): + gate_up = _w8a8_gemm(x, self._fused_gu_w, self._fused_gu_s) + split = self._gu_split + hidden = nn.silu(gate_up[..., :split]) * gate_up[..., split:] + if hasattr(self.down_proj, '_w8a8_w'): + return _w8a8_gemm(hidden, self.down_proj._w8a8_w, self.down_proj._w8a8_s) + return self.down_proj(hidden) + return self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x)) + + +# ── Model conversion ─────────────────────────────────────────── +def convert_model_fused(model, verbose=True): + """Convert Qwen3-VL language model to fused W8A8 hybrid (in-place). + + Fuses QKV projections and gate+up projections for prefill acceleration. + Decode path remains unchanged (original QuantizedLinear). + + Args: + model: Loaded Qwen3-VL model (from mlx_vlm.utils.load). + verbose: Print progress. + + Returns: + dict with conversion stats. + """ + _ensure_ext() + t0 = time.perf_counter() + + layers = model.language_model.model.layers + n_layers = len(layers) + + for i, layer in enumerate(layers): + attn = layer.self_attn + mlp = layer.mlp + + # ── Fuse QKV ── + w_q = _dequant_ql(attn.q_proj) + w_k = _dequant_ql(attn.k_proj) + w_v = _dequant_ql(attn.v_proj) + w_qkv = mx.concatenate([w_q, w_k, w_v], axis=0) + mx.eval(w_qkv) + attn._fused_qkv_w, attn._fused_qkv_s = _make_w8a8(w_qkv) + attn._qkv_split = (w_q.shape[0], w_k.shape[0], w_v.shape[0]) + del w_q, w_k, w_v, w_qkv + + # ── W8A8 for o_proj ── + attn.o_proj._w8a8_w, attn.o_proj._w8a8_s = _make_w8a8(_dequant_ql(attn.o_proj)) + + # ── Fuse gate+up ── + w_gate = _dequant_ql(mlp.gate_proj) + w_up = _dequant_ql(mlp.up_proj) + w_gu = mx.concatenate([w_gate, w_up], axis=0) + mx.eval(w_gu) + mlp._fused_gu_w, mlp._fused_gu_s = _make_w8a8(w_gu) + mlp._gu_split = w_gate.shape[0] + del w_gate, w_up, w_gu + + # ── W8A8 for down_proj ── + mlp.down_proj._w8a8_w, mlp.down_proj._w8a8_s = _make_w8a8(_dequant_ql(mlp.down_proj)) + + # Periodic eval to prevent graph explosion + if (i + 1) % 7 == 0 or i == n_layers - 1: + mx.eval(model.parameters()) + gc.collect() + + # ── Class swap for __call__ override ── + from mlx_vlm.models.qwen3_vl.language import Attention as OrigAttention + from mlx_vlm.models.qwen3_vl.language import MLP as OrigMLP + + class FusedAttention(OrigAttention): + __call__ = _fused_attention_call + + class FusedMLP(OrigMLP): + __call__ = _fused_mlp_call + + for layer in layers: + layer.self_attn.__class__ = FusedAttention + layer.mlp.__class__ = FusedMLP + + elapsed = time.perf_counter() - t0 + stats = { + "n_layers": n_layers, + "fused_qkv": n_layers, + "fused_gate_up": n_layers, + "w8a8_single": n_layers * 2, # o_proj + down_proj + "elapsed_s": elapsed, + } + if verbose: + print(f"[W8A8 Fused] Converted {n_layers} layers " + f"({n_layers} QKV fused + {n_layers} gate_up fused + {n_layers*2} single) " + f"in {elapsed:.1f}s") + return stats diff --git a/cider/kernels/cider_sdpa_vector.h b/cider/kernels/cider_sdpa_vector.h new file mode 100644 index 0000000..b7a68cd --- /dev/null +++ b/cider/kernels/cider_sdpa_vector.h @@ -0,0 +1,337 @@ +// cider_sdpa_vector.h — Cider v9 optimized SDPA vector kernels +// +// Improvements over stock MLX: +// 1. Contiguous chunk layout (cache-friendly sequential access) +// 2. FlashInfer-style register tiling (TILE=4 unroll) +// 3. Adaptive blocks selection per GQA ratio +// +// Template parameters: +// T — data type (float, half, bfloat) +// D — head dimension (64, 96, 128, 256) +// BLOCKS — number of blocks for 2-pass (32, 64, 128) +// +// Three kernels: +// cider_sdpa_vector — 1-pass (short N) +// cider_sdpa_vector_2pass_1 — 2-pass pass1 (per-block partials) +// cider_sdpa_vector_2pass_2 — 2-pass pass2 (reduce) + +#pragma once + +#include + +using namespace metal; + +// ═══════════════════════════════════════════════════════════════ +// 1-pass kernel (N <= threshold, no intermediate buffer) +// ═══════════════════════════════════════════════════════════════ +template +[[kernel]] void cider_sdpa_vector( + const device T* queries [[buffer(0)]], + const device T* keys [[buffer(1)]], + const device T* values [[buffer(2)]], + device T* out [[buffer(3)]], + const constant int& gqa_factor [[buffer(4)]], + const constant int& N [[buffer(5)]], + const constant size_t& k_head_stride [[buffer(6)]], + const constant size_t& k_seq_stride [[buffer(7)]], + const constant size_t& v_head_stride [[buffer(8)]], + const constant size_t& v_seq_stride [[buffer(9)]], + const constant float& scale [[buffer(10)]], + uint3 tid [[threadgroup_position_in_grid]], + uint simd_gid [[simdgroup_index_in_threadgroup]], + uint simd_lid [[thread_index_in_simdgroup]]) { + + constexpr int BN = 32; + constexpr int BD = 32; + constexpr int qk_per_thread = D / BD; + constexpr int v_per_thread = D / BD; + + typedef float U; + + const int inner_k_stride = BN * int(k_seq_stride); + const int inner_v_stride = BN * int(v_seq_stride); + + thread U q_reg[qk_per_thread]; + thread U k_reg[qk_per_thread]; + thread U o_reg[v_per_thread]; + + threadgroup U tg_outputs[BN * BD]; + threadgroup U tg_max_scores[BN]; + threadgroup U tg_sum_exp_scores[BN]; + + const int q_batch_head_idx = tid.x; + const int kv_head_idx = q_batch_head_idx / gqa_factor; + const int o_offset = q_batch_head_idx; + + const device T* q_ptr = queries + o_offset * D + simd_lid * qk_per_thread; + const device T* k_ptr = keys + kv_head_idx * int(k_head_stride) + + simd_gid * int(k_seq_stride) + simd_lid * qk_per_thread; + const device T* v_ptr = values + kv_head_idx * int(v_head_stride) + + simd_gid * int(v_seq_stride) + simd_lid * v_per_thread; + + for (int i = 0; i < qk_per_thread; i++) { + q_reg[i] = static_cast(scale) * static_cast(q_ptr[i]); + } + for (int i = 0; i < v_per_thread; i++) { + o_reg[i] = 0; + } + + U max_score = -1e38f; + U sum_exp_score = 0; + + for (int i = simd_gid; i < N; i += BN) { + for (int j = 0; j < qk_per_thread; j++) { + k_reg[j] = static_cast(k_ptr[j]); + } + U score = 0; + for (int j = 0; j < qk_per_thread; j++) { + score += q_reg[j] * k_reg[j]; + } + score = simd_sum(score); + + U new_max = max(max_score, score); + U factor = fast::exp(max_score - new_max); + U exp_score = fast::exp(score - new_max); + max_score = new_max; + sum_exp_score = sum_exp_score * factor + exp_score; + + for (int j = 0; j < v_per_thread; j++) { + o_reg[j] = o_reg[j] * factor + exp_score * static_cast(v_ptr[j]); + } + + k_ptr += inner_k_stride; + v_ptr += inner_v_stride; + } + + if (simd_lid == 0) { + tg_max_scores[simd_gid] = max_score; + tg_sum_exp_scores[simd_gid] = sum_exp_score; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + max_score = tg_max_scores[simd_lid]; + U new_max_final = simd_max(max_score); + U factor = fast::exp(max_score - new_max_final); + sum_exp_score = simd_sum(tg_sum_exp_scores[simd_lid] * factor); + + for (int i = 0; i < v_per_thread; i++) { + tg_outputs[simd_lid * BD + simd_gid] = o_reg[i]; + threadgroup_barrier(mem_flags::mem_threadgroup); + o_reg[i] = simd_sum(tg_outputs[simd_gid * BD + simd_lid] * factor); + o_reg[i] = sum_exp_score == 0 ? o_reg[i] : (o_reg[i] / sum_exp_score); + threadgroup_barrier(mem_flags::mem_threadgroup); + } + + if (simd_lid == 0) { + device T* o_ptr = out + o_offset * D + simd_gid * v_per_thread; + for (int i = 0; i < v_per_thread; i++) { + o_ptr[i] = static_cast(o_reg[i]); + } + } +} + + +// ═══════════════════════════════════════════════════════════════ +// 2-pass pass1: per-block partial results +// v9: contiguous chunks + TILE=4 register tiling +// ═══════════════════════════════════════════════════════════════ +template +[[kernel]] void cider_sdpa_vector_2pass_1( + const device T* queries [[buffer(0)]], + const device T* keys [[buffer(1)]], + const device T* values [[buffer(2)]], + device T* partials [[buffer(3)]], + device float* sums [[buffer(4)]], + device float* maxs [[buffer(5)]], + const constant int& N [[buffer(7)]], + const constant size_t& k_head_stride [[buffer(8)]], + const constant size_t& k_seq_stride [[buffer(9)]], + const constant size_t& v_head_stride [[buffer(10)]], + const constant size_t& v_seq_stride [[buffer(11)]], + const constant float& scale [[buffer(12)]], + uint3 tptg [[threads_per_threadgroup]], + uint3 tidtg [[thread_position_in_threadgroup]], + uint3 tid [[threadgroup_position_in_grid]], + uint3 tpg [[threadgroups_per_grid]], + uint simd_lid [[thread_index_in_simdgroup]]) { + + constexpr int BD = 32; + constexpr int qk_per_thread = D / BD; + constexpr int v_per_thread = D / BD; + constexpr int TILE = 4; + + typedef float U; + + thread U q_reg[qk_per_thread]; + thread U o_reg[v_per_thread] = {0}; + + const int kv_head_idx = tid.x; + const int batch_idx = tid.y; + const int block_idx = tid.z; + const int gqa_factor = tptg.y; + const int q_head_idx = gqa_factor * kv_head_idx + tidtg.y; + const int num_kv_heads = tpg.x; + const int num_q_heads = num_kv_heads * gqa_factor; + const int q_batch_head_idx = batch_idx * num_q_heads + q_head_idx; + const int o_offset = q_batch_head_idx; + + queries += o_offset * D + simd_lid * qk_per_thread; + + const int kv_batch_head_idx = batch_idx * num_kv_heads + kv_head_idx; + + // v9: contiguous chunk layout (not interleaved) + const int chunk_size = (N + BLOCKS - 1) / BLOCKS; + const int kv_start = block_idx * chunk_size; + const int kv_end = min(kv_start + chunk_size, N); + + const device T* k_ptr = keys + kv_batch_head_idx * int(k_head_stride) + + kv_start * int(k_seq_stride) + simd_lid * qk_per_thread; + const device T* v_ptr = values + kv_batch_head_idx * int(v_head_stride) + + kv_start * int(v_seq_stride) + simd_lid * v_per_thread; + + device T* o_ptr = partials + o_offset * BLOCKS * D + block_idx * D + + simd_lid * v_per_thread; + + // Read query + for (int i = 0; i < qk_per_thread; i++) { + q_reg[i] = static_cast(scale) * static_cast(queries[i]); + } + + U max_score = -1e38f; + U sum_exp_score = 0; + + const int kss = int(k_seq_stride); + const int vss = int(v_seq_stride); + + // Main loop with TILE=4 unrolling + int pos = kv_start; + const int tiled_end = kv_start + ((kv_end - kv_start) / TILE) * TILE; + + for (; pos < tiled_end; pos += TILE) { + U scores[TILE]; + for (int t = 0; t < TILE; t++) { + U score = 0; + const device T* kt = k_ptr + t * kss; + for (int j = 0; j < qk_per_thread; j++) { + score += q_reg[j] * static_cast(kt[j]); + } + scores[t] = simd_sum(score); + } + for (int t = 0; t < TILE; t++) { + U new_max = max(max_score, scores[t]); + U factor = fast::exp(max_score - new_max); + U exp_score = fast::exp(scores[t] - new_max); + max_score = new_max; + sum_exp_score = sum_exp_score * factor + exp_score; + + const device T* vt = v_ptr + t * vss; + for (int j = 0; j < v_per_thread; j++) { + o_reg[j] = o_reg[j] * factor + exp_score * static_cast(vt[j]); + } + } + k_ptr += TILE * kss; + v_ptr += TILE * vss; + } + // Remainder + for (; pos < kv_end; pos++) { + U score = 0; + for (int j = 0; j < qk_per_thread; j++) { + score += q_reg[j] * static_cast(k_ptr[j]); + } + score = simd_sum(score); + + U new_max = max(max_score, score); + U factor = fast::exp(max_score - new_max); + U exp_score = fast::exp(score - new_max); + max_score = new_max; + sum_exp_score = sum_exp_score * factor + exp_score; + + for (int j = 0; j < v_per_thread; j++) { + o_reg[j] = o_reg[j] * factor + exp_score * static_cast(v_ptr[j]); + } + k_ptr += kss; + v_ptr += vss; + } + + // Write partial results + if (simd_lid == 0) { + sums[o_offset * BLOCKS + block_idx] = sum_exp_score; + maxs[o_offset * BLOCKS + block_idx] = max_score; + } + for (int i = 0; i < v_per_thread; i++) { + o_ptr[i] = static_cast(o_reg[i]); + } +} + + +// ═══════════════════════════════════════════════════════════════ +// 2-pass pass2: reduce partial results across blocks +// ═══════════════════════════════════════════════════════════════ +template +[[kernel]] void cider_sdpa_vector_2pass_2( + const device T* partials [[buffer(0)]], + const device float* sums [[buffer(1)]], + const device float* maxs [[buffer(2)]], + device T* out [[buffer(3)]], + uint3 tid [[threadgroup_position_in_grid]], + uint simd_gid [[simdgroup_index_in_threadgroup]], + uint simd_lid [[thread_index_in_simdgroup]]) { + + constexpr int BN = 32; + constexpr int BD = 32; + constexpr int elem_per_thread = D / BD; + + typedef float U; + + thread U o_reg[elem_per_thread] = {0}; + threadgroup U tg_outputs[BN * BD]; + + const int head_idx = tid.x; + + const device T* p_ptr = partials + head_idx * BLOCKS * D + + simd_gid * D + simd_lid * elem_per_thread; + const device float* s_ptr = sums + head_idx * BLOCKS; + const device float* m_ptr = maxs + head_idx * BLOCKS; + + // Reduce max + U max_score = -1e38f; + for (int b = 0; b < BLOCKS / BN; ++b) { + max_score = max(max_score, m_ptr[simd_lid + BN * b]); + } + max_score = simd_max(max_score); + + // Reduce sum_exp + U sum_exp_score = 0; + for (int b = 0; b < BLOCKS / BN; ++b) { + U factor = fast::exp(m_ptr[simd_lid + BN * b] - max_score); + sum_exp_score += factor * s_ptr[simd_lid + BN * b]; + } + sum_exp_score = simd_sum(sum_exp_score); + + // Reduce partials + const device float* m_walk = m_ptr; + for (int b = 0; b < BLOCKS / BN; ++b) { + U factor = fast::exp(m_walk[simd_gid] - max_score); + for (int i = 0; i < elem_per_thread; i++) { + o_reg[i] += factor * static_cast(p_ptr[i]); + } + m_walk += BN; + p_ptr += BN * D; + } + + // Transpose + reduce via shared memory + for (int i = 0; i < elem_per_thread; i++) { + tg_outputs[simd_lid * BD + simd_gid] = o_reg[i]; + threadgroup_barrier(mem_flags::mem_threadgroup); + o_reg[i] = simd_sum(tg_outputs[simd_gid * BD + simd_lid]); + o_reg[i] = sum_exp_score == 0 ? o_reg[i] : (o_reg[i] / sum_exp_score); + threadgroup_barrier(mem_flags::mem_threadgroup); + } + + if (simd_lid == 0) { + device T* o_ptr = out + head_idx * D + simd_gid * elem_per_thread; + for (int i = 0; i < elem_per_thread; i++) { + o_ptr[i] = static_cast(o_reg[i]); + } + } +} diff --git a/cider/kernels/cider_sdpa_vector.metal b/cider/kernels/cider_sdpa_vector.metal new file mode 100644 index 0000000..bf7f8db --- /dev/null +++ b/cider/kernels/cider_sdpa_vector.metal @@ -0,0 +1,42 @@ +// cider_sdpa_vector.metal — Template instantiations for Cider v9 SDPA kernels +// +// Naming convention: +// 1-pass: cider_sdpa_vector_{type}_{D} +// 2-pass1: cider_sdpa_vector_2pass_1_{type}_{D}_b{BLOCKS} +// 2-pass2: cider_sdpa_vector_2pass_2_{type}_{D}_b{BLOCKS} + +#include +#include "cider_sdpa_vector.h" + +using namespace metal; + +// ── Helper macros ──────────────────────────────────────────────── +#define instantiate_1pass(type, type_name, D) \ + template [[host_name("cider_sdpa_vector_" #type_name "_" #D)]] \ + [[kernel]] decltype(cider_sdpa_vector) cider_sdpa_vector; + +#define instantiate_2pass(type, type_name, D, B) \ + template [[host_name("cider_sdpa_vector_2pass_1_" #type_name "_" #D "_b" #B)]] \ + [[kernel]] decltype(cider_sdpa_vector_2pass_1) cider_sdpa_vector_2pass_1; \ + template [[host_name("cider_sdpa_vector_2pass_2_" #type_name "_" #D "_b" #B)]] \ + [[kernel]] decltype(cider_sdpa_vector_2pass_2) cider_sdpa_vector_2pass_2; + +#define instantiate_all_blocks(type, type_name, D) \ + instantiate_2pass(type, type_name, D, 32) \ + instantiate_2pass(type, type_name, D, 64) \ + instantiate_2pass(type, type_name, D, 128) + +#define instantiate_heads(type, type_name) \ + instantiate_1pass(type, type_name, 64) \ + instantiate_1pass(type, type_name, 96) \ + instantiate_1pass(type, type_name, 128) \ + instantiate_1pass(type, type_name, 256) \ + instantiate_all_blocks(type, type_name, 64) \ + instantiate_all_blocks(type, type_name, 96) \ + instantiate_all_blocks(type, type_name, 128) \ + instantiate_all_blocks(type, type_name, 256) + +// ── Instantiate for all types ──────────────────────────────────── +instantiate_heads(float, float32) +instantiate_heads(half, float16) +instantiate_heads(bfloat, bfloat16) diff --git a/cider/kernels/pergroup_int8_gemm.metal b/cider/kernels/pergroup_int8_gemm.metal new file mode 100644 index 0000000..a843994 --- /dev/null +++ b/cider/kernels/pergroup_int8_gemm.metal @@ -0,0 +1,223 @@ +// ============================================================ +// Per-group INT8 TensorOps GEMM — symmetric quantization (bias=0) +// Target: Apple M5 (G17G), Metal 4 +// +// Weight layout: B is [N, K] int8 (per-group symmetric quantized) +// - scales_w: [num_groups, N] float32 (TRANSPOSED for coalesced access) +// - scales_a: [M] float32 (per-token activation scales) +// +// Computes: C[m,n] = scale_a[m] * Sigma_g { float(dot_int32[m,n,g]) * scale_w[g,n] } +// where dot_int32[m,n,g] = Sigma_{k in group g} A_int8[m,k] * B_int8[n,k] +// +// Supported group_size: 64, 128, 256 +// V2: scale_w transposed [num_groups, N] for coalesced SIMD access +// ============================================================ + +#include +#include +using namespace metal; + +// -- NAXFrag layout constants +constant constexpr short kEPF = 8; +constant constexpr short kEC = 4; +constant constexpr short kERJ = 8; + +// -- NAXFrag coordinate mapping +inline short2 nax_coord(ushort lid) { + short qid = short(lid >> 2); + short fm = ((qid & 4) | ((short(lid) >> 1) & 3)); + short fn = ((qid & 2) | (short(lid) & 1)) * 4; + return short2{fn, fm}; +} + +// -- Fragment load: device -> register +template +inline void frag_load(thread T *dst, const device T *src, int ld, short2 sc, + short off_m = 0, short off_n = 0) { + src += (sc.y + off_m) * ld + (sc.x + off_n); + for (short i = 0; i < 2; i++) { + for (short j = 0; j < kEC; j++) { + dst[i * kEC + j] = src[(i * kERJ) * ld + j]; + } + } +} + +// -- Per-group GEMM implementation (scale_w transposed: [num_groups, N]) +template +void pergroup_gemm_impl(const device int8_t *A, const device int8_t *B, + device half *C, uint M, uint N, uint K, + const device float *scale_a, + const device float *scale_w, // [num_groups, N] transposed + const device half *bias, // [N] half + uint swizzle_log, uint tiles_m, uint tiles_n, + uint2 tgid, uint sgid, uint lid) { + constexpr int SM = BM / WM; + constexpr int SN = BN / WN; + constexpr short TM = SM / 16; + constexpr short TN = SN / 16; + constexpr short TK = SK / 16; + + uint tid_y = (tgid.y << swizzle_log) + (tgid.x & ((1u << swizzle_log) - 1u)); + uint tid_x = tgid.x >> swizzle_log; + if (tid_x >= tiles_n || tid_y >= tiles_m) { + return; + } + + short2 sc = nax_coord(ushort(lid)); + uint sg_row = sgid / WN; + uint sg_col = sgid % WN; + uint m_base = tid_y * BM + sg_row * SM; + uint n_base = tid_x * BN + sg_col * SN; + + const device int8_t *sg_A = A + m_base * K; + const device int8_t *sg_B = B + n_base * K; + + uint num_groups = K / BK; + + constexpr auto desc = mpp::tensor_ops::matmul2d_descriptor( + 16, 32, 16, false, true, true, + mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate); + mpp::tensor_ops::matmul2d gemm_op; + + auto ct_a = + gemm_op.get_left_input_cooperative_tensor(); + auto ct_b = + gemm_op.get_right_input_cooperative_tensor(); + auto ct_c = + gemm_op.get_destination_cooperative_tensor(); + + // Float accumulator (across all groups) + float acc[TM * TN][kEPF]; + for (int f = 0; f < TM * TN; f++) { + for (int i = 0; i < kEPF; i++) { + acc[f][i] = 0.0f; + } + } + + // -- Main K loop: one iteration per group + for (uint g = 0; g < num_groups; g++) { + // INT32 accumulator for this group + int32_t c_frags[TM * TN][kEPF]; + for (int f = 0; f < TM * TN; f++) { + for (int i = 0; i < kEPF; i++) { + c_frags[f][i] = 0; + } + } + + // Inner loop within group + for (int kk1 = 0; kk1 < BK; kk1 += SK) { + int8_t a_frags[TM][TK][kEPF]; + int8_t b_frags[TN][TK][kEPF]; + + for (short mm = 0; mm < TM; mm++) { + for (short kk = 0; kk < TK; kk++) { + frag_load(a_frags[mm][kk], sg_A + kk1, int(K), sc, short(mm * 16), + short(kk * 16)); + } + } + + for (short nn = 0; nn < TN; nn++) { + for (short kk = 0; kk < TK; kk++) { + frag_load(b_frags[nn][kk], sg_B + kk1, int(K), sc, short(nn * 16), + short(kk * 16)); + } + } + + for (short mm = 0; mm < TM; mm++) { + for (short nn = 0; nn < TN; nn += 2) { + for (short kk = 0; kk < TK; kk++) { + for (short i = 0; i < kEPF; i++) { + ct_a[i] = a_frags[mm][kk][i]; + } + for (short i = 0; i < kEPF; i++) { + ct_b[i] = b_frags[nn][kk][i]; + ct_b[kEPF + i] = b_frags[nn + 1][kk][i]; + } + short c0 = mm * TN + nn, c1 = c0 + 1; + for (short i = 0; i < kEPF; i++) { + ct_c[i] = c_frags[c0][i]; + ct_c[kEPF + i] = c_frags[c1][i]; + } + gemm_op.run(ct_a, ct_b, ct_c); + for (short i = 0; i < kEPF; i++) { + c_frags[c0][i] = ct_c[i]; + c_frags[c1][i] = ct_c[kEPF + i]; + } + } + } + } + } + + // -- Flush: int32 * scale_w[g, n] -> accumulate + // scale_w is [num_groups, N]: scale_w[g * N + n_idx] is coalesced for adjacent n + for (short mm = 0; mm < TM; mm++) { + for (short nn = 0; nn < TN; nn++) { + short fidx = mm * TN + nn; + float sw[kEC]; + for (short j = 0; j < kEC; j++) { + uint n_idx = n_base + uint(sc.x) + uint(nn * 16) + uint(j); + sw[j] = (n_idx < N) ? scale_w[g * N + n_idx] : 0.0f; + } + for (short i = 0; i < 2; i++) { + for (short j = 0; j < kEC; j++) { + acc[fidx][i * kEC + j] += float(c_frags[fidx][i * kEC + j]) * sw[j]; + } + } + } + } + + sg_A += BK; + sg_B += BK; + } + + // -- Store: acc * scale_a + bias -> half + device half *D = C + m_base * N + n_base; + for (short mm = 0; mm < TM; mm++) { + for (short nn = 0; nn < TN; nn++) { + short fidx = mm * TN + nn; + for (short i = 0; i < 2; i++) { + for (short j = 0; j < kEC; j++) { + uint mi = m_base + uint(sc.y) + uint(mm * 16 + i * kERJ); + uint ni = n_base + uint(sc.x) + uint(nn * 16 + j); + if (mi < M && ni < N) { + float val = acc[fidx][i * kEC + j] * scale_a[mi] + float(bias[ni]); + D[(sc.y + mm * 16 + i * kERJ) * int(N) + (sc.x + nn * 16 + j)] = + half(val); + } + } + } + } + } +} + +// ============================================================ +// Kernel entry points +// ============================================================ + +#define GEMM_ENTRY(SUFFIX, BM_V, BN_V, BK_V, WM_V, WN_V) \ + kernel void pergroup_int8_gemm_##SUFFIX( \ + const device int8_t *A [[buffer(0)]], \ + const device int8_t *B [[buffer(1)]], device half *C [[buffer(2)]], \ + constant uint &M [[buffer(3)]], constant uint &N [[buffer(4)]], \ + constant uint &K [[buffer(5)]], \ + const device float *scale_a [[buffer(6)]], \ + const device float *scale_w [[buffer(7)]], \ + constant uint &swizzle_log [[buffer(8)]], \ + constant uint &tiles_m [[buffer(9)]], \ + constant uint &tiles_n [[buffer(10)]], \ + const device half *bias [[buffer(11)]], \ + uint2 tgid [[threadgroup_position_in_grid]], \ + uint sgid [[simdgroup_index_in_threadgroup]], \ + uint lid [[thread_index_in_simdgroup]]) { \ + pergroup_gemm_impl( \ + A, B, C, M, N, K, scale_a, scale_w, bias, swizzle_log, tiles_m, \ + tiles_n, tgid, sgid, lid); \ + } + +GEMM_ENTRY(g64, 128, 128, 64, 4, 4) +GEMM_ENTRY(g64_small, 32, 128, 64, 1, 4) +GEMM_ENTRY(g128, 128, 128, 128, 4, 4) +GEMM_ENTRY(g128_small, 32, 128, 128, 1, 4) +GEMM_ENTRY(g256, 128, 128, 256, 4, 4) +GEMM_ENTRY(g256_small, 32, 128, 256, 1, 4) diff --git a/cider/kernels/pergroup_int8_mv.metal b/cider/kernels/pergroup_int8_mv.metal new file mode 100644 index 0000000..90cc42e --- /dev/null +++ b/cider/kernels/pergroup_int8_mv.metal @@ -0,0 +1,158 @@ +// ============================================================ +// Per-group INT8 MV kernel V5 — symmetric-only fast path +// +// Match MLX qmv_fast: 2 SG, VPT=8, block=256 +// Symmetric quantization: new_bias is always zero, skip correction. +// +// Formula: +// y[n] = Sigma_g { scale_w[g,n] * dot_g(w[n], x) } + bias[n] +// +// V2: scale_w transposed [num_groups, N] for coalesced access +// ============================================================ + +#include +using namespace metal; + +constant constexpr int SIMD_SIZE = 32; +constant constexpr int NUM_SIMDGROUPS = 2; +constant constexpr int RESULTS_PER_SG = 4; +constant constexpr int VPT = 8; +constant constexpr int BLOCK_K = VPT * SIMD_SIZE; // 256 + +template +inline void pergroup_mv_v5_impl( + const device half *x, + const device int8_t *W, + device half *y, + const device float *scale_w, // [num_groups, N] transposed + const device float *new_bias [[maybe_unused]], + constant uint &N, constant uint &K, + const device half *bias, + uint tgid [[threadgroup_position_in_grid]], + uint sgid [[simdgroup_index_in_threadgroup]], + uint lid [[thread_index_in_simdgroup]]) +{ + constexpr int rows_per_tg = NUM_SIMDGROUPS * RESULTS_PER_SG; // 8 + const uint out_row = tgid * rows_per_tg + sgid * RESULTS_PER_SG; + const uint num_groups = K / GROUP_SIZE; + + float result[RESULTS_PER_SG] = {0.0f, 0.0f, 0.0f, 0.0f}; + + // Pointers + const device half *xp = x + lid * VPT; + + const device int8_t *wrows[RESULTS_PER_SG]; + for (int r = 0; r < RESULTS_PER_SG; r++) { + uint n_idx = out_row + r; + wrows[r] = (n_idx < N) ? (W + n_idx * K + lid * VPT) : W; + } + + for (uint k = 0; k < K; k += BLOCK_K) { + // Load x tile + float xv[VPT]; + for (int i = 0; i < VPT; i++) { + xv[i] = float(xp[i]); + } + + // Group index for this thread data + uint g = (k + lid * VPT) / GROUP_SIZE; + + for (int r = 0; r < RESULTS_PER_SG; r++) { + uint n_idx = out_row + r; + if (n_idx >= N) continue; + + // Load 8 int8 weights as 2x uint32 + uint32_t packed0 = *reinterpret_cast(wrows[r]); + uint32_t packed1 = *reinterpret_cast(wrows[r] + 4); + + float b0 = float(int8_t(packed0 & 0xFF)); + float b1 = float(int8_t((packed0 >> 8) & 0xFF)); + float b2 = float(int8_t((packed0 >> 16) & 0xFF)); + float b3 = float(int8_t((packed0 >> 24) & 0xFF)); + float b4 = float(int8_t(packed1 & 0xFF)); + float b5 = float(int8_t((packed1 >> 8) & 0xFF)); + float b6 = float(int8_t((packed1 >> 16) & 0xFF)); + float b7 = float(int8_t((packed1 >> 24) & 0xFF)); + + float dot = xv[0]*b0 + xv[1]*b1 + xv[2]*b2 + xv[3]*b3 + + xv[4]*b4 + xv[5]*b5 + xv[6]*b6 + xv[7]*b7; + + // Transposed: scale_w[g * N + n_idx] (coalesced for adjacent n) + float sw = scale_w[g * N + (out_row + r)]; + result[r] += sw * dot; + } + + // Advance pointers + xp += BLOCK_K; + for (int r = 0; r < RESULTS_PER_SG; r++) { + wrows[r] += BLOCK_K; + } + } + + // Reduce across simdgroup + for (int r = 0; r < RESULTS_PER_SG; r++) { + result[r] = simd_sum(result[r]); + } + + // Write result + if (lid == 0) { + for (int r = 0; r < RESULTS_PER_SG; r++) { + uint n_idx = out_row + r; + if (n_idx < N) { + y[n_idx] = half(result[r] + float(bias[n_idx])); + } + } + } +} + +// ============================================================ +// Entry points +// ============================================================ + +kernel void pergroup_int8_mv_g64( + const device half *x [[buffer(0)]], + const device int8_t *W [[buffer(1)]], + device half *y [[buffer(2)]], + const device float *scale_w [[buffer(3)]], + constant uint &N [[buffer(4)]], + constant uint &K [[buffer(5)]], + const device half *bias [[buffer(6)]], + const device float *new_bias [[buffer(7)]], + uint tgid [[threadgroup_position_in_grid]], + uint sgid [[simdgroup_index_in_threadgroup]], + uint lid [[thread_index_in_simdgroup]]) +{ + pergroup_mv_v5_impl<64>(x, W, y, scale_w, new_bias, N, K, bias, tgid, sgid, lid); +} + +kernel void pergroup_int8_mv_g128( + const device half *x [[buffer(0)]], + const device int8_t *W [[buffer(1)]], + device half *y [[buffer(2)]], + const device float *scale_w [[buffer(3)]], + constant uint &N [[buffer(4)]], + constant uint &K [[buffer(5)]], + const device half *bias [[buffer(6)]], + const device float *new_bias [[buffer(7)]], + uint tgid [[threadgroup_position_in_grid]], + uint sgid [[simdgroup_index_in_threadgroup]], + uint lid [[thread_index_in_simdgroup]]) +{ + pergroup_mv_v5_impl<128>(x, W, y, scale_w, new_bias, N, K, bias, tgid, sgid, lid); +} + +kernel void pergroup_int8_mv_g256( + const device half *x [[buffer(0)]], + const device int8_t *W [[buffer(1)]], + device half *y [[buffer(2)]], + const device float *scale_w [[buffer(3)]], + constant uint &N [[buffer(4)]], + constant uint &K [[buffer(5)]], + const device half *bias [[buffer(6)]], + const device float *new_bias [[buffer(7)]], + uint tgid [[threadgroup_position_in_grid]], + uint sgid [[simdgroup_index_in_threadgroup]], + uint lid [[thread_index_in_simdgroup]]) +{ + pergroup_mv_v5_impl<256>(x, W, y, scale_w, new_bias, N, K, bias, tgid, sgid, lid); +} diff --git a/cider/kernels/w4a8_matmul.metal b/cider/kernels/w4a8_matmul.metal new file mode 100644 index 0000000..d7ce4d8 --- /dev/null +++ b/cider/kernels/w4a8_matmul.metal @@ -0,0 +1,288 @@ +// ============================================================ +// W4A8 INT4-weight × INT8-activation → FP16 TensorOps GEMM +// Target: Apple M5, Metal 4 +// +// V3: Optimized inline unpack with precomputed base pointers. +// Key insight: for fragment's 8 elements (2 rows × 4 cols), +// the 2 rows are at k and k+8. For packed [K/2, N]: +// - Row 0 (k): byte at (k/2)*N + n, use k&1 to select nibble +// - Row 1 (k+8): byte at ((k+8)/2)*N + n, use (k+8)&1 to select nibble +// Since k+8 has same parity as k (8 is even), both rows use same nibble. +// => Can share nibble selection logic. +// +// Further: read 4 consecutive bytes at once per row (cols are contiguous +// in N dimension), then extract 4 nibbles. This maximizes memory bandwidth. +// ============================================================ + +#include +#include +using namespace metal; + +constant constexpr short kElemsPerFrag = 8; +constant constexpr short kElemCols = 4; +constant constexpr short kElemRowsJump = 8; + +inline short2 nax_get_coord(ushort lid) { + short qid = short(lid >> 2); + short fm = ((qid & 4) | ((short(lid) >> 1) & 3)); + short fn = ((qid & 2) | (short(lid) & 1)) * 4; + return short2{fn, fm}; +} + +// ── Fragment load A from device memory ────────────────────────── +inline void frag_load_a(thread int8_t *dst, const device int8_t *src, int ld, + short2 sc, short off_m, short off_n) { + src += (sc.y + off_m) * ld + (sc.x + off_n); + for (short i = 0; i < 2; i++) { + for (short j = 0; j < kElemCols; j++) { + dst[i * kElemCols + j] = src[(i * kElemRowsJump) * ld + j]; + } + } +} + +// ── Fragment load B with optimized W4 unpack ──────────────────── +// Fragment maps to 2 rows (k, k+8) × 4 cols (n, n+1, n+2, n+3). +// packed_w layout: [K/2, N] uint8, high nibble = even k, low nibble = odd k. +// Pre-compute row pointers and read 4 consecutive bytes per row. +inline void frag_load_b_w4(thread int8_t *dst, const device uint8_t *packed_w, + uint N, uint k_base, uint n_base, short2 sc) { + uint k0 = k_base + uint(sc.y); // first row + uint k1 = k0 + kElemRowsJump; // second row (k+8) + uint n = n_base + uint(sc.x); // column start + + // Both k0 and k1 have same parity (differ by 8) + bool use_low = (k0 & 1u); + + const device uint8_t *row0 = packed_w + (k0 >> 1) * N + n; + const device uint8_t *row1 = packed_w + (k1 >> 1) * N + n; + + if (use_low) { + for (short j = 0; j < kElemCols; j++) { + dst[j] = int8_t(row0[j] & 0xF) - 8; + dst[kElemCols + j] = int8_t(row1[j] & 0xF) - 8; + } + } else { + for (short j = 0; j < kElemCols; j++) { + dst[j] = int8_t(row0[j] >> 4) - 8; + dst[kElemCols + j] = int8_t(row1[j] >> 4) - 8; + } + } +} + +// ── Fragment store with fused dequant ─────────────────────────── +inline void nax_frag_store_dequant(const thread int32_t *src, device half *dst, + int ld, short2 sc, short off_m, short off_n, + uint M, uint N, uint m_base, uint n_base, + const device float *scale_a, + const device float *scale_w) { + for (short i = 0; i < 2; i++) { + for (short j = 0; j < kElemCols; j++) { + uint mi = m_base + sc.y + off_m + i * kElemRowsJump; + uint ni = n_base + sc.x + off_n + j; + if (mi < M && ni < N) { + float val = float(src[i * kElemCols + j]) * scale_a[mi] * scale_w[ni]; + dst[(sc.y + off_m + i * kElemRowsJump) * ld + (sc.x + off_n + j)] = + half(val); + } + } + } +} + +// ── W4A8 GEMM V3 ─────────────────────────────────────────────── +template +void w4a8_gemm_impl(const device int8_t *A, const device uint8_t *packed_w, + device half *C, uint M, uint N, uint K, + const device float *scale_a, const device float *scale_w, + uint swizzle_log, uint tiles_m, uint tiles_n, uint2 tgid, + uint sgid, uint lid) { + constexpr int SM = BM / WM; + constexpr int SN = BN / WN; + constexpr short TM = SM / 16; + constexpr short TN = SN / 16; + constexpr short TK = SK / 16; + + uint tid_y = (tgid.y << swizzle_log) + (tgid.x & ((1u << swizzle_log) - 1u)); + uint tid_x = tgid.x >> swizzle_log; + if (tid_x >= tiles_n || tid_y >= tiles_m) { + return; + } + + short2 sc = nax_get_coord(ushort(lid)); + uint sg_row = sgid / WN; + uint sg_col = sgid % WN; + uint m_base = tid_y * BM + sg_row * SM; + uint n_base = tid_x * BN + sg_col * SN; + + const device int8_t *sg_A = A + m_base * K; + + constexpr auto matmul_desc = mpp::tensor_ops::matmul2d_descriptor( + 16, 32, 16, false, false, true, + mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate); + mpp::tensor_ops::matmul2d gemm_op; + + auto ct_a = + gemm_op.get_left_input_cooperative_tensor(); + auto ct_b = + gemm_op.get_right_input_cooperative_tensor(); + auto ct_c = + gemm_op.get_destination_cooperative_tensor(); + + int32_t c_frags[TM * TN][kElemsPerFrag]; + for (int f = 0; f < TM * TN; f++) { + for (int e = 0; e < kElemsPerFrag; e++) { + c_frags[f][e] = 0; + } + } + int gemm_k_iters = int(K) / BK; + + for (int kk0 = 0; kk0 < gemm_k_iters; kk0++) { + uint k_offset = uint(kk0) * BK; + + for (int kk1 = 0; kk1 < BK; kk1 += SK) { + int8_t a_frags[TM][TK][kElemsPerFrag]; + int8_t b_frags[TK][TN][kElemsPerFrag]; + volatile int compiler_barrier; + + for (short mm = 0; mm < TM; mm++) { + for (short kk = 0; kk < TK; kk++) { + frag_load_a(a_frags[mm][kk], sg_A + k_offset + kk1, int(K), sc, + short(mm * 16), short(kk * 16)); + } + } + + for (short kk = 0; kk < TK; kk++) { + for (short nn = 0; nn < TN; nn++) { + frag_load_b_w4(b_frags[kk][nn], packed_w, N, + k_offset + uint(kk1) + uint(kk * 16), + n_base + uint(nn * 16), sc); + } + } + + for (short mm = 0; mm < TM; mm++) { + for (short nn = 0; nn < TN; nn += 2) { + for (short kk = 0; kk < TK; kk++) { + for (short i = 0; i < kElemsPerFrag; i++) { + ct_a[i] = a_frags[mm][kk][i]; + } + for (short i = 0; i < kElemsPerFrag; i++) { + ct_b[i] = b_frags[kk][nn][i]; + ct_b[kElemsPerFrag + i] = b_frags[kk][nn + 1][i]; + } + short c0 = mm * TN + nn, c1 = c0 + 1; + for (short i = 0; i < kElemsPerFrag; i++) { + ct_c[i] = c_frags[c0][i]; + ct_c[kElemsPerFrag + i] = c_frags[c1][i]; + } + gemm_op.run(ct_a, ct_b, ct_c); + for (short i = 0; i < kElemsPerFrag; i++) { + c_frags[c0][i] = ct_c[i]; + c_frags[c1][i] = ct_c[kElemsPerFrag + i]; + } + } + } + } + (void)compiler_barrier; + } + } + + // Remainder + int rem_k = int(K) - gemm_k_iters * BK; + for (int kk1 = 0; kk1 < rem_k; kk1 += 16) { + int8_t a_frag[TM][kElemsPerFrag]; + int8_t b_frag[TN][kElemsPerFrag]; + short psk = short(max(0, rem_k - kk1)); + uint k_abs = uint(gemm_k_iters * BK + kk1); + + for (short mm = 0; mm < TM; mm++) { + const device int8_t *ptr = sg_A + k_abs + (sc.y + mm * 16) * K + sc.x; + for (short i = 0; i < 2; i++) { + for (short j = 0; j < kElemCols; j++) { + short ki = short(sc.x + j); + a_frag[mm][i * kElemCols + j] = + (ki < psk) ? ptr[(i * kElemRowsJump) * K + j] : int8_t(0); + } + } + } + + for (short nn = 0; nn < TN; nn++) { + uint n = n_base + uint(nn * 16); + for (short i = 0; i < 2; i++) { + for (short j = 0; j < kElemCols; j++) { + short ki = short(sc.y + i * kElemRowsJump); + if (ki < psk) { + uint k = k_abs + uint(ki); + uint ni = n + uint(sc.x) + uint(j); + uint byte_row = k >> 1; + uint8_t packed = packed_w[byte_row * N + ni]; + uint8_t nibble = (k & 1) == 0 ? (packed >> 4) : (packed & 0xF); + b_frag[nn][i * kElemCols + j] = int8_t(nibble) - 8; + } else { + b_frag[nn][i * kElemCols + j] = 0; + } + } + } + } + + for (short mm = 0; mm < TM; mm++) { + for (short i = 0; i < kElemsPerFrag; i++) { + ct_a[i] = a_frag[mm][i]; + } + for (short i = 0; i < kElemsPerFrag; i++) { + ct_b[i] = b_frag[0][i]; + ct_b[kElemsPerFrag + i] = b_frag[1][i]; + } + short c0 = mm * TN, c1 = c0 + 1; + for (short i = 0; i < kElemsPerFrag; i++) { + ct_c[i] = c_frags[c0][i]; + ct_c[kElemsPerFrag + i] = c_frags[c1][i]; + } + gemm_op.run(ct_a, ct_b, ct_c); + for (short i = 0; i < kElemsPerFrag; i++) { + c_frags[c0][i] = ct_c[i]; + c_frags[c1][i] = ct_c[kElemsPerFrag + i]; + } + } + } + + device half *D = C + m_base * N + n_base; + for (short mm = 0; mm < TM; mm++) { + for (short nn = 0; nn < TN; nn++) { + nax_frag_store_dequant(c_frags[mm * TN + nn], D, int(N), sc, + short(mm * 16), short(nn * 16), M, N, m_base, + n_base, scale_a, scale_w); + } + } +} + +kernel void w4a8_matmul_fused_dequant( + const device int8_t *A [[buffer(0)]], + const device uint8_t *packed_w [[buffer(1)]], device half *C [[buffer(2)]], + constant uint &M [[buffer(3)]], constant uint &N [[buffer(4)]], + constant uint &K [[buffer(5)]], const device float *scale_a [[buffer(6)]], + const device float *scale_w [[buffer(7)]], + constant uint &swizzle_log [[buffer(8)]], + constant uint &tiles_m [[buffer(9)]], constant uint &tiles_n [[buffer(10)]], + uint2 tgid [[threadgroup_position_in_grid]], + uint sgid [[simdgroup_index_in_threadgroup]], + uint lid [[thread_index_in_simdgroup]]) { + w4a8_gemm_impl<128, 128, 512, 32, 4, 4>(A, packed_w, C, M, N, K, scale_a, + scale_w, swizzle_log, tiles_m, + tiles_n, tgid, sgid, lid); +} + +kernel void w4a8_matmul_fused_dequant_small( + const device int8_t *A [[buffer(0)]], + const device uint8_t *packed_w [[buffer(1)]], device half *C [[buffer(2)]], + constant uint &M [[buffer(3)]], constant uint &N [[buffer(4)]], + constant uint &K [[buffer(5)]], const device float *scale_a [[buffer(6)]], + const device float *scale_w [[buffer(7)]], + constant uint &swizzle_log [[buffer(8)]], + constant uint &tiles_m [[buffer(9)]], constant uint &tiles_n [[buffer(10)]], + uint2 tgid [[threadgroup_position_in_grid]], + uint sgid [[simdgroup_index_in_threadgroup]], + uint lid [[thread_index_in_simdgroup]]) { + w4a8_gemm_impl<32, 128, 512, 32, 1, 4>(A, packed_w, C, M, N, K, scale_a, + scale_w, swizzle_log, tiles_m, tiles_n, + tgid, sgid, lid); +} diff --git a/cider/kernels/w8a8_int8_mv.metal b/cider/kernels/w8a8_int8_mv.metal new file mode 100644 index 0000000..1fc41d5 --- /dev/null +++ b/cider/kernels/w8a8_int8_mv.metal @@ -0,0 +1,71 @@ +// w8a8_int8_mv.metal — Per-channel symmetric INT8 MV (optimized) +// Matches MLX qmv_fast structure: 2 SG, VPT=8, block=256 +// y[n] = scale_w[n] * dot(W_int8[n,:], x[:]) + bias[n] + +#include +using namespace metal; + +constant constexpr int SIMD_SIZE = 32; +constant constexpr int NUM_SIMDGROUPS = 2; +constant constexpr int RESULTS_PER_SG = 4; +constant constexpr int VPT = 8; +constant constexpr int BLOCK_K = VPT * SIMD_SIZE; // 256 + +kernel void w8a8_int8_mv( + const device half *x [[buffer(0)]], + const device int8_t *W [[buffer(1)]], + device half *y [[buffer(2)]], + const device float *scale_w [[buffer(3)]], + constant uint &N [[buffer(4)]], + constant uint &K [[buffer(5)]], + const device half *bias [[buffer(6)]], + uint tgid [[threadgroup_position_in_grid]], + uint sgid [[simdgroup_index_in_threadgroup]], + uint lid [[thread_index_in_simdgroup]]) +{ + const uint rows_per_tg = NUM_SIMDGROUPS * RESULTS_PER_SG; // 8 + const uint out_row = tgid * rows_per_tg + sgid * RESULTS_PER_SG; + + float result[RESULTS_PER_SG] = {0.0f, 0.0f, 0.0f, 0.0f}; + + // Pointers: each thread handles VPT consecutive elements per block + const device half *xp = x + lid * VPT; + + for (uint k = 0; k < K; k += BLOCK_K) { + // Load x tile into registers + float xv[VPT]; + for (int i = 0; i < VPT; i++) { + xv[i] = float(xp[i]); + } + + // Dot product with weight rows + for (int r = 0; r < RESULTS_PER_SG; r++) { + uint n_idx = out_row + r; + if (n_idx < N) { + const device int8_t *wp = W + n_idx * K + k + lid * VPT; + float dot = 0.0f; + for (int i = 0; i < VPT; i++) { + dot += float(wp[i]) * xv[i]; + } + result[r] += dot; + } + } + xp += BLOCK_K; + } + + // Reduce across simdgroup + for (int r = 0; r < RESULTS_PER_SG; r++) { + result[r] = simd_sum(result[r]); + } + + // Write back: y = scale * dot + bias + if (lid == 0) { + for (int r = 0; r < RESULTS_PER_SG; r++) { + uint n_idx = out_row + r; + if (n_idx < N) { + float val = scale_w[n_idx] * result[r] + float(bias[n_idx]); + y[n_idx] = half(val); + } + } + } +} diff --git a/cider/kernels/w8a8_matmul.metal b/cider/kernels/w8a8_matmul.metal new file mode 100644 index 0000000..cfd9c11 --- /dev/null +++ b/cider/kernels/w8a8_matmul.metal @@ -0,0 +1,483 @@ +// ============================================================ +// W8A8 INT8×INT8→INT32 TensorOps GEMM +// Target: Apple M5 (G17G), Metal 4 +// +// Weight layout: B is [N, K] (row-major), transpose_b=true +// Computes: C[M,N] = A[M,K] × B[N,K]^T +// +// Variants: +// - fused dequant: INT8×INT8→FP16, with per-token/per-channel scales +// - raw INT32: INT8×INT8→INT32, no scale (pure integer GEMM) +// Multi-config: large (BM=128) and small (BM=32) tiles +// Swizzle dispatch for L2 cache locality +// +// matmul2d(16,32,16) via MPP cooperative_tensor +// ============================================================ + +#include +#include +using namespace metal; + +// ── NAXFrag layout constants ──────────────────────────────────── +constant constexpr short kElemsPerFrag = 8; +constant constexpr short kElemCols = 4; +constant constexpr short kElemRowsJump = 8; + +// ── NAXFrag coordinate mapping ────────────────────────────────── +inline short2 nax_get_coord(ushort lid) { + short qid = short(lid >> 2); + short fm = ((qid & 4) | ((short(lid) >> 1) & 3)); + short fn = ((qid & 2) | (short(lid) & 1)) * 4; + return short2{fn, fm}; +} + +// ── Fragment load: device → register ──────────────────────────── +template +inline void nax_frag_load(thread T *dst, const device T *src, int ld, short2 sc, + short off_m = 0, short off_n = 0) { + src += (sc.y + off_m) * ld + (sc.x + off_n); + for (short i = 0; i < 2; i++) { + for (short j = 0; j < kElemCols; j++) { + dst[i * kElemCols + j] = src[(i * kElemRowsJump) * ld + j]; + } + } +} + +// ── Fragment store: raw INT32 (no dequant) ─────────────────── +inline void nax_frag_store_int32(const thread int32_t *src, device int32_t *dst, + int ld, short2 sc, short off_m, short off_n, + uint M, uint N, uint m_base, uint n_base) { + for (short i = 0; i < 2; i++) { + for (short j = 0; j < kElemCols; j++) { + uint mi = m_base + sc.y + off_m + i * kElemRowsJump; + uint ni = n_base + sc.x + off_n + j; + if (mi < M && ni < N) { + dst[(sc.y + off_m + i * kElemRowsJump) * ld + (sc.x + off_n + j)] = + src[i * kElemCols + j]; + } + } + } +} + +// ── Fragment store with bounds check and dequant ──────────────── +inline void nax_frag_store_dequant(const thread int32_t *src, device half *dst, + int ld, short2 sc, short off_m, short off_n, + uint M, uint N, uint m_base, uint n_base, + const device float *scale_a, + const device float *scale_w, + const device half *bias) { + for (short i = 0; i < 2; i++) { + for (short j = 0; j < kElemCols; j++) { + uint mi = m_base + sc.y + off_m + i * kElemRowsJump; + uint ni = n_base + sc.x + off_n + j; + if (mi < M && ni < N) { + float val = float(src[i * kElemCols + j]) * scale_a[mi] * scale_w[ni] + + float(bias[ni]); + dst[(sc.y + off_m + i * kElemRowsJump) * ld + (sc.x + off_n + j)] = + half(val); + } + } + } +} + +// ── Generic GEMM kernel (B is [N, K], transpose_b) ───────────── +// Computes: C[M,N] = A[M,K] × B[N,K]^T +// A is [M, K] row-major, B is [N, K] row-major +// B fragments loaded as [N_tile, K_tile] and hardware transposes via +// transpose_b=true +template +void w8a8_gemm_impl(const device int8_t *A, const device int8_t *B, + device half *C, uint M, uint N, uint K, + const device float *scale_a, const device float *scale_w, + const device half *bias, uint swizzle_log, uint tiles_m, + uint tiles_n, uint2 tgid, uint sgid, uint lid) { + constexpr int SM = BM / WM; // 32 + constexpr int SN = BN / WN; // 32 + constexpr short TM = SM / 16; // 2 + constexpr short TN = SN / 16; // 2 + constexpr short TK = SK / 16; // 2 + + // Swizzle decode + uint tid_y = (tgid.y << swizzle_log) + (tgid.x & ((1u << swizzle_log) - 1u)); + uint tid_x = tgid.x >> swizzle_log; + + if (tid_x >= tiles_n || tid_y >= tiles_m) { + return; + } + + short2 sc = nax_get_coord(ushort(lid)); + uint sg_row = sgid / WN; + uint sg_col = sgid % WN; + uint m_base = tid_y * BM + sg_row * SM; + uint n_base = tid_x * BN + sg_col * SN; + + // A: [M, K] row-major — same as before + const device int8_t *sg_A = A + m_base * K; + // B: [N, K] row-major — pointer to start of n_base-th row + const device int8_t *sg_B = B + n_base * K; + + // transpose_b=true: right operand is [N_frag, K_frag], hardware transposes + constexpr auto desc = mpp::tensor_ops::matmul2d_descriptor( + 16, 32, 16, false, true, true, + mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate); + mpp::tensor_ops::matmul2d gemm_op; + + auto ct_a = + gemm_op.get_left_input_cooperative_tensor(); + auto ct_b = + gemm_op.get_right_input_cooperative_tensor(); + auto ct_c = + gemm_op.get_destination_cooperative_tensor(); + + int32_t c_frags[TM * TN][kElemsPerFrag]; + for (int f = 0; f < TM * TN; f++) { + for (int i = 0; i < kElemsPerFrag; i++) { + c_frags[f][i] = 0; + } + } + + // ── Main K loop ───────────────────────────────────────────── + int gemm_k_iters = int(K) / BK; + + for (int kk0 = 0; kk0 < gemm_k_iters; kk0++) { + threadgroup_barrier(mem_flags::mem_none); + + for (int kk1 = 0; kk1 < BK; kk1 += SK) { + int8_t a_frags[TM][TK][kElemsPerFrag]; + int8_t b_frags[TN][TK][kElemsPerFrag]; // [N_tile, K_tile] for transpose_b + volatile int compiler_barrier; + + // Load A fragments: [M_tile, K_tile], ld=K + for (short mm = 0; mm < TM; mm++) { + for (short kk = 0; kk < TK; kk++) { + nax_frag_load(a_frags[mm][kk], sg_A + kk1, int(K), sc, short(mm * 16), + short(kk * 16)); + } + } + // Load B fragments: [N_tile, K_tile], ld=K (B is [N, K] row-major) + for (short nn = 0; nn < TN; nn++) { + for (short kk = 0; kk < TK; kk++) { + nax_frag_load(b_frags[nn][kk], sg_B + kk1, int(K), sc, short(nn * 16), + short(kk * 16)); + } + } + // Compute: ct_a=[M_frag, K_frag], ct_b=[N_frag, K_frag] (transpose_b) + for (short mm = 0; mm < TM; mm++) { + for (short nn = 0; nn < TN; nn += 2) { + for (short kk = 0; kk < TK; kk++) { + for (short i = 0; i < kElemsPerFrag; i++) { + ct_a[i] = a_frags[mm][kk][i]; + } + for (short i = 0; i < kElemsPerFrag; i++) { + ct_b[i] = b_frags[nn][kk][i]; + ct_b[kElemsPerFrag + i] = b_frags[nn + 1][kk][i]; + } + short c0 = mm * TN + nn, c1 = c0 + 1; + for (short i = 0; i < kElemsPerFrag; i++) { + ct_c[i] = c_frags[c0][i]; + ct_c[kElemsPerFrag + i] = c_frags[c1][i]; + } + gemm_op.run(ct_a, ct_b, ct_c); + for (short i = 0; i < kElemsPerFrag; i++) { + c_frags[c0][i] = ct_c[i]; + c_frags[c1][i] = ct_c[kElemsPerFrag + i]; + } + } + } + } + (void)compiler_barrier; + } + + sg_A += BK; + sg_B += BK; // B is [N, K]: K advances by BK along columns + } + + // ── Remainder K ───────────────────────────────────────────── + int rem_k = int(K) - gemm_k_iters * BK; + for (int kk1 = 0; kk1 < rem_k; kk1 += 16) { + int8_t a_frag[TM][kElemsPerFrag]; + int8_t b_frag[TN][kElemsPerFrag]; + short psk = short(max(0, rem_k - kk1)); + + // Load A remainder: same as before + for (short mm = 0; mm < TM; mm++) { + const device int8_t *ptr = sg_A + kk1 + (sc.y + mm * 16) * K + sc.x; + for (short i = 0; i < 2; i++) { + for (short j = 0; j < kElemCols; j++) { + short ki = short(sc.x + j); + a_frag[mm][i * kElemCols + j] = + (ki < psk) ? ptr[(i * kElemRowsJump) * K + j] : int8_t(0); + } + } + } + + // Load B remainder: B is [N, K], reading [N_tile, K_rem] + for (short nn = 0; nn < TN; nn++) { + const device int8_t *ptr = sg_B + kk1 + (sc.y + nn * 16) * K + sc.x; + for (short i = 0; i < 2; i++) { + for (short j = 0; j < kElemCols; j++) { + short ki = short(sc.x + j); + b_frag[nn][i * kElemCols + j] = + (ki < psk) ? ptr[(i * kElemRowsJump) * K + j] : int8_t(0); + } + } + } + + for (short mm = 0; mm < TM; mm++) { + for (short i = 0; i < kElemsPerFrag; i++) { + ct_a[i] = a_frag[mm][i]; + } + for (short i = 0; i < kElemsPerFrag; i++) { + ct_b[i] = b_frag[0][i]; + ct_b[kElemsPerFrag + i] = b_frag[1][i]; + } + short c0 = mm * TN, c1 = c0 + 1; + for (short i = 0; i < kElemsPerFrag; i++) { + ct_c[i] = c_frags[c0][i]; + ct_c[kElemsPerFrag + i] = c_frags[c1][i]; + } + gemm_op.run(ct_a, ct_b, ct_c); + for (short i = 0; i < kElemsPerFrag; i++) { + c_frags[c0][i] = ct_c[i]; + c_frags[c1][i] = ct_c[kElemsPerFrag + i]; + } + } + } + + // ── Store with fused dequant ──────────────────────────────── + device half *D = C + m_base * N + n_base; + for (short mm = 0; mm < TM; mm++) { + for (short nn = 0; nn < TN; nn++) { + nax_frag_store_dequant(c_frags[mm * TN + nn], D, int(N), sc, + short(mm * 16), short(nn * 16), M, N, m_base, + n_base, scale_a, scale_w, bias); + } + } +} + +// ============================================================ +// Kernel entry points — fused dequant +// ============================================================ + +kernel void w8a8_matmul_fused_dequant( + const device int8_t *A [[buffer(0)]], const device int8_t *B [[buffer(1)]], + device half *C [[buffer(2)]], constant uint &M [[buffer(3)]], + constant uint &N [[buffer(4)]], constant uint &K [[buffer(5)]], + const device float *scale_a [[buffer(6)]], + const device float *scale_w [[buffer(7)]], + constant uint &swizzle_log [[buffer(8)]], + constant uint &tiles_m [[buffer(9)]], constant uint &tiles_n [[buffer(10)]], + const device half *bias [[buffer(11)]], + uint2 tgid [[threadgroup_position_in_grid]], + uint sgid [[simdgroup_index_in_threadgroup]], + uint lid [[thread_index_in_simdgroup]]) { + w8a8_gemm_impl<128, 128, 512, 32, 4, 4>(A, B, C, M, N, K, scale_a, scale_w, + bias, swizzle_log, tiles_m, tiles_n, + tgid, sgid, lid); +} + +kernel void w8a8_matmul_fused_dequant_small( + const device int8_t *A [[buffer(0)]], const device int8_t *B [[buffer(1)]], + device half *C [[buffer(2)]], constant uint &M [[buffer(3)]], + constant uint &N [[buffer(4)]], constant uint &K [[buffer(5)]], + const device float *scale_a [[buffer(6)]], + const device float *scale_w [[buffer(7)]], + constant uint &swizzle_log [[buffer(8)]], + constant uint &tiles_m [[buffer(9)]], constant uint &tiles_n [[buffer(10)]], + const device half *bias [[buffer(11)]], + uint2 tgid [[threadgroup_position_in_grid]], + uint sgid [[simdgroup_index_in_threadgroup]], + uint lid [[thread_index_in_simdgroup]]) { + w8a8_gemm_impl<32, 128, 512, 32, 1, 4>(A, B, C, M, N, K, scale_a, scale_w, + bias, swizzle_log, tiles_m, tiles_n, + tgid, sgid, lid); +} + +// ============================================================ +// Raw INT32 GEMM impl (B is [N, K], transpose_b=true) +// ============================================================ + +template +void w8a8_gemm_int32_impl(const device int8_t *A, const device int8_t *B, + device int32_t *C, uint M, uint N, uint K, + uint swizzle_log, uint tiles_m, uint tiles_n, + uint2 tgid, uint sgid, uint lid) { + constexpr int SM = BM / WM; + constexpr int SN = BN / WN; + constexpr short TM = SM / 16; + constexpr short TN = SN / 16; + constexpr short TK = SK / 16; + + uint tid_y = (tgid.y << swizzle_log) + (tgid.x & ((1u << swizzle_log) - 1u)); + uint tid_x = tgid.x >> swizzle_log; + if (tid_x >= tiles_n || tid_y >= tiles_m) { + return; + } + + short2 sc = nax_get_coord(ushort(lid)); + uint sg_row = sgid / WN; + uint sg_col = sgid % WN; + uint m_base = tid_y * BM + sg_row * SM; + uint n_base = tid_x * BN + sg_col * SN; + + const device int8_t *sg_A = A + m_base * K; + const device int8_t *sg_B = B + n_base * K; + + constexpr auto desc = mpp::tensor_ops::matmul2d_descriptor( + 16, 32, 16, false, true, true, + mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate); + mpp::tensor_ops::matmul2d gemm_op; + + auto ct_a = + gemm_op.get_left_input_cooperative_tensor(); + auto ct_b = + gemm_op.get_right_input_cooperative_tensor(); + auto ct_c = + gemm_op.get_destination_cooperative_tensor(); + + int32_t c_frags[TM * TN][kElemsPerFrag]; + for (int f = 0; f < TM * TN; f++) { + for (int i = 0; i < kElemsPerFrag; i++) { + c_frags[f][i] = 0; + } + } + + int gemm_k_iters = int(K) / BK; + for (int kk0 = 0; kk0 < gemm_k_iters; kk0++) { + threadgroup_barrier(mem_flags::mem_none); + for (int kk1 = 0; kk1 < BK; kk1 += SK) { + int8_t a_frags[TM][TK][kElemsPerFrag]; + int8_t b_frags[TN][TK][kElemsPerFrag]; + volatile int compiler_barrier; + + for (short mm = 0; mm < TM; mm++) { + for (short kk = 0; kk < TK; kk++) { + nax_frag_load(a_frags[mm][kk], sg_A + kk1, int(K), sc, short(mm * 16), + short(kk * 16)); + } + } + + for (short nn = 0; nn < TN; nn++) { + for (short kk = 0; kk < TK; kk++) { + nax_frag_load(b_frags[nn][kk], sg_B + kk1, int(K), sc, short(nn * 16), + short(kk * 16)); + } + } + + for (short mm = 0; mm < TM; mm++) { + for (short nn = 0; nn < TN; nn += 2) { + for (short kk = 0; kk < TK; kk++) { + for (short i = 0; i < kElemsPerFrag; i++) { + ct_a[i] = a_frags[mm][kk][i]; + } + for (short i = 0; i < kElemsPerFrag; i++) { + ct_b[i] = b_frags[nn][kk][i]; + ct_b[kElemsPerFrag + i] = b_frags[nn + 1][kk][i]; + } + short c0 = mm * TN + nn, c1 = c0 + 1; + for (short i = 0; i < kElemsPerFrag; i++) { + ct_c[i] = c_frags[c0][i]; + ct_c[kElemsPerFrag + i] = c_frags[c1][i]; + } + gemm_op.run(ct_a, ct_b, ct_c); + for (short i = 0; i < kElemsPerFrag; i++) { + c_frags[c0][i] = ct_c[i]; + c_frags[c1][i] = ct_c[kElemsPerFrag + i]; + } + } + } + } + (void)compiler_barrier; + } + sg_A += BK; + sg_B += BK; + } + + // Remainder K + int rem_k = int(K) - gemm_k_iters * BK; + for (int kk1 = 0; kk1 < rem_k; kk1 += 16) { + int8_t a_frag[TM][kElemsPerFrag]; + int8_t b_frag[TN][kElemsPerFrag]; + short psk = short(max(0, rem_k - kk1)); + + for (short mm = 0; mm < TM; mm++) { + const device int8_t *ptr = sg_A + kk1 + (sc.y + mm * 16) * K + sc.x; + for (short i = 0; i < 2; i++) + for (short j = 0; j < kElemCols; j++) { + short ki = short(sc.x + j); + a_frag[mm][i * kElemCols + j] = + (ki < psk) ? ptr[(i * kElemRowsJump) * K + j] : int8_t(0); + } + } + for (short nn = 0; nn < TN; nn++) { + const device int8_t *ptr = sg_B + kk1 + (sc.y + nn * 16) * K + sc.x; + for (short i = 0; i < 2; i++) { + for (short j = 0; j < kElemCols; j++) { + short ki = short(sc.x + j); + b_frag[nn][i * kElemCols + j] = + (ki < psk) ? ptr[(i * kElemRowsJump) * K + j] : int8_t(0); + } + } + } + for (short mm = 0; mm < TM; mm++) { + for (short i = 0; i < kElemsPerFrag; i++) { + ct_a[i] = a_frag[mm][i]; + } + for (short i = 0; i < kElemsPerFrag; i++) { + ct_b[i] = b_frag[0][i]; + ct_b[kElemsPerFrag + i] = b_frag[1][i]; + } + short c0 = mm * TN, c1 = c0 + 1; + for (short i = 0; i < kElemsPerFrag; i++) { + ct_c[i] = c_frags[c0][i]; + ct_c[kElemsPerFrag + i] = c_frags[c1][i]; + } + gemm_op.run(ct_a, ct_b, ct_c); + for (short i = 0; i < kElemsPerFrag; i++) { + c_frags[c0][i] = ct_c[i]; + c_frags[c1][i] = ct_c[kElemsPerFrag + i]; + } + } + } + + // Store raw INT32 + device int32_t *D = C + m_base * N + n_base; + for (short mm = 0; mm < TM; mm++) { + for (short nn = 0; nn < TN; nn++) { + nax_frag_store_int32(c_frags[mm * TN + nn], D, int(N), sc, short(mm * 16), + short(nn * 16), M, N, m_base, n_base); + } + } +} + +// ============================================================ +// Kernel entry points — raw INT32 output +// ============================================================ + +kernel void int8_matmul_int32( + const device int8_t *A [[buffer(0)]], const device int8_t *B [[buffer(1)]], + device int32_t *C [[buffer(2)]], constant uint &M [[buffer(3)]], + constant uint &N [[buffer(4)]], constant uint &K [[buffer(5)]], + constant uint &swizzle_log [[buffer(6)]], + constant uint &tiles_m [[buffer(7)]], constant uint &tiles_n [[buffer(8)]], + uint2 tgid [[threadgroup_position_in_grid]], + uint sgid [[simdgroup_index_in_threadgroup]], + uint lid [[thread_index_in_simdgroup]]) { + w8a8_gemm_int32_impl<128, 128, 512, 32, 4, 4>( + A, B, C, M, N, K, swizzle_log, tiles_m, tiles_n, tgid, sgid, lid); +} + +kernel void int8_matmul_int32_small( + const device int8_t *A [[buffer(0)]], const device int8_t *B [[buffer(1)]], + device int32_t *C [[buffer(2)]], constant uint &M [[buffer(3)]], + constant uint &N [[buffer(4)]], constant uint &K [[buffer(5)]], + constant uint &swizzle_log [[buffer(6)]], + constant uint &tiles_m [[buffer(7)]], constant uint &tiles_n [[buffer(8)]], + uint2 tgid [[threadgroup_position_in_grid]], + uint sgid [[simdgroup_index_in_threadgroup]], + uint lid [[thread_index_in_simdgroup]]) { + w8a8_gemm_int32_impl<32, 128, 512, 32, 1, 4>( + A, B, C, M, N, K, swizzle_log, tiles_m, tiles_n, tgid, sgid, lid); +} diff --git a/cider/kernels/w8a8_quantize.metal b/cider/kernels/w8a8_quantize.metal new file mode 100644 index 0000000..8ca335e --- /dev/null +++ b/cider/kernels/w8a8_quantize.metal @@ -0,0 +1,87 @@ +// ============================================================ +// Per-token quantization: FP16 → INT8 + float32 scale +// Target: Apple M5, Metal 4 +// +// Each threadgroup handles one row (one token). +// Threads cooperate to find absmax via simdgroup reduce, +// then quantize in parallel. +// +// Host dispatch: +// threadgroup = (min(256, ceil(K/32)*32), 1, 1) +// grid = (M, 1, 1) +// ============================================================ + +#include +using namespace metal; + +kernel void +quantize_per_token(const device half *X [[buffer(0)]], // [M, K] FP16 input + device int8_t *A [[buffer(1)]], // [M, K] INT8 output + device float *scale [[buffer(2)]], // [M] float32 scale + constant uint &M [[buffer(3)]], + constant uint &K [[buffer(4)]], + uint gid [[threadgroup_position_in_grid]], // row index + uint lid [[thread_index_in_threadgroup]], + uint tg_size [[threads_per_threadgroup]]) { + if (gid >= M) { + return; + } + + const device half *row_in = X + gid * K; + device int8_t *row_out = A + gid * K; + + // Step 1: Find local absmax + float local_max = 0.0f; + for (uint i = lid; i < K; i += tg_size) { + float v = abs(float(row_in[i])); + local_max = max(local_max, v); + } + + // Step 2: Simdgroup reduce max + float sg_max = simd_max(local_max); + + // Step 3: Threadgroup reduce across simdgroups via shared memory + threadgroup float sg_maxes[8]; // up to 8 simdgroups (256/32) + threadgroup float shared_scale; + threadgroup float shared_inv_scale; + uint sg_id = lid / 32; + uint sg_lid = lid % 32; + if (sg_lid == 0) { + sg_maxes[sg_id] = sg_max; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + // Final reduce (first simdgroup only) + if (sg_id == 0) { + float row_max = 0.0f; + uint num_sgs = (tg_size + 31) / 32; + if (sg_lid < num_sgs) { + row_max = sg_maxes[sg_lid]; + } + row_max = simd_max(row_max); + + // Compute and broadcast scale + float s = row_max / 127.0f; + if (s == 0.0f) { + s = 1.0f; + } + + if (sg_lid == 0) { + shared_scale = s; + shared_inv_scale = 1.0f / s; + // Store scale to output + scale[gid] = s; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + // Step 4: All threads read broadcasted scale + float inv_s = shared_inv_scale; + + // Step 5: Quantize + for (uint i = lid; i < K; i += tg_size) { + float v = float(row_in[i]) * inv_s; + v = clamp(round(v), -128.0f, 127.0f); + row_out[i] = int8_t(v); + } +} diff --git a/cider/nn.py b/cider/nn.py new file mode 100644 index 0000000..93e134d --- /dev/null +++ b/cider/nn.py @@ -0,0 +1,241 @@ +"""cider.nn — CiderLinear: unified INT8 kernel, zero double storage. + +Two internal paths (transparent to caller): + - per_channel (gs=0): perchannel_linear (prefill GEMM + decode MV, both per-channel) + - per_group (gs∈{64,128,256}): pergroup_linear (prefill GEMM + decode MV, per-group) + +Both paths: only one copy of int8 weights in memory. +Conversion: from_float() does symmetric requant. + +Usage: + from cider import convert_model + convert_model(model) + # Done. Both prefill and decode use INT8 kernels. +""" + +import mlx.core as mx +import mlx.nn as nn +import numpy as np + +from . import ops + + +# ── Backward compat stubs (no-op) ────────────────────────────── +def set_mode(mode: str): + """No-op. Kept for backward compatibility.""" + pass + +def get_mode() -> str: + """Always returns 'auto'.""" + return "auto" + + +# ── Per-group symmetric quantization helper ──────────────────── +def _symmetric_quantize_pergroup(w_fp: np.ndarray, group_size: int): + """Quantize [N, K] float weights to per-group symmetric INT8. + + Args: + w_fp: [N, K] float32 numpy array + group_size: elements per group (K must be divisible) + + Returns: + w_int8: [N, K] int8 numpy + scale_w: [N, num_groups] float32 numpy + """ + N, K = w_fp.shape + assert K % group_size == 0, f"K={K} not divisible by group_size={group_size}" + num_groups = K // group_size + + w_reshaped = w_fp.reshape(N, num_groups, group_size) + group_max = np.max(np.abs(w_reshaped), axis=2) # [N, num_groups] + scale = group_max / 127.0 + scale = np.where(scale == 0, 1.0, scale) # avoid div by zero + + w_int8 = np.clip( + np.round(w_reshaped / scale[:, :, np.newaxis]), + -128, 127 + ).astype(np.int8).reshape(N, K) + + return w_int8, scale.astype(np.float32) + + +# ── CiderLinear ───────────────────────────────────────────────── +class CiderLinear(nn.Module): + """Unified INT8 Linear: both prefill and decode use custom kernels. + + Internal dispatch (transparent to caller): + per_channel (gs=0): ops.perchannel_linear — per-channel INT8 pipeline + per_group (gs∈{64,128,256}): ops.pergroup_linear — per-group with group scales + + Both paths auto-dispatch GEMM (M>1) vs MV (M==1) internally. + No double weight storage. Only int8 weights + scales in memory. + """ + + def __init__( + self, + w_int8: mx.array, # [N, K] int8 + scale_w: mx.array, # per-channel: [N], per-group: [num_groups, N] (row-major contiguous) + group_size: int, # 0 = per-channel, 64/128/256 = per-group + in_features: int, + out_features: int, + bias: mx.array = None, # [N] float16, default zeros + ): + super().__init__() + self.w_int8 = w_int8 + # Per-group: physically transpose scale_w from [N, num_groups] to [num_groups, N] + # for coalesced SIMD access in Metal kernels. + # scale_w for per-group must already be [num_groups, N] physically contiguous. + # This is ensured by from_float() and convert.py at construction time (numpy transpose). + self.scale_w = scale_w + self.group_size = group_size + self._in_features = in_features + self._out_features = out_features + self.bias = bias if bias is not None else mx.zeros((out_features,), dtype=mx.float16) + + if group_size == 0: + self._mode = "per_channel" + elif group_size in (64, 128, 256): + self._mode = "per_group" + else: + raise ValueError(f"Unsupported group_size={group_size}. Use 0 (per-channel) or 64/128/256.") + + def __call__(self, x: mx.array) -> mx.array: + orig_shape = x.shape + x_2d = x.reshape(-1, self._in_features) + + if self._mode == "per_channel": + y = ops.perchannel_linear(x_2d, self.w_int8, self.scale_w, self.bias) + else: + + y = ops.pergroup_linear(x_2d, self.w_int8, self.scale_w, self.group_size, self.bias) + + y = y.reshape(*orig_shape[:-1], self._out_features) + return y + + @staticmethod + def from_float(layer: nn.Module, target_group_size: int = None, clip_percentile: float = None) -> "CiderLinear": + """Create from nn.Linear or nn.QuantizedLinear. + + For QuantizedLinear (8-bit, gs∈{64,128,256}): dequant → symmetric requant per-group. + For QuantizedLinear (non-8-bit or unsupported gs): dequant → per-channel requant. + For Linear: per-channel requant. + + Args: + layer: Source nn.Linear or nn.QuantizedLinear + target_group_size: Override group_size for conversion. + """ + if isinstance(layer, nn.QuantizedLinear): + bits = layer.bits + gs = layer.group_size + out_f = layer.scales.shape[0] + in_f = layer.scales.shape[1] * gs + lin_bias = getattr(layer, "bias", None) + + # Determine target group_size + if target_group_size is not None: + tgs = target_group_size + elif bits == 8 and gs in (64, 128, 256): + tgs = gs + else: + tgs = 0 # per-channel + + # Dequant + w_fp = mx.dequantize( + layer.weight, layer.scales, + getattr(layer, "biases", None), + gs, bits, + ) + w_np = np.array(w_fp.astype(mx.float32)) + + if tgs == 0: + # Per-channel symmetric requant + w_int8_np, scale_np = ops.quantize_weight_int8(w_np, clip_percentile=clip_percentile) + return CiderLinear( + w_int8=mx.array(w_int8_np), + scale_w=mx.array(scale_np), + group_size=0, + in_features=in_f, + out_features=out_f, + bias=lin_bias, + ) + else: + # Per-group symmetric requant + w_int8_np, scale_np = _symmetric_quantize_pergroup(w_np, tgs) + return CiderLinear( + w_int8=mx.array(w_int8_np), + scale_w=mx.array(scale_np.T.copy()), # [N,ng] -> [ng,N] contiguous + group_size=tgs, + in_features=in_f, + out_features=out_f, + bias=lin_bias, + ) + + elif hasattr(layer, "weight"): + # FP Linear → per-channel + out_f, in_f = layer.weight.shape + lin_bias = getattr(layer, "bias", None) + tgs = target_group_size or 0 + w_np = np.array(layer.weight.astype(mx.float32)) + + if tgs == 0: + w_int8_np, scale_np = ops.quantize_weight_int8(w_np, clip_percentile=clip_percentile) + else: + w_int8_np, scale_np = _symmetric_quantize_pergroup(w_np, tgs) + + return CiderLinear( + w_int8=mx.array(w_int8_np), + scale_w=mx.array(scale_np.T.copy() if tgs > 0 else scale_np), # [N,ng] -> [ng,N] + group_size=tgs, + in_features=in_f, + out_features=out_f, + bias=lin_bias, + ) + else: + raise TypeError(f"Unsupported layer: {type(layer)}") + + @property + def input_dims(self) -> int: + return self._in_features + + @property + def output_dims(self) -> int: + return self._out_features + + def __repr__(self): + return ( + f"CiderLinear(in={self._in_features}, out={self._out_features}, " + f"mode={self._mode}, gs={self.group_size})" + ) + + +# Backward compatibility alias +W8A8Linear = CiderLinear + + +# ── W4A8Linear ────────────────────────────────────────────────── +class W4A8Linear(nn.Module): + """Packed INT4 weight × INT8 activation linear layer.""" + + def __init__(self, packed_weight: mx.array, scale: mx.array, K: int): + super().__init__() + self.packed_weight = packed_weight + self.scale = scale + self._K = K + + def __call__(self, x: mx.array) -> mx.array: + return ops.w4a8_linear(x, self.packed_weight, self.scale) + + @staticmethod + def from_weights(w: np.ndarray, zero_point: int = 8) -> "W4A8Linear": + """Create from FP16/FP32 numpy weight [K, N].""" + K = w.shape[0] + packed, scale = ops.pack_weight_int4(w, zero_point) + return W4A8Linear(mx.array(packed), mx.array(scale), K) + + @property + def input_dims(self) -> int: + return self._K + + @property + def output_dims(self) -> int: + return self.packed_weight.shape[1] diff --git a/cider/ops.py b/cider/ops.py new file mode 100644 index 0000000..c34d407 --- /dev/null +++ b/cider/ops.py @@ -0,0 +1,266 @@ +"""cider.ops — Low-level primitive API for W8A8 / W4A8 / per-group linear. + +These functions return lazy mx.array nodes. Computation happens when +you call mx.eval() — fully compatible with MLX's graph-based execution. +""" + +import re +import subprocess +from pathlib import Path +from typing import Optional + +import mlx.core as mx +import numpy as np + +# ── Kernel directory (shipped with the package) ───────────────── +_KERNEL_DIR: Optional[str] = None + + +def kernel_dir() -> str: + """Return the absolute path to the bundled Metal kernels.""" + global _KERNEL_DIR + if _KERNEL_DIR is None: + _KERNEL_DIR = str(Path(__file__).parent / "kernels") + return _KERNEL_DIR + + +# ── Extension loader ──────────────────────────────────────────── +_ext = None + + +def _load_ext(): + global _ext + if _ext is not None: + return _ext + import sys + lib_dir = str(Path(__file__).parent / "lib") + if lib_dir not in sys.path: + sys.path.insert(0, lib_dir) + try: + import _cider_prim + _ext = _cider_prim + return _ext + except ImportError: + raise RuntimeError( + "Cider C++ extension not available. INT8 TensorOps require Apple M5+. " + "On M4 and below, use standard MLX inference instead." + ) + + +# ── Hardware detection ────────────────────────────────────────── + +def is_available() -> bool: + """Check if INT8 TensorOps are available (Apple M5+, Metal 4).""" + try: + chip = subprocess.run( + ["sysctl", "-n", "machdep.cpu.brand_string"], + capture_output=True, text=True, timeout=5 + ).stdout.strip() + except Exception: + return False + m = re.match(r"Apple M(\d+)", chip) + if not m or int(m.group(1)) < 5: + return False + try: + _load_ext() + return True + except (ImportError, RuntimeError, Exception): + return False + + +# ── Weight quantization helpers ───────────────────────────────── + +def quantize_weight_int8( + w: np.ndarray, + clip_percentile: float = None, +) -> tuple: + """Quantize FP16/FP32 weights to per-row symmetric INT8. + + Args: + w: Weight matrix [N, K] as numpy array (N=out_features, K=in_features). + clip_percentile: If set (e.g. 99.9), use per-row percentile instead of + absmax to compute scale. Clips outliers before quantization, improving + precision for layers with extreme weight outliers. Default None (absmax). + + Returns: + (w_int8, scale_w) where w_int8 is [N, K] int8 and + scale_w is [N] float32 (one scale per output channel). + """ + w = w.astype(np.float32) + if clip_percentile is not None: + row_clip = np.percentile(np.abs(w), clip_percentile, axis=1) # [N] + row_clip = np.where(row_clip == 0, np.max(np.abs(w), axis=1), row_clip) + scale = row_clip / 127.0 + scale = np.where(scale == 0, 1.0, scale) + w = np.clip(w, -row_clip[:, np.newaxis], row_clip[:, np.newaxis]) + else: + row_max = np.max(np.abs(w), axis=1) # [N] + scale = row_max / 127.0 + scale = np.where(scale == 0, 1.0, scale) + w_int8 = np.clip(np.round(w / scale[:, np.newaxis]), -128, 127).astype(np.int8) + return w_int8, scale.astype(np.float32) + + +def pack_weight_int4( + w: np.ndarray, + zero_point: int = 8, +) -> tuple: + """Quantize FP16/FP32 weights to packed INT4 (symmetric, per-column). + + Args: + w: Weight matrix [K, N] as numpy array. K must be even. + zero_point: INT4 zero point (default 8 for signed range [-8, 7]). + + Returns: + (packed_w, scale_w) where packed_w is [K//2, N] uint8 + (high nibble = even k, low nibble = odd k) and + scale_w is [N] float32. + """ + K, N = w.shape + assert K % 2 == 0, f"K must be even, got {K}" + w = w.astype(np.float32) + col_max = np.max(np.abs(w), axis=0) + scale = col_max / 7.0 + scale = np.where(scale == 0, 1.0, scale) + w_q = np.clip(np.round(w / scale[np.newaxis, :]) + zero_point, 0, 15).astype(np.uint8) + packed = (w_q[0::2, :] << 4) | w_q[1::2, :] + return packed, scale.astype(np.float32) + + +# ── Primitive API ─────────────────────────────────────────────── + +def perchannel_linear( + x: mx.array, + w: mx.array, + scale_w: mx.array, + bias: Optional[mx.array] = None, + stream: Optional[mx.Stream] = None, +) -> mx.array: + """W8A8 per-channel quantized linear: y = dequant(quant_a(x) @ w_int8) + bias. + + Args: + x: Input activations [M, K] float16 or bfloat16. + w: INT8 weights [N, K] int8 (per-row quantized). + scale_w: Per-row weight scales [N] float32. + stream: Optional MLX stream. + + Returns: + Output [M, N] matching input dtype. + """ + ext = _load_ext() + out_dtype = x.dtype + kw = {"stream": stream} if stream is not None else {} + N = w.shape[0] + if bias is None: + bias = mx.zeros((N,), dtype=mx.float16) + result = ext.perchannel_linear(x, w, scale_w, bias, kernel_dir(), **kw) + if out_dtype != mx.float16: + result = result.astype(out_dtype, **kw) + return result + + +# Shared placeholder for new_bias (V5 kernel ignores it; Metal needs valid buffer) +_shared_new_bias_cache = {} + +def _get_shared_new_bias_placeholder(N: int, num_groups: int): + key = (N, num_groups) + if key not in _shared_new_bias_cache: + _shared_new_bias_cache[key] = mx.zeros((N, num_groups), dtype=mx.float32) + return _shared_new_bias_cache[key] + + +def pergroup_linear( + x: mx.array, + w: mx.array, + scale_w: mx.array, + group_size: int, + bias: Optional[mx.array] = None, + new_bias: Optional[mx.array] = None, + stream: Optional[mx.Stream] = None, +) -> mx.array: + """ + mlx native quantize format asymmetric affine + quantize: q = clip(round((w - biases) / scales), 0, 2^b - 1), b = bits + dequantize: w = q*scale + bias + """ + """Per-group INT8 linear with optional bias. + + Dispatches internally: + M > 1 → per-group GEMM (activation quantize + INT8 TensorOps) + M == 1 → per-group MV (FP activation, weight dequant on-the-fly) + + Args: + x: Input activations [M, K] float16 or bfloat16. + w: INT8 weights [N, K] int8 (per-group symmetric quantized). + scale_w: Per-group weight scales [num_groups, N] float32 (physically transposed for coalesced GPU access). + group_size: Group size (64, 128, or 256). + bias: Optional bias [N] float16. Default zeros. + stream: Optional MLX stream. + + Returns: + Output [M, N] matching input dtype. + """ + ext = _load_ext() + N = w.shape[0] + num_groups = scale_w.shape[0] if scale_w.ndim == 2 else 1 # scale_w is [num_groups, N] + if bias is None: + bias = mx.zeros((N,), dtype=mx.float16) + if new_bias is None: + # V5 kernel ignores new_bias (symmetric quantization), but Metal + # requires a valid buffer binding. Use a tiny shared placeholder + # instead of allocating (N, num_groups) every forward call. + new_bias = _get_shared_new_bias_placeholder(N, num_groups) + out_dtype = x.dtype + kw = {"stream": stream} if stream is not None else {} + # scale_w layout: [num_groups, N] physically contiguous. Kernel indexes as scale_w[g * N + n]. + + result = ext.pergroup_linear(x, w, scale_w, bias, new_bias, group_size, kernel_dir(), **kw) + if out_dtype != mx.float16: + result = result.astype(out_dtype, **kw) + return result + + +def w4a8_linear( + x: mx.array, + packed_w: mx.array, + scale_w: mx.array, + stream: Optional[mx.Stream] = None, +) -> mx.array: + """W4A8 quantized linear: y = dequant(quant_a(x) @ unpack4(w)). + + Args: + x: Input activations [M, K] float16. + packed_w: Packed INT4 weights [K//2, N] uint8. + scale_w: Per-column weight scales [N] float32. + stream: Optional MLX stream. + + Returns: + Output [M, N] float16. + """ + ext = _load_ext() + out_dtype = x.dtype + kw = {"stream": stream} if stream is not None else {} + result = ext.w4a8_linear(x, packed_w, scale_w, kernel_dir(), **kw) + if out_dtype != mx.float16: + result = result.astype(out_dtype, **kw) + return result + + +def int8_matmul_int32( + a: mx.array, + b: mx.array, + stream=None, +) -> mx.array: + """Raw INT8×INT8→INT32 matmul (bit-exact, no dequant). + + Args: + a: INT8 matrix [M, K]. + b: INT8 matrix [N, K] (transposed weight layout). + stream: Optional MLX stream. + + Returns: + INT32 result [M, N]. + """ + ext = _load_ext() + kw = {"stream": stream} if stream is not None else {} + return ext.int8_matmul_int32(a, b, kernel_dir(), **kw) diff --git a/config/config.yaml b/config/config.yaml new file mode 100644 index 0000000..7292758 --- /dev/null +++ b/config/config.yaml @@ -0,0 +1,14 @@ +model_name_or_path: /path/to/model +sampling: + max_new_tokens: 1024 + repetition_penalty: 1.0 + temperature: 1.0 + top_k: null + top_p: 1.0 +server: + host: 0.0.0.0 + max_image_buffer_size: 2 + port: 8341 + ttl: 1800 +w8a8: + mode: 'off' diff --git a/csrc/include/pergroup_primitive.h b/csrc/include/pergroup_primitive.h new file mode 100644 index 0000000..151438b --- /dev/null +++ b/csrc/include/pergroup_primitive.h @@ -0,0 +1,46 @@ +// Per-group INT8 GEMM/MV primitive header — symmetric quant (no bias) +#pragma once + +#include "mlx/ops.h" +#include "mlx/primitives.h" + +#include + +namespace cider { + +namespace mx = mlx::core; + +// Per-group INT8 GEMM (prefill) / MV (decode) — symmetric quantization +// Inputs: +// x: [M, K] float16/bfloat16 — activation +// w: [N, K] int8 — per-group symmetric quantized weight +// scale_w: [N, num_groups] float32 — per-group weight scales +// group_size: 64, 128, or 256 +class PerGroupLinear : public mx::Primitive { +public: + PerGroupLinear(mx::Stream s, const std::string &kernel_dir, int group_size) + : mx::Primitive(s), kernel_dir_(kernel_dir), group_size_(group_size) {} + + void eval_cpu(const std::vector &, + std::vector &) override { + throw std::runtime_error("PerGroupLinear: CPU not supported"); + } + void eval_gpu(const std::vector &inputs, + std::vector &outputs) override; + + const char *name() const override { return "PerGroupLinear"; } + bool is_equivalent(const mx::Primitive &other) const override { return true; } + +private: + std::string kernel_dir_; + int group_size_; +}; + +// Python-facing function +mx::array pergroup_linear(const mx::array &x, const mx::array &w, + const mx::array &scale_w, const mx::array &bias, + const mx::array &new_bias, + int group_size, const std::string &kernel_dir, + mx::StreamOrDevice s = {}); + +} // namespace cider diff --git a/csrc/include/sdpa_primitive.h b/csrc/include/sdpa_primitive.h new file mode 100644 index 0000000..4fe8142 --- /dev/null +++ b/csrc/include/sdpa_primitive.h @@ -0,0 +1,143 @@ +// Cider SDPA — Custom Primitive for v9 optimized attention +// +// Dispatches 1-pass or 2-pass SDPA kernels via MLX's CommandEncoder. +// All kernel selection (1pass vs 2pass, blocks value) is decided by the +// Python caller and passed as arguments. + +#pragma once + +#include "mlx/ops.h" +#include "mlx/primitives.h" + +#include + +namespace cider { + +namespace mx = mlx::core; + +// ── 1-pass SDPA Primitive ──────────────────────────────────────── +// Inputs: [Q(B*H, D), K(...), V(...)] +// Output: [O(B*H, D)] +class SDPAVector : public mx::Primitive { + public: + SDPAVector(mx::Stream stream, + const std::string& kernel_dir, + int gqa_factor, + int N, + size_t k_head_stride, + size_t k_seq_stride, + size_t v_head_stride, + size_t v_seq_stride, + float scale) + : mx::Primitive(stream), + kernel_dir_(kernel_dir), + gqa_factor_(gqa_factor), + N_(N), + k_head_stride_(k_head_stride), + k_seq_stride_(k_seq_stride), + v_head_stride_(v_head_stride), + v_seq_stride_(v_seq_stride), + scale_(scale) {} + + void eval_cpu(const std::vector& inputs, + std::vector& outputs) override { + throw std::runtime_error("SDPAVector: CPU not supported"); + } + + void eval_gpu(const std::vector& inputs, + std::vector& outputs) override; + + const char* name() const override { return "CiderSDPAVector"; } + bool is_equivalent(const mx::Primitive& other) const override { return true; } + + private: + std::string kernel_dir_; + int gqa_factor_; + int N_; + size_t k_head_stride_; + size_t k_seq_stride_; + size_t v_head_stride_; + size_t v_seq_stride_; + float scale_; +}; + +// ── 2-pass SDPA Primitive ──────────────────────────────────────── +// Inputs: [Q(B*H, D), K(...), V(...)] +// Output: [O(B*H, D)] +// Internally allocates partials/sums/maxs as temporaries. +class SDPAVector2Pass : public mx::Primitive { + public: + SDPAVector2Pass(mx::Stream stream, + const std::string& kernel_dir, + int gqa_factor, + int N, + int blocks, + int num_kv_heads, + int batch_size, + size_t k_head_stride, + size_t k_seq_stride, + size_t v_head_stride, + size_t v_seq_stride, + float scale) + : mx::Primitive(stream), + kernel_dir_(kernel_dir), + gqa_factor_(gqa_factor), + N_(N), + blocks_(blocks), + num_kv_heads_(num_kv_heads), + batch_size_(batch_size), + k_head_stride_(k_head_stride), + k_seq_stride_(k_seq_stride), + v_head_stride_(v_head_stride), + v_seq_stride_(v_seq_stride), + scale_(scale) {} + + void eval_cpu(const std::vector& inputs, + std::vector& outputs) override { + throw std::runtime_error("SDPAVector2Pass: CPU not supported"); + } + + void eval_gpu(const std::vector& inputs, + std::vector& outputs) override; + + const char* name() const override { return "CiderSDPAVector2Pass"; } + bool is_equivalent(const mx::Primitive& other) const override { return true; } + + private: + std::string kernel_dir_; + int gqa_factor_; + int N_; + int blocks_; + int num_kv_heads_; + int batch_size_; + size_t k_head_stride_; + size_t k_seq_stride_; + size_t v_head_stride_; + size_t v_seq_stride_; + float scale_; +}; + +// ── Public API ─────────────────────────────────────────────────── + +// 1-pass SDPA (short sequences) +mx::array cider_sdpa_1pass( + const mx::array& queries, // [B*H, D] + const mx::array& keys, // [B*Hkv, N, D] + const mx::array& values, // [B*Hkv, N, D] + int gqa_factor, + float scale, + const std::string& kernel_dir, + mx::StreamOrDevice s = {}); + +// 2-pass SDPA (long sequences) +mx::array cider_sdpa_2pass( + const mx::array& queries, // [B, H, 1, D] + const mx::array& keys, // [B, Hkv, N, D] + const mx::array& values, // [B, Hkv, N, D] + int gqa_factor, + int blocks, + float scale, + const std::string& kernel_dir, + mx::StreamOrDevice s = {}); + +} // namespace cider diff --git a/csrc/include/w4a8_primitive.h b/csrc/include/w4a8_primitive.h new file mode 100644 index 0000000..3cea058 --- /dev/null +++ b/csrc/include/w4a8_primitive.h @@ -0,0 +1,55 @@ +// W4A8 Linear as mlx Custom Primitive. +// Packed INT4 weights × INT8 activations via TensorOps. + +#pragma once + +#include "mlx/ops.h" +#include "mlx/primitives.h" + +#include + +namespace cider { + +namespace mx = mlx::core; + +// Inputs: [x(M,K) float16, packed_w(K/2,N) uint8, scale_w(N) float32] +// Output: [y(M,N) float16] +// +// Weight layout: [K/2, N] uint8 — packed INT4 symmetric (zero_point=8) +// high nibble = even k index, low nibble = odd k index +// scale_w: per-column scale (includes group scale pre-folded) +class W4A8Linear : public mx::Primitive { + public: + explicit W4A8Linear(mx::Stream stream, const std::string& kernel_dir) + : mx::Primitive(stream), kernel_dir_(kernel_dir) {} + + void eval_cpu( + const std::vector& inputs, + std::vector& outputs) override { + throw std::runtime_error("W4A8Linear: CPU not supported"); + } + + void eval_gpu( + const std::vector& inputs, + std::vector& outputs) override; + + const char* name() const override { + return "W4A8Linear"; + } + + bool is_equivalent(const mx::Primitive& other) const override { + return true; + } + + private: + std::string kernel_dir_; +}; + +mx::array w4a8_linear( + const mx::array& x, // [M, K] float16 + const mx::array& packed_w, // [K/2, N] uint8 (packed INT4) + const mx::array& scale_w, // [N] float32 + const std::string& kernel_dir, + mx::StreamOrDevice s = {}); + +} // namespace cider diff --git a/csrc/include/w8a8_primitive.h b/csrc/include/w8a8_primitive.h new file mode 100644 index 0000000..24d933e --- /dev/null +++ b/csrc/include/w8a8_primitive.h @@ -0,0 +1,102 @@ +// W8A8 Linear as mlx Custom Primitive. +// Dispatches quantize + INT8 matmul (prefill) or FP MV (decode) via mlx's CommandEncoder. + +#pragma once + +#include "mlx/ops.h" +#include "mlx/primitives.h" + +#include +#include +#include + +namespace MTL { +class ComputePipelineState; +} + +namespace cider { + +namespace mx = mlx::core; + +// ── Custom Primitive ───────────────────────────────────────────── +// Inputs: [x(M,K) float16, w(N,K) int8, scale_w(N) float32, bias(N) float16] +// Output: [y(M,N) float16] +// +// M > 1: quantize activation + INT8 GEMM (prefill) +// M == 1: FP activation × dequant weight MV (decode) +class W8A8Linear : public mx::Primitive { + public: + explicit W8A8Linear(mx::Stream stream, const std::string& kernel_dir) + : mx::Primitive(stream), kernel_dir_(kernel_dir) {} + + void eval_cpu( + const std::vector& inputs, + std::vector& outputs) override { + throw std::runtime_error("W8A8Linear: CPU not supported"); + } + + void eval_gpu( + const std::vector& inputs, + std::vector& outputs) override; + + const char* name() const override { + return "W8A8Linear"; + } + + bool is_equivalent(const mx::Primitive& other) const override { + return true; + } + + private: + std::string kernel_dir_; +}; + + +// ── Raw INT32 Matmul Primitive ─────────────────────────────────── +// Inputs: [A(M,K) int8, B(K,N) int8] +// Output: [C(M,N) int32] (bit-exact, no dequant) +class Int8MatMulInt32 : public mx::Primitive { + public: + explicit Int8MatMulInt32(mx::Stream stream, const std::string& kernel_dir) + : mx::Primitive(stream), kernel_dir_(kernel_dir) {} + + void eval_cpu( + const std::vector& inputs, + std::vector& outputs) override { + throw std::runtime_error("Int8MatMulInt32: CPU not supported"); + } + + void eval_gpu( + const std::vector& inputs, + std::vector& outputs) override; + + const char* name() const override { + return "Int8MatMulInt32"; + } + + bool is_equivalent(const mx::Primitive& other) const override { + return true; + } + + private: + std::string kernel_dir_; +}; + +// ── Public API ─────────────────────────────────────────────────── +mx::array perchannel_linear( + const mx::array& x, // [M, K] float16 + const mx::array& w, // [N, K] int8 + const mx::array& scale_w, // [N] float32 + const mx::array& bias, // [N] float16 + const std::string& kernel_dir, + mx::StreamOrDevice s = {}); + + +// ── Raw INT32 matmul (for bit-exact testing) ───────────────────── +mx::array int8_matmul_int32( + const mx::array& a, // [M, K] int8 + const mx::array& b, // [K, N] int8 + const std::string& kernel_dir, + mx::StreamOrDevice s = {}); + +} // namespace cider diff --git a/csrc/src/pergroup_primitive.mm b/csrc/src/pergroup_primitive.mm new file mode 100644 index 0000000..e9086c1 --- /dev/null +++ b/csrc/src/pergroup_primitive.mm @@ -0,0 +1,314 @@ +// Per-group INT8 GEMM primitive implementation +// +// Dispatches: +// 1. quantize_per_token: x[M,K] float16 → a_int8[M,K] int8 + scale_a[M] +// float32 +// 2. pergroup_int8_gemm_gXX: A[M,K] int8 × B[N,K] int8 → C[M,N] float16 +// with per-group scale_w[N, num_groups] and per-token scale_a[M] + +#include "pergroup_primitive.h" +#include "mlx/backend/metal/device.h" +#include "mlx/backend/metal/utils.h" + +#import +#import + +#include +#include +#include +#include + +namespace cider { + +using namespace mlx::core; + +// ── Pipeline cache for per-group kernels ──────────────────────── +class PerGroupPipelineCache { +public: + static PerGroupPipelineCache &instance() { + static PerGroupPipelineCache cache; + return cache; + } + + void ensure_init(const std::string &kernel_dir) { + std::lock_guard lock(mutex_); + if (initialized_ && kernel_dir_ == kernel_dir) + return; + auto &dev = metal::device(mlx::core::Device::gpu); + auto *mtl_device = dev.mtl_device(); + pergroup_lib_ = + compile_source(mtl_device, kernel_dir + "/pergroup_int8_gemm.metal"); + mv_lib_ = + compile_source(mtl_device, kernel_dir + "/pergroup_int8_mv.metal"); + quantize_lib_ = + compile_source(mtl_device, kernel_dir + "/w8a8_quantize.metal"); + kernel_dir_ = kernel_dir; + pipelines_.clear(); + initialized_ = true; + } + + MTL::ComputePipelineState *get(const std::string &name) { + std::lock_guard lock(mutex_); + auto it = pipelines_.find(name); + if (it != pipelines_.end()) { + return it->second; + } + auto &dev = metal::device(mlx::core::Device::gpu); + auto *pso = make_pipeline(dev.mtl_device(), name); + pipelines_[name] = pso; + return pso; + } + +private: + PerGroupPipelineCache() = default; + bool initialized_ = false; + std::string kernel_dir_; + std::unordered_map pipelines_; + MTL::Library *pergroup_lib_ = nullptr; + MTL::Library *mv_lib_ = nullptr; + MTL::Library *quantize_lib_ = nullptr; + std::mutex mutex_; + + static std::string read_file(const std::string &path) { + std::ifstream f(path); + if (!f.is_open()) { + throw std::runtime_error("Cannot open: " + path); + } + std::stringstream ss; + ss << f.rdbuf(); + return ss.str(); + } + + MTL::Library *compile_source(MTL::Device *mtl_device, + const std::string &source_path) { + std::string source = read_file(source_path); + @autoreleasepool { + NSString *src = [NSString stringWithUTF8String:source.c_str()]; + MTLCompileOptions *opts = [[MTLCompileOptions alloc] init]; + opts.languageVersion = MTLLanguageVersion4_0; + NSError *error = nil; + id device_objc = (__bridge id)mtl_device; + id lib_objc = [device_objc newLibraryWithSource:src + options:opts + error:&error]; + if (!lib_objc) { + std::string err = + error ? [[error localizedDescription] UTF8String] : "Unknown error"; + throw std::runtime_error("Metal compile failed (" + source_path + + "): " + err); + } + return (__bridge MTL::Library *)(void *)CFBridgingRetain(lib_objc); + } + } + + MTL::ComputePipelineState *make_pipeline(MTL::Device *mtl_device, + const std::string &name) { + @autoreleasepool { + NSString *fn_name = [NSString stringWithUTF8String:name.c_str()]; + // Try pergroup lib, then mv lib, then quantize lib + id func = nil; + id lib_objc = (__bridge id)pergroup_lib_; + func = [lib_objc newFunctionWithName:fn_name]; + if (!func) { + lib_objc = (__bridge id)mv_lib_; + func = [lib_objc newFunctionWithName:fn_name]; + } + if (!func) { + lib_objc = (__bridge id)quantize_lib_; + func = [lib_objc newFunctionWithName:fn_name]; + } + if (!func) { + throw std::runtime_error("Kernel not found: " + name); + } + NSError *error = nil; + id device_objc = (__bridge id)mtl_device; + id pso = + [device_objc newComputePipelineStateWithFunction:func error:&error]; + if (!pso) { + std::string err = + error ? [[error localizedDescription] UTF8String] : "Unknown error"; + throw std::runtime_error("Pipeline failed for " + name + ": " + err); + } + return (__bridge MTL::ComputePipelineState *)(void *)CFBridgingRetain( + pso); + } + } +}; + +void PerGroupLinear::eval_gpu(const std::vector &inputs, + std::vector &outputs) { + auto &x = inputs[0]; // [M, K] float16 + auto &w = inputs[1]; // [N, K] int8 + auto &scale_w = inputs[2]; // [N, num_groups] float32 + auto &bias = inputs[3]; // [N] float16 + auto &new_bias = inputs[4]; // [N, num_groups] float32 (asymmetric correction) + auto &out = outputs[0]; // [M, N] float16 + + uint32_t M = static_cast(x.shape(0)); + uint32_t N = static_cast(w.shape(0)); + uint32_t K = static_cast(w.shape(1)); + + out.set_data(allocator::malloc(out.nbytes())); + + auto &cache = PerGroupPipelineCache::instance(); + cache.ensure_init(kernel_dir_); + + auto &s = stream(); + auto &enc = metal::get_command_encoder(s); + + if (M == 1) { + // ── Decode path: per-group MV kernel ── + // No activation quantization needed; directly use FP16 activation + std::string kname; + if (group_size_ == 64) { + kname = "pergroup_int8_mv_g64"; + } else if (group_size_ == 128) { + kname = "pergroup_int8_mv_g128"; + } else { + kname = "pergroup_int8_mv_g256"; + } + + auto *pso = cache.get(kname); + enc.set_compute_pipeline_state(pso); + + constexpr uint32_t TOTAL_ROWS = 8; // NUM_SIMDGROUPS(2) * RESULTS_PER_SG(4) + uint32_t threadgroups = (N + TOTAL_ROWS - 1) / TOTAL_ROWS; + uint32_t threads = 64; // 2 simdgroups x 32 threads + + enc.set_input_array(x, 0); // [1, K] float16 → pass as [K] + enc.set_input_array(w, 1); // [N, K] int8 + enc.set_output_array(out, 2); // [1, N] float16 → write as [N] + enc.set_input_array(scale_w, 3); // [N, num_groups] float32 + enc.set_bytes(N, 4); + enc.set_bytes(K, 5); + enc.set_input_array(bias, 6); // [N] float16 + enc.set_input_array(new_bias, 7); // [N, num_groups] float32 (correction) + + enc.dispatch_threadgroups(MTL::Size::Make(threadgroups, 1, 1), + MTL::Size::Make(threads, 1, 1)); + } else { + // ── Prefill path: per-group GEMM with activation quantization ── + + // Scratch: activation quantization + size_t a_bytes = static_cast(M) * K; + size_t sa_bytes = static_cast(M) * sizeof(float); + + array a_int8({static_cast(M), static_cast(K)}, int8, nullptr, {}); + a_int8.set_data(allocator::malloc(a_bytes)); + + array sa({static_cast(M)}, float32, nullptr, {}); + sa.set_data(allocator::malloc(sa_bytes)); + + // DISPATCH 1: quantize_per_token + { + auto *pso = cache.get("quantize_per_token"); + enc.set_compute_pipeline_state(pso); + enc.set_input_array(x, 0); + enc.set_output_array(a_int8, 1); + enc.set_output_array(sa, 2); + enc.set_bytes(M, 3); + enc.set_bytes(K, 4); + uint32_t tg = std::min(256u, std::max(32u, ((K + 31) / 32) * 32)); + enc.dispatch_threadgroups(MTL::Size::Make(M, 1, 1), + MTL::Size::Make(tg, 1, 1)); + } + + enc.barrier(); + + // DISPATCH 2: per-group GEMM + { + std::string kname; + bool use_small = (M <= 64); + uint32_t BM, BN, threads; + + if (group_size_ == 64) { + kname = use_small ? "pergroup_int8_gemm_g64_small" + : "pergroup_int8_gemm_g64"; + } else if (group_size_ == 128) { + kname = use_small ? "pergroup_int8_gemm_g128_small" + : "pergroup_int8_gemm_g128"; + } else { + kname = use_small ? "pergroup_int8_gemm_g256_small" + : "pergroup_int8_gemm_g256"; + } + + BM = use_small ? 32 : 128; + BN = 128; + threads = use_small ? 128 : 512; + + auto *pso = cache.get(kname); + enc.set_compute_pipeline_state(pso); + + uint32_t tiles_n = (N + BN - 1) / BN; + uint32_t tiles_m = (M + BM - 1) / BM; + uint32_t swizzle_log; + if (tiles_m <= 3) { + swizzle_log = 0; + } else if (tiles_m <= 6) { + swizzle_log = 1; + } else { + swizzle_log = 2; + } + + uint32_t tile = 1u << swizzle_log; + uint32_t grid_x = tiles_n * tile; + uint32_t grid_y = (tiles_m + tile - 1) / tile; + + enc.set_input_array(a_int8, 0); + enc.set_input_array(w, 1); + enc.set_output_array(out, 2); + enc.set_bytes(M, 3); + enc.set_bytes(N, 4); + enc.set_bytes(K, 5); + enc.set_input_array(sa, 6); + enc.set_input_array(scale_w, 7); + enc.set_bytes(swizzle_log, 8); + enc.set_bytes(tiles_m, 9); + enc.set_bytes(tiles_n, 10); + enc.set_input_array(bias, 11); + + enc.dispatch_threadgroups(MTL::Size::Make(grid_x, grid_y, 1), + MTL::Size::Make(threads, 1, 1)); + } + + enc.add_temporary(a_int8); + enc.add_temporary(sa); + } +} + +array pergroup_linear(const array &x, const array &w, const array &scale_w, + const array &bias, const array &new_bias, int group_size, + const std::string &kernel_dir, StreamOrDevice s) { + if (x.ndim() != 2) { + throw std::invalid_argument("pergroup_linear: x must be 2D [M,K]"); + } + if (w.ndim() != 2) { + throw std::invalid_argument("pergroup_linear: w must be 2D [N,K]"); + } + if (scale_w.ndim() != 2) { + throw std::invalid_argument( + "pergroup_linear: scale_w must be 2D [N,num_groups]"); + } + if (group_size != 64 && group_size != 128 && group_size != 256) { + throw std::invalid_argument( + "pergroup_linear: group_size must be 64, 128, or 256"); + } + + int M = x.shape(0); + int N = w.shape(0); + auto stream = to_stream(s); + + auto result = + array({M, N}, float16, + std::make_shared(stream, kernel_dir, group_size), + {astype(x, float16, stream), astype(w, int8, stream), + astype(scale_w, float32, stream), astype(bias, float16, stream), + astype(new_bias, float32, stream)}); + + if (x.dtype() == bfloat16) { + return astype(result, bfloat16, stream); + } + return result; +} + +} // namespace cider diff --git a/csrc/src/prim_bindings.cpp b/csrc/src/prim_bindings.cpp new file mode 100644 index 0000000..4c268ba --- /dev/null +++ b/csrc/src/prim_bindings.cpp @@ -0,0 +1,49 @@ +#include +#include +#include +#include + +#include "mlx/ops.h" +#include "pergroup_primitive.h" +#include "sdpa_primitive.h" +#include "w4a8_primitive.h" +#include "w8a8_primitive.h" + +namespace nb = nanobind; +using namespace nb::literals; +namespace mx = mlx::core; + +NB_MODULE(_cider_prim, m) { + m.doc() = "cider: W8A8 + W4A8 INT8 TensorOps + SDPA primitives for Apple M5+"; + + m.def("perchannel_linear", &cider::perchannel_linear, "x"_a, "w"_a, "scale_w"_a, + "bias"_a, "kernel_dir"_a, nb::kw_only(), "stream"_a = nb::none(), + "W8A8 quantized linear: y = dequant(quant_a(x) @ w_int8) + bias"); + + m.def("w4a8_linear", &cider::w4a8_linear, "x"_a, "packed_w"_a, "scale_w"_a, + "kernel_dir"_a, nb::kw_only(), "stream"_a = nb::none(), + "W4A8 quantized linear: y = dequant(quant_a(x) @ unpack4(w))"); + + m.def("int8_matmul_int32", &cider::int8_matmul_int32, "a"_a, "b"_a, + "kernel_dir"_a, nb::kw_only(), "stream"_a = nb::none(), + "Raw INT8xINT8->INT32 matmul (bit-exact, no dequant)"); + + m.def("pergroup_linear", &cider::pergroup_linear, "x"_a, "w"_a, "scale_w"_a, + "bias"_a, "new_bias"_a, "group_size"_a, "kernel_dir"_a, nb::kw_only(), + "stream"_a = nb::none(), + "Per-group INT8 linear with bias: prefill GEMM or decode MV with " + "per-group scales"); + + // ── SDPA ── + m.def("cider_sdpa_1pass", &cider::cider_sdpa_1pass, + "queries"_a, "keys"_a, "values"_a, + "gqa_factor"_a, "scale"_a, "kernel_dir"_a, + nb::kw_only(), "stream"_a = nb::none(), + "Cider v9 SDPA 1-pass (short sequences)"); + + m.def("cider_sdpa_2pass", &cider::cider_sdpa_2pass, + "queries"_a, "keys"_a, "values"_a, + "gqa_factor"_a, "blocks"_a, "scale"_a, "kernel_dir"_a, + nb::kw_only(), "stream"_a = nb::none(), + "Cider v9 SDPA 2-pass with contiguous chunks + TILE=4 tiling"); +} diff --git a/csrc/src/sdpa_primitive.mm b/csrc/src/sdpa_primitive.mm new file mode 100644 index 0000000..b7ba6a1 --- /dev/null +++ b/csrc/src/sdpa_primitive.mm @@ -0,0 +1,344 @@ +// sdpa_primitive.mm — Cider v9 SDPA Custom Primitive implementation +// +// Dispatches precompiled Metal kernels for optimized attention. +// Two primitives: SDPAVector (1-pass) and SDPAVector2Pass (2-pass). + +#include "sdpa_primitive.h" +#include "mlx/backend/metal/device.h" +#include "mlx/backend/metal/utils.h" + +#import +#import + +#include +#include +#include +#include + +namespace cider { + +using namespace mlx::core; + +// ── Pipeline cache for SDPA kernels ───────────────────────────── +class SDPAPipelineCache { +public: + static SDPAPipelineCache &instance() { + static SDPAPipelineCache cache; + return cache; + } + + void ensure_init(const std::string &kernel_dir) { + std::lock_guard lock(mutex_); + if (initialized_ && kernel_dir_ == kernel_dir) + return; + + auto &dev = metal::device(mlx::core::Device::gpu); + auto *mtl_device = dev.mtl_device(); + + // Read and compile the SDPA metal source + std::string source = read_file(kernel_dir + "/cider_sdpa_vector.metal"); + + // Prepend the include path — the header is in the same directory + // We inline the header by reading it and prepending + std::string header = read_file(kernel_dir + "/cider_sdpa_vector.h"); + + // Replace the #include directive with the actual header content + std::string combined; + combined.reserve(header.size() + source.size() + 256); + // Add metal_stdlib include + combined += "#include \n"; + combined += "#include \n"; + combined += "using namespace metal;\n\n"; + combined += "// === Inlined cider_sdpa_vector.h ===\n"; + combined += header; + combined += "\n// === End inlined header ===\n\n"; + + // Append the .metal source but strip its own #include lines + std::istringstream iss(source); + std::string line; + while (std::getline(iss, line)) { + if (line.find("#include") != std::string::npos) + continue; + if (line.find("using namespace metal") != std::string::npos) + continue; + combined += line + "\n"; + } + + lib_ = compile_source(mtl_device, combined); + kernel_dir_ = kernel_dir; + pipelines_.clear(); + initialized_ = true; + } + + MTL::ComputePipelineState *get(const std::string &name) { + std::lock_guard lock(mutex_); + auto it = pipelines_.find(name); + if (it != pipelines_.end()) + return it->second; + + auto &dev = metal::device(mlx::core::Device::gpu); + auto *pso = make_pipeline(dev.mtl_device(), name); + pipelines_[name] = pso; + return pso; + } + +private: + SDPAPipelineCache() = default; + bool initialized_ = false; + std::string kernel_dir_; + std::unordered_map pipelines_; + MTL::Library *lib_ = nullptr; + std::mutex mutex_; + + static std::string read_file(const std::string &path) { + std::ifstream f(path); + if (!f.is_open()) { + throw std::runtime_error("Cannot open: " + path); + } + std::stringstream ss; + ss << f.rdbuf(); + return ss.str(); + } + + MTL::Library *compile_source(MTL::Device *mtl_device, + const std::string &source) { + @autoreleasepool { + NSString *src = [NSString stringWithUTF8String:source.c_str()]; + MTLCompileOptions *opts = [[MTLCompileOptions alloc] init]; + opts.languageVersion = MTLLanguageVersion3_1; + NSError *error = nil; + id device_objc = (__bridge id)mtl_device; + id lib_objc = [device_objc newLibraryWithSource:src + options:opts + error:&error]; + if (!lib_objc) { + std::string err = + error ? [[error localizedDescription] UTF8String] : "Unknown error"; + throw std::runtime_error("SDPA Metal compile failed: " + err); + } + return (__bridge MTL::Library *)(void *)CFBridgingRetain(lib_objc); + } + } + + MTL::ComputePipelineState *make_pipeline(MTL::Device *mtl_device, + const std::string &name) { + @autoreleasepool { + NSString *fn_name = [NSString stringWithUTF8String:name.c_str()]; + id lib_objc = (__bridge id)lib_; + id func = [lib_objc newFunctionWithName:fn_name]; + if (!func) { + throw std::runtime_error("SDPA kernel not found: " + name); + } + NSError *error = nil; + id device_objc = (__bridge id)mtl_device; + id pso = + [device_objc newComputePipelineStateWithFunction:func error:&error]; + if (!pso) { + std::string err = + error ? [[error localizedDescription] UTF8String] : "Unknown error"; + throw std::runtime_error("SDPA pipeline failed for " + name + ": " + + err); + } + return (__bridge MTL::ComputePipelineState *)(void *)CFBridgingRetain( + pso); + } + } +}; + +// ── Helper: dtype string ──────────────────────────────────────── +static std::string dtype_str(Dtype dt) { + switch (dt) { + case float32: + return "float32"; + case float16: + return "float16"; + case bfloat16: + return "bfloat16"; + default: + throw std::runtime_error("SDPA: unsupported dtype"); + } +} + +// ═══════════════════════════════════════════════════════════════ +// SDPAVector (1-pass) eval_gpu +// ═══════════════════════════════════════════════════════════════ +void SDPAVector::eval_gpu(const std::vector &inputs, + std::vector &outputs) { + auto &queries = inputs[0]; // [B*H, D] + auto &keys = inputs[1]; + auto &values = inputs[2]; + auto &out = outputs[0]; + + int D = queries.shape(-1); + int total_heads = queries.shape(0); + + out.set_data(allocator::malloc(out.nbytes())); + + auto &cache = SDPAPipelineCache::instance(); + cache.ensure_init(kernel_dir_); + + // Kernel name: cider_sdpa_vector_{type}_{D} + std::string kname = "cider_sdpa_vector_" + dtype_str(queries.dtype()) + "_" + + std::to_string(D); + auto *pso = cache.get(kname); + + auto &s = stream(); + auto &enc = metal::get_command_encoder(s); + enc.set_compute_pipeline_state(pso); + + enc.set_input_array(queries, 0); + enc.set_input_array(keys, 1); + enc.set_input_array(values, 2); + enc.set_output_array(out, 3); + enc.set_bytes(gqa_factor_, 4); + enc.set_bytes(N_, 5); + enc.set_bytes(k_head_stride_, 6); + enc.set_bytes(k_seq_stride_, 7); + enc.set_bytes(v_head_stride_, 8); + enc.set_bytes(v_seq_stride_, 9); + enc.set_bytes(scale_, 10); + + // 1-pass: grid=(total_heads, 1, 1), group=(1024, 1, 1) + MTL::Size grid_dims(total_heads, 1, 1); + MTL::Size group_dims(1024, 1, 1); + enc.dispatch_threadgroups(grid_dims, group_dims); +} + +// ═══════════════════════════════════════════════════════════════ +// SDPAVector2Pass eval_gpu +// ═══════════════════════════════════════════════════════════════ +void SDPAVector2Pass::eval_gpu(const std::vector &inputs, + std::vector &outputs) { + auto &queries = inputs[0]; // [B, H, 1, D] + auto &keys = inputs[1]; // [B, Hkv, N, D] + auto &values = inputs[2]; // [B, Hkv, N, D] + auto &out = outputs[0]; // [B, H, 1, D] + + int D = queries.shape(-1); + int num_q_heads = num_kv_heads_ * gqa_factor_; + int total_q_heads = batch_size_ * num_q_heads; + + out.set_data(allocator::malloc(out.nbytes())); + + auto &cache = SDPAPipelineCache::instance(); + cache.ensure_init(kernel_dir_); + + auto &s = stream(); + auto &enc = metal::get_command_encoder(s); + + // ── Allocate intermediates ── + int partial_total = total_q_heads * blocks_; + size_t partial_bytes = + static_cast(partial_total) * D * size_of(queries.dtype()); + size_t stats_bytes = static_cast(partial_total) * sizeof(float); + + array intermediate({partial_total, D}, queries.dtype(), nullptr, {}); + array sums({partial_total}, float32, nullptr, {}); + array maxs({partial_total}, float32, nullptr, {}); + intermediate.set_data(allocator::malloc(intermediate.nbytes())); + sums.set_data(allocator::malloc(sums.nbytes())); + maxs.set_data(allocator::malloc(maxs.nbytes())); + enc.add_temporary(intermediate); + enc.add_temporary(sums); + enc.add_temporary(maxs); + + // ── Pass 1 ── + std::string kname1 = "cider_sdpa_vector_2pass_1_" + + dtype_str(queries.dtype()) + "_" + std::to_string(D) + + "_b" + std::to_string(blocks_); + auto *pso1 = cache.get(kname1); + enc.set_compute_pipeline_state(pso1); + + enc.set_input_array(queries, 0); + enc.set_input_array(keys, 1); + enc.set_input_array(values, 2); + enc.set_output_array(intermediate, 3); + enc.set_output_array(sums, 4); + enc.set_output_array(maxs, 5); + // buffer 6 is skipped (matching MLX convention) + enc.set_bytes(N_, 7); + enc.set_bytes(k_head_stride_, 8); + enc.set_bytes(k_seq_stride_, 9); + enc.set_bytes(v_head_stride_, 10); + enc.set_bytes(v_seq_stride_, 11); + enc.set_bytes(scale_, 12); + + // grid=(num_kv_heads, batch_size, blocks), group=(32, gqa_factor, 1) + MTL::Size grid1(num_kv_heads_, batch_size_, blocks_); + MTL::Size group1(32, gqa_factor_, 1); + enc.dispatch_threadgroups(grid1, group1); + + // ── Pass 2 ── + std::string kname2 = "cider_sdpa_vector_2pass_2_" + + dtype_str(queries.dtype()) + "_" + std::to_string(D) + + "_b" + std::to_string(blocks_); + auto *pso2 = cache.get(kname2); + enc.set_compute_pipeline_state(pso2); + + enc.set_input_array(intermediate, 0); + enc.set_input_array(sums, 1); + enc.set_input_array(maxs, 2); + enc.set_output_array(out, 3); + + // grid=(total_q_heads, 1, 1), group=(1024, 1, 1) + MTL::Size grid2(total_q_heads, 1, 1); + MTL::Size group2(1024, 1, 1); + enc.dispatch_threadgroups(grid2, group2); +} + +// ═══════════════════════════════════════════════════════════════ +// Public API +// ═══════════════════════════════════════════════════════════════ + +array cider_sdpa_1pass(const array &queries, const array &keys, + const array &values, int gqa_factor, float scale, + const std::string &kernel_dir, StreamOrDevice s) { + auto stream = to_stream(s); + + int total_heads = queries.shape(0); + int D = queries.shape(-1); + int N = keys.shape(-2); + + // Compute strides from key/value shapes + // keys: contiguous [B*Hkv, N, D] + size_t k_head_stride = static_cast(N) * D; + size_t k_seq_stride = D; + size_t v_head_stride = k_head_stride; + size_t v_seq_stride = D; + + return array({total_heads, D}, queries.dtype(), + std::make_shared(stream, kernel_dir, gqa_factor, N, + k_head_stride, k_seq_stride, + v_head_stride, v_seq_stride, scale), + {queries, keys, values}); +} + +array cider_sdpa_2pass(const array &queries, const array &keys, + const array &values, int gqa_factor, int blocks, + float scale, const std::string &kernel_dir, + StreamOrDevice s) { + auto stream = to_stream(s); + + int B = queries.shape(0); + int H = queries.shape(1); + int D = queries.shape(-1); + int Hkv = keys.shape(1); + int N = keys.shape(2); + + // Strides for 4D layout [B, Hkv, N, D] + size_t k_head_stride = static_cast(N) * D; + size_t k_seq_stride = D; + size_t v_head_stride = k_head_stride; + size_t v_seq_stride = D; + + auto out_shape = queries.shape(); + + return array(out_shape, queries.dtype(), + std::make_shared( + stream, kernel_dir, gqa_factor, N, blocks, Hkv, B, + k_head_stride, k_seq_stride, v_head_stride, v_seq_stride, + scale), + {queries, keys, values}); +} + +} // namespace cider diff --git a/csrc/src/w4a8_primitive.mm b/csrc/src/w4a8_primitive.mm new file mode 100644 index 0000000..2030c75 --- /dev/null +++ b/csrc/src/w4a8_primitive.mm @@ -0,0 +1,285 @@ +// W4A8 Linear V1 — Packed INT4 weights × INT8 activations +// Reuses W8A8's scratch buffer pool for activation quantization. +// Kernel: unpack W4→INT8 in threadgroup mem, then TensorOps matmul2d. + +#include "w4a8_primitive.h" +#include "mlx/backend/metal/device.h" +#include "mlx/backend/metal/utils.h" + +#import +#import + +#include +#include +#include +#include + +namespace cider { + +using namespace mlx::core; + +struct TileLarge { + static constexpr uint32_t BM = 128, BN = 128, THREADS = 512; +}; +struct TileSmall { + static constexpr uint32_t BM = 32, BN = 128, THREADS = 128; +}; + +static std::string read_file(const std::string &path) { + std::ifstream f(path); + if (!f.is_open()) { + throw std::runtime_error("Cannot open: " + path); + } + std::stringstream ss; + ss << f.rdbuf(); + f.close(); + return ss.str(); +} + +// Scratch buffer pool (same pattern as W8A8) +struct ScratchPool { + allocator::Buffer a_int8_buf{nullptr}; + size_t a_int8_size = 0; + allocator::Buffer scale_a_buf{nullptr}; + size_t scale_a_size = 0; + + allocator::Buffer get_a_int8(size_t needed) { + if (a_int8_buf.ptr() && a_int8_size >= needed) { + return a_int8_buf; + } + if (a_int8_buf.ptr()) { + allocator::free(a_int8_buf); + } + a_int8_buf = allocator::malloc(needed); + a_int8_size = needed; + return a_int8_buf; + } + allocator::Buffer get_scale_a(size_t needed) { + if (scale_a_buf.ptr() && scale_a_size >= needed) { + return scale_a_buf; + } + if (scale_a_buf.ptr()) { + allocator::free(scale_a_buf); + } + scale_a_buf = allocator::malloc(needed); + scale_a_size = needed; + return scale_a_buf; + } + ~ScratchPool() { + if (a_int8_buf.ptr()) { + allocator::free(a_int8_buf); + } + if (scale_a_buf.ptr()) { + allocator::free(scale_a_buf); + } + } +}; + +class W4A8PipelineCache { +public: + static W4A8PipelineCache &instance() { + static W4A8PipelineCache cache; + return cache; + } + + void ensure_init(const std::string &kernel_dir) { + std::lock_guard lock(mutex_); + if (initialized_ && kernel_dir_ == kernel_dir) { + return; + } + auto &dev = metal::device(mlx::core::Device::gpu); + auto *mtl_device = dev.mtl_device(); + w4a8_lib_ = compile_source(mtl_device, kernel_dir + "/w4a8_matmul.metal"); + quantize_lib_ = + compile_source(mtl_device, kernel_dir + "/w8a8_quantize.metal"); + kernel_dir_ = kernel_dir; + pipelines_.clear(); + initialized_ = true; + } + + MTL::ComputePipelineState *get(const std::string &name) { + std::lock_guard lock(mutex_); + auto it = pipelines_.find(name); + if (it != pipelines_.end()) { + return it->second; + } + auto &dev = metal::device(mlx::core::Device::gpu); + auto *pso = make_pipeline(dev.mtl_device(), name); + pipelines_[name] = pso; + return pso; + } + + ScratchPool &scratch() { return scratch_pool_; } + +private: + W4A8PipelineCache() = default; + bool initialized_ = false; + std::string kernel_dir_; + std::unordered_map pipelines_; + MTL::Library *w4a8_lib_ = nullptr; + MTL::Library *quantize_lib_ = nullptr; + ScratchPool scratch_pool_; + std::mutex mutex_; + + MTL::Library *compile_source(MTL::Device *mtl_device, + const std::string &source_path) { + std::string source = read_file(source_path); + @autoreleasepool { + NSString *src = [NSString stringWithUTF8String:source.c_str()]; + MTLCompileOptions *opts = [[MTLCompileOptions alloc] init]; + opts.languageVersion = MTLLanguageVersion4_0; + NSError *error = nil; + id device_objc = (__bridge id)mtl_device; + id lib_objc = [device_objc newLibraryWithSource:src + options:opts + error:&error]; + if (!lib_objc) { + std::string err = + error ? [[error localizedDescription] UTF8String] : "Unknown error"; + throw std::runtime_error("Metal compile failed (" + source_path + + "): " + err); + } + return (__bridge MTL::Library *)(void *)CFBridgingRetain(lib_objc); + } + } + + MTL::ComputePipelineState *make_pipeline(MTL::Device *mtl_device, + const std::string &name) { + @autoreleasepool { + NSString *fn_name = [NSString stringWithUTF8String:name.c_str()]; + // Try w4a8 lib first, then quantize lib + id lib_objc = (__bridge id)w4a8_lib_; + id func = [lib_objc newFunctionWithName:fn_name]; + if (!func) { + lib_objc = (__bridge id)quantize_lib_; + func = [lib_objc newFunctionWithName:fn_name]; + } + if (!func) { + throw std::runtime_error("Kernel not found: " + name); + } + NSError *error = nil; + id device_objc = (__bridge id)mtl_device; + id pso = + [device_objc newComputePipelineStateWithFunction:func error:&error]; + if (!pso) { + std::string err = + error ? [[error localizedDescription] UTF8String] : "Unknown error"; + throw std::runtime_error("Pipeline failed for " + name + ": " + err); + } + return (__bridge MTL::ComputePipelineState *)(void *)CFBridgingRetain( + pso); + } + } +}; + +void W4A8Linear::eval_gpu(const std::vector &inputs, + std::vector &outputs) { + auto &x = inputs[0]; // [M, K] float16 + auto &packed_w = inputs[1]; // [K/2, N] uint8 + auto &scale_w = inputs[2]; // [N] float32 + auto &out = outputs[0]; // [M, N] float16 + + uint32_t M = static_cast(x.shape(0)); + uint32_t K = static_cast(x.shape(1)); + uint32_t N = static_cast(packed_w.shape(1)); + + out.set_data(allocator::malloc(out.nbytes())); + + auto &cache = W4A8PipelineCache::instance(); + cache.ensure_init(kernel_dir_); + auto &pool = cache.scratch(); + + size_t a_bytes = static_cast(M) * K; + size_t sa_bytes = static_cast(M) * sizeof(float); + + allocator::Buffer a_int8_buf = pool.get_a_int8(a_bytes); + allocator::Buffer sa_buf = pool.get_scale_a(sa_bytes); + + auto noop = [](allocator::Buffer) {}; + array a_int8(a_int8_buf, {static_cast(M), static_cast(K)}, int8, + noop); + array sa(sa_buf, {static_cast(M)}, float32, noop); + + auto &s = stream(); + // auto& dev removed — use free function get_command_encoder + auto &enc = metal::get_command_encoder(s); + + // DISPATCH 1: quantize activations (reuse w8a8_quantize.metal) + { + auto *pso = cache.get("quantize_per_token"); + enc.set_compute_pipeline_state(pso); + enc.set_input_array(x, 0); + enc.set_output_array(a_int8, 1); + enc.set_output_array(sa, 2); + enc.set_bytes(M, 3); + enc.set_bytes(K, 4); + uint32_t tg = std::min(256u, std::max(32u, ((K + 31) / 32) * 32)); + enc.dispatch_threadgroups(MTL::Size::Make(M, 1, 1), + MTL::Size::Make(tg, 1, 1)); + } + + enc.barrier(); + + // DISPATCH 2: W4A8 matmul with fused dequant + { + bool use_small = (M <= 64); + const char *kname = use_small ? "w4a8_matmul_fused_dequant_small" + : "w4a8_matmul_fused_dequant"; + uint32_t BM = use_small ? TileSmall::BM : TileLarge::BM; + uint32_t BN = use_small ? TileSmall::BN : TileLarge::BN; + uint32_t threads = use_small ? TileSmall::THREADS : TileLarge::THREADS; + + auto *pso = cache.get(kname); + enc.set_compute_pipeline_state(pso); + + uint32_t tiles_n = (N + BN - 1) / BN; + uint32_t tiles_m = (M + BM - 1) / BM; + uint32_t swizzle_log; + if (tiles_m <= 3) { + swizzle_log = 0; + } else if (tiles_m <= 6) { + swizzle_log = 1; + } else { + swizzle_log = 2; + } + uint32_t tile = 1u << swizzle_log; + uint32_t grid_x = tiles_n * tile; + uint32_t grid_y = (tiles_m + tile - 1) / tile; + + enc.set_input_array(a_int8, 0); // A: INT8 activations + enc.set_input_array(packed_w, 1); // B: packed W4 + enc.set_output_array(out, 2); // C: FP16 output + enc.set_bytes(M, 3); + enc.set_bytes(N, 4); + enc.set_bytes(K, 5); + enc.set_input_array(sa, 6); // scale_a + enc.set_input_array(scale_w, 7); // scale_w + enc.set_bytes(swizzle_log, 8); + enc.set_bytes(tiles_m, 9); + enc.set_bytes(tiles_n, 10); + + enc.dispatch_threadgroups(MTL::Size::Make(grid_x, grid_y, 1), + MTL::Size::Make(threads, 1, 1)); + } +} + +array w4a8_linear(const array &x, const array &packed_w, const array &scale_w, + const std::string &kernel_dir, StreamOrDevice s) { + if (x.ndim() != 2) { + throw std::invalid_argument("w4a8_linear: x must be 2D [M,K]"); + } + if (packed_w.ndim() != 2) { + throw std::invalid_argument("w4a8_linear: packed_w must be 2D [K/2,N]"); + } + + int M = x.shape(0); + int N = packed_w.shape(1); + auto stream = to_stream(s); + + return array( + {M, N}, float16, std::make_shared(stream, kernel_dir), + {astype(x, float16, stream), astype(packed_w, mlx::core::uint8, stream), + astype(scale_w, float32, stream)}); +} + +} // namespace cider diff --git a/csrc/src/w8a8_primitive.mm b/csrc/src/w8a8_primitive.mm new file mode 100644 index 0000000..2425ff5 --- /dev/null +++ b/csrc/src/w8a8_primitive.mm @@ -0,0 +1,344 @@ +// W8A8 Linear V9 — Unified prefill (GEMM) + decode (FP MV) with bias +// +// M > 1: quantize activation → INT8 GEMM → y = acc * scale_a * scale_w + bias +// M == 1: FP MV → y = scale_w * sum(float(w_int8) * float(x)) + bias +// Per-channel: single scale per row, no group switching overhead. + +#include "w8a8_primitive.h" +#include "mlx/backend/metal/device.h" +#include "mlx/backend/metal/utils.h" + +#import +#import + +#include +#include +#include +#include + +namespace cider { + +using namespace mlx::core; + +struct TileLarge { + static constexpr uint32_t BM = 128, BN = 128, THREADS = 512; +}; +struct TileSmall { + static constexpr uint32_t BM = 32, BN = 128, THREADS = 128; +}; + +static std::string read_file(const std::string &path) { + std::ifstream f(path); + if (!f.is_open()) { + throw std::runtime_error("Cannot open: " + path); + } + std::stringstream ss; + ss << f.rdbuf(); + f.close(); + return ss.str(); +} + +// ── Pipeline cache ────────────────────────────────────────────── +class PipelineCache { +public: + static PipelineCache &instance() { + static PipelineCache cache; + return cache; + } + + void ensure_init(const std::string &kernel_dir) { + std::lock_guard lock(mutex_); + if (initialized_ && kernel_dir_ == kernel_dir) { + return; + } + auto &dev = metal::device(mlx::core::Device::gpu); + auto *mtl_device = dev.mtl_device(); + matmul_lib_ = compile_source(mtl_device, kernel_dir + "/w8a8_matmul.metal"); + quantize_lib_ = compile_source(mtl_device, kernel_dir + "/w8a8_quantize.metal"); + mv_lib_ = compile_source(mtl_device, kernel_dir + "/w8a8_int8_mv.metal"); + kernel_dir_ = kernel_dir; + pipelines_.clear(); + initialized_ = true; + } + + MTL::ComputePipelineState *get(const std::string &name) { + std::lock_guard lock(mutex_); + auto it = pipelines_.find(name); + if (it != pipelines_.end()) { + return it->second; + } + auto &dev = metal::device(mlx::core::Device::gpu); + auto *pso = make_pipeline(dev.mtl_device(), name); + pipelines_[name] = pso; + return pso; + } + +private: + PipelineCache() = default; + bool initialized_ = false; + std::string kernel_dir_; + std::unordered_map pipelines_; + MTL::Library *matmul_lib_ = nullptr; + MTL::Library *quantize_lib_ = nullptr; + MTL::Library *mv_lib_ = nullptr; + std::mutex mutex_; + + MTL::Library *compile_source(MTL::Device *mtl_device, + const std::string &source_path) { + std::string source = read_file(source_path); + @autoreleasepool { + NSString *src = [NSString stringWithUTF8String:source.c_str()]; + MTLCompileOptions *opts = [[MTLCompileOptions alloc] init]; + opts.languageVersion = MTLLanguageVersion4_0; + NSError *error = nil; + id device_objc = (__bridge id)mtl_device; + id lib_objc = [device_objc newLibraryWithSource:src + options:opts + error:&error]; + if (!lib_objc) { + std::string err = + error ? [[error localizedDescription] UTF8String] : "Unknown error"; + throw std::runtime_error("Metal compile failed (" + source_path + + "): " + err); + } + return (__bridge MTL::Library *)(void *)CFBridgingRetain(lib_objc); + } + } + + MTL::ComputePipelineState *make_pipeline(MTL::Device *mtl_device, + const std::string &name) { + @autoreleasepool { + NSString *fn_name = [NSString stringWithUTF8String:name.c_str()]; + id func = nil; + for (auto *lib : {matmul_lib_, quantize_lib_, mv_lib_}) { + if (!lib) continue; + id lib_objc = (__bridge id)lib; + func = [lib_objc newFunctionWithName:fn_name]; + if (func) break; + } + if (!func) { + throw std::runtime_error("Kernel not found: " + name); + } + NSError *error = nil; + id device_objc = (__bridge id)mtl_device; + id pso = + [device_objc newComputePipelineStateWithFunction:func error:&error]; + if (!pso) { + std::string err = + error ? [[error localizedDescription] UTF8String] : "Unknown error"; + throw std::runtime_error("Pipeline failed for " + name + ": " + err); + } + return (__bridge MTL::ComputePipelineState *)(void *)CFBridgingRetain( + pso); + } + } +}; + +void W8A8Linear::eval_gpu(const std::vector &inputs, + std::vector &outputs) { + auto &x = inputs[0]; // [M, K] float16 + auto &w = inputs[1]; // [N, K] int8 + auto &scale_w = inputs[2]; // [N] float32 + auto &bias = inputs[3]; // [N] float16 + auto &out = outputs[0]; // [M, N] float16 + + uint32_t M = static_cast(x.shape(0)); + uint32_t N = static_cast(w.shape(0)); + uint32_t K = static_cast(w.shape(1)); + + out.set_data(allocator::malloc(out.nbytes())); + + auto &cache = PipelineCache::instance(); + cache.ensure_init(kernel_dir_); + + auto &s = stream(); + auto &enc = metal::get_command_encoder(s); + + if (M == 1) { + // ── Decode path: FP MV (no activation quantization) ── + // w8a8_int8_mv.metal buffer layout: + // x[0], W[1], y[2], scale_w[3], N[4], K[5], bias[6] + auto *pso = cache.get("w8a8_int8_mv"); + enc.set_compute_pipeline_state(pso); + enc.set_input_array(x, 0); // [1, K] float16 + enc.set_input_array(w, 1); // [N, K] int8 + enc.set_output_array(out, 2); // [1, N] float16 + enc.set_input_array(scale_w, 3); // [N] float32 + enc.set_bytes(N, 4); + enc.set_bytes(K, 5); + enc.set_input_array(bias, 6); // [N] float16 + + constexpr uint32_t NUM_SIMDGROUPS = 2; + constexpr uint32_t RESULTS_PER_SG = 4; + constexpr uint32_t SIMD_SIZE = 32; + uint32_t rows_per_tg = NUM_SIMDGROUPS * RESULTS_PER_SG; // 16 + uint32_t grid_n = (N + rows_per_tg - 1) / rows_per_tg; + + enc.dispatch_threadgroups(MTL::Size::Make(grid_n, 1, 1), + MTL::Size::Make(SIMD_SIZE * NUM_SIMDGROUPS, 1, 1)); + } else { + // ── Prefill path: quantize + INT8 GEMM ── + size_t a_bytes = static_cast(M) * K; + size_t sa_bytes = static_cast(M) * sizeof(float); + + array a_int8({static_cast(M), static_cast(K)}, int8, nullptr, {}); + a_int8.set_data(allocator::malloc(a_bytes)); + + array sa({static_cast(M)}, float32, nullptr, {}); + sa.set_data(allocator::malloc(sa_bytes)); + + // DISPATCH 1: quantize_per_token + { + auto *pso = cache.get("quantize_per_token"); + enc.set_compute_pipeline_state(pso); + enc.set_input_array(x, 0); + enc.set_output_array(a_int8, 1); + enc.set_output_array(sa, 2); + enc.set_bytes(M, 3); + enc.set_bytes(K, 4); + uint32_t tg = std::min(256u, std::max(32u, ((K + 31) / 32) * 32)); + enc.dispatch_threadgroups(MTL::Size::Make(M, 1, 1), + MTL::Size::Make(tg, 1, 1)); + } + + enc.barrier(); + + // DISPATCH 2: matmul_dequant + { + bool use_small = (M <= 64); + const char *kname = use_small ? "w8a8_matmul_fused_dequant_small" + : "w8a8_matmul_fused_dequant"; + uint32_t BM = use_small ? TileSmall::BM : TileLarge::BM; + uint32_t BN = use_small ? TileSmall::BN : TileLarge::BN; + uint32_t threads = use_small ? TileSmall::THREADS : TileLarge::THREADS; + + auto *pso = cache.get(kname); + enc.set_compute_pipeline_state(pso); + + uint32_t tiles_n = (N + BN - 1) / BN; + uint32_t tiles_m = (M + BM - 1) / BM; + uint32_t swizzle_log; + if (tiles_m <= 3) { + swizzle_log = 0; + } else if (tiles_m <= 6) { + swizzle_log = 1; + } else { + swizzle_log = 2; + } + uint32_t tile = 1u << swizzle_log; + uint32_t grid_x = tiles_n * tile; + uint32_t grid_y = (tiles_m + tile - 1) / tile; + + enc.set_input_array(a_int8, 0); + enc.set_input_array(w, 1); + enc.set_output_array(out, 2); + enc.set_bytes(M, 3); + enc.set_bytes(N, 4); + enc.set_bytes(K, 5); + enc.set_input_array(sa, 6); + enc.set_input_array(scale_w, 7); + enc.set_bytes(swizzle_log, 8); + enc.set_bytes(tiles_m, 9); + enc.set_bytes(tiles_n, 10); + enc.set_input_array(bias, 11); + + enc.dispatch_threadgroups(MTL::Size::Make(grid_x, grid_y, 1), + MTL::Size::Make(threads, 1, 1)); + } + + enc.add_temporary(a_int8); + enc.add_temporary(sa); + } +} + +array perchannel_linear(const array &x, const array &w, const array &scale_w, + const array &bias, const std::string &kernel_dir, + StreamOrDevice s) { + if (x.ndim() != 2) { + throw std::invalid_argument("perchannel_linear: x must be 2D [M,K]"); + } + if (w.ndim() != 2) { + throw std::invalid_argument("perchannel_linear: w must be 2D [N,K]"); + } + + int M = x.shape(0); + int N = w.shape(0); + auto stream = to_stream(s); + + auto result = + array({M, N}, float16, std::make_shared(stream, kernel_dir), + {astype(x, float16, stream), astype(w, int8, stream), + astype(scale_w, float32, stream), astype(bias, float16, stream)}); + if (x.dtype() == bfloat16) { + return astype(result, bfloat16, stream); + } + return result; +} + +// ── Int8MatMulInt32 primitive (raw INT32 output, no dequant) ──── + +void Int8MatMulInt32::eval_gpu(const std::vector &inputs, + std::vector &outputs) { + auto &a = inputs[0]; // [M, K] int8 + auto &b = inputs[1]; // [N, K] int8 + + uint32_t M = static_cast(a.shape(0)); + uint32_t N = static_cast(b.shape(0)); + uint32_t K = static_cast(b.shape(1)); + + auto &out = outputs[0]; + out.set_data(mx::allocator::malloc(out.nbytes())); + + auto &cache = PipelineCache::instance(); + cache.ensure_init(kernel_dir_); + + bool use_small = (M <= 64); + auto *pso = use_small ? cache.get("int8_matmul_int32_small") + : cache.get("int8_matmul_int32"); + + uint32_t BM = use_small ? 32 : 128; + uint32_t BN = 128; + uint32_t threads = use_small ? 128 : 512; + + uint32_t tiles_n = (N + BN - 1) / BN; + uint32_t tiles_m = (M + BM - 1) / BM; + uint32_t swizzle_log; + if (tiles_m <= 3) { + swizzle_log = 0; + } else if (tiles_m <= 6) { + swizzle_log = 1; + } else { + swizzle_log = 2; + } + uint32_t tile = 1u << swizzle_log; + uint32_t grid_x = tiles_n * tile; + uint32_t grid_y = (tiles_m + tile - 1) / tile; + + auto &s = stream(); + auto &enc = metal::get_command_encoder(s); + enc.set_compute_pipeline_state(pso); + enc.set_input_array(a, 0); + enc.set_input_array(b, 1); + enc.set_output_array(out, 2); + enc.set_bytes(&M, sizeof(M), 3); + enc.set_bytes(&N, sizeof(N), 4); + enc.set_bytes(&K, sizeof(K), 5); + enc.set_bytes(&swizzle_log, sizeof(swizzle_log), 6); + enc.set_bytes(&tiles_m, sizeof(tiles_m), 7); + enc.set_bytes(&tiles_n, sizeof(tiles_n), 8); + enc.dispatch_threadgroups(MTL::Size(grid_x, grid_y, 1), + MTL::Size(threads, 1, 1)); +} + +mx::array int8_matmul_int32(const mx::array &a, const mx::array &b, + const std::string &kernel_dir, + mx::StreamOrDevice s) { + uint32_t M = static_cast(a.shape(0)); + uint32_t N = static_cast(b.shape(0)); + return mx::array( + {static_cast(M), static_cast(N)}, mx::int32, + std::make_shared(mx::to_stream(s), kernel_dir), + {astype(a, int8, mx::to_stream(s)), astype(b, int8, mx::to_stream(s))}); +} + +} // namespace cider diff --git a/dev/step_kernels/bench_final.py b/dev/step_kernels/bench_final.py new file mode 100644 index 0000000..bb0df85 --- /dev/null +++ b/dev/step_kernels/bench_final.py @@ -0,0 +1,131 @@ +#!/usr/bin/env python3 +"""Step-by-step kernel benchmark. Run each step sequentially with gc.""" +import sys, os, time, gc +import numpy as np +sys.path.insert(0, "~/work/cider/cider/lib") +import mlx.core as mx +import _cider_prim as ext + +DEV = "~/work/cider/dev/step_kernels" +STEPS = ["step1", "step2", "step3", "step4", "step5"] +WARMUP = 5 +REPEAT = 50 + +# Large tile shapes (M%128==0) +LARGE = [ + (128, 4096, 4096), + (256, 4096, 4096), + (128, 3584, 18944), + (256, 3584, 18944), +] + +# Small tile shapes (step3+ only) +SMALL = [ + (1, 4096, 4096), + (16, 4096, 4096), + (32, 4096, 4096), + (64, 4096, 4096), +] + +def bench_one(A_mx, B_mx, kernel_dir): + """Warmup + benchmark, return median ms.""" + for _ in range(WARMUP): + mx.eval(ext.int8_matmul_int32(A_mx, B_mx, kernel_dir)) + times = [] + for _ in range(REPEAT): + t0 = time.perf_counter() + mx.eval(ext.int8_matmul_int32(A_mx, B_mx, kernel_dir)) + t1 = time.perf_counter() + times.append((t1 - t0) * 1000) + times.sort() + return times[len(times) // 2] + +# === Large tile benchmark === +print("=" * 90) +print("Part 1: Large Tile (M>=128)") +print("=" * 90) + +large_results = {s: {} for s in STEPS} + +for M, K, N in LARGE: + key = f"({M},{K},{N})" + np.random.seed(42) + A_np = np.random.randint(-127, 128, (M, K), dtype=np.int8) + B_np = np.random.randint(-127, 128, (K, N), dtype=np.int8) + A_mx = mx.array(A_np) + B_mx = mx.array(B_np) + mx.eval(A_mx, B_mx) + del A_np, B_np + gc.collect() + + for step in STEPS: + kdir = os.path.join(DEV, step) + med = bench_one(A_mx, B_mx, kdir) + large_results[step][key] = med + print(f" {step} {key}: {med:.4f} ms") + + del A_mx, B_mx + mx.clear_cache() + gc.collect() + +# Print large table +print() +hdr = f"{'Shape':<22}" +for s in STEPS: + hdr += f" | {s:>10}" +print(hdr) +print("-" * 82) +for M, K, N in LARGE: + key = f"({M},{K},{N})" + row = f"{key:<22}" + for s in STEPS: + v = large_results[s].get(key, -1) + row += f" | {v:>10.3f}" + print(row) + +# === Small tile benchmark (step3/4/5 only) === +print() +print("=" * 60) +print("Part 2: Small Tile (M<128, step3/4/5)") +print("=" * 60) + +small_steps = ["step3", "step4", "step5"] +small_results = {s: {} for s in small_steps} + +for M, K, N in SMALL: + key = f"({M},{K},{N})" + np.random.seed(42) + A_np = np.random.randint(-127, 128, (M, K), dtype=np.int8) + B_np = np.random.randint(-127, 128, (K, N), dtype=np.int8) + A_mx = mx.array(A_np) + B_mx = mx.array(B_np) + mx.eval(A_mx, B_mx) + del A_np, B_np + gc.collect() + + for step in small_steps: + kdir = os.path.join(DEV, step) + med = bench_one(A_mx, B_mx, kdir) + small_results[step][key] = med + print(f" {step} {key}: {med:.4f} ms") + + del A_mx, B_mx + mx.clear_cache() + gc.collect() + +# Print small table +print() +hdr = f"{'Shape':<22}" +for s in small_steps: + hdr += f" | {s:>10}" +print(hdr) +print("-" * 58) +for M, K, N in SMALL: + key = f"({M},{K},{N})" + row = f"{key:<22}" + for s in small_steps: + v = small_results[s].get(key, -1) + row += f" | {v:>10.3f}" + print(row) + +print("\nDone!") diff --git a/dev/step_kernels/step1/w8a8_matmul.metal b/dev/step_kernels/step1/w8a8_matmul.metal new file mode 100644 index 0000000..b918ad5 --- /dev/null +++ b/dev/step_kernels/step1/w8a8_matmul.metal @@ -0,0 +1,381 @@ +// ============================================================ +// Step 1: Naive INT8 TensorOps GEMM with Threadgroup Memory Staging +// - BM=128, BN=128, BK=128, SK=32 +// - Only large tile (WM=4, WN=4, 512 threads) +// - Threadgroup memory staging: load A/B tiles to TG memory first +// - Swizzle decode included (required by C++ dispatch) +// ============================================================ + +#include +#include +using namespace metal; + +constant constexpr short kElemsPerFrag = 8; +constant constexpr short kElemCols = 4; +constant constexpr short kElemRowsJump = 8; + +inline short2 nax_get_coord(ushort lid) { + short qid = short(lid >> 2); + short fm = ((qid & 4) | ((short(lid) >> 1) & 3)); + short fn = ((qid & 2) | (short(lid) & 1)) * 4; + return short2{fn, fm}; +} + +// Swizzle decode: convert swizzled grid coords to real tile coords +inline void swizzle_decode(uint2 tgid, uint swizzle_log, uint tiles_n, + thread uint &tid_y, thread uint &tid_x) { + uint tile = 1u << swizzle_log; + uint tgid_y_raw = tgid.y * tile + (tgid.x % tile); + uint tgid_x_raw = tgid.x / tile; + tid_y = tgid_y_raw; + tid_x = tgid_x_raw; +} + +template +inline void nax_frag_load_tg(thread T *dst, const threadgroup T *src, int ld, + short2 sc, short off_m = 0, short off_n = 0) { + src += (sc.y + off_m) * ld + (sc.x + off_n); + for (short i = 0; i < 2; i++) + for (short j = 0; j < kElemCols; j++) + dst[i * kElemCols + j] = src[(i * kElemRowsJump) * ld + j]; +} + +inline void nax_frag_store_int32(const thread int32_t *src, device int32_t *dst, + int ld, short2 sc, short off_m, short off_n, + uint M, uint N, uint m_base, uint n_base) { + for (short i = 0; i < 2; i++) + for (short j = 0; j < kElemCols; j++) { + uint mi = m_base + sc.y + off_m + i * kElemRowsJump; + uint ni = n_base + sc.x + off_n + j; + if (mi < M && ni < N) + dst[(sc.y + off_m + i * kElemRowsJump) * ld + (sc.x + off_n + j)] = + src[i * kElemCols + j]; + } +} + +inline void nax_frag_store_dequant(const thread int32_t *src, device half *dst, + int ld, short2 sc, short off_m, short off_n, + uint M, uint N, uint m_base, uint n_base, + const device float *scale_a, + const device float *scale_w) { + for (short i = 0; i < 2; i++) + for (short j = 0; j < kElemCols; j++) { + uint mi = m_base + sc.y + off_m + i * kElemRowsJump; + uint ni = n_base + sc.x + off_n + j; + if (mi < M && ni < N) { + float val = float(src[i * kElemCols + j]) * scale_a[mi] * scale_w[ni]; + dst[(sc.y + off_m + i * kElemRowsJump) * ld + (sc.x + off_n + j)] = + half(val); + } + } +} + +// Step 1: TG-staged GEMM, BM=128 only +template +void step1_gemm_int32_impl( + const device int8_t *A, const device int8_t *B, + device int32_t *C, uint M, uint N, uint K, + uint swizzle_log, uint tiles_m, uint tiles_n, + threadgroup int8_t *tg_A, + threadgroup int8_t *tg_B, + uint2 tgid, uint sgid, uint lid, uint tid_in_tg) { + + constexpr int SM = BM / WM; // 32 + constexpr int SN = BN / WN; // 32 + constexpr short TM = SM / 16; // 2 + constexpr short TN = SN / 16; // 2 + constexpr short TK = SK / 16; // 2 + constexpr int NUM_THREADS = WM * WN * 32; // 512 + + uint tid_y, tid_x; + swizzle_decode(tgid, swizzle_log, tiles_n, tid_y, tid_x); + if (tid_x >= tiles_n || tid_y >= tiles_m) return; + + short2 sc = nax_get_coord(ushort(lid)); + uint sg_row = sgid / WN; + uint sg_col = sgid % WN; + uint m_base = tid_y * BM + sg_row * SM; + uint n_base = tid_x * BN + sg_col * SN; + + constexpr auto desc = mpp::tensor_ops::matmul2d_descriptor( + 16, 32, 16, false, false, true, + mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate); + mpp::tensor_ops::matmul2d gemm_op; + + auto ct_a = gemm_op.get_left_input_cooperative_tensor(); + auto ct_b = gemm_op.get_right_input_cooperative_tensor(); + auto ct_c = gemm_op.get_destination_cooperative_tensor(); + + int32_t c_frags[TM * TN][kElemsPerFrag]; + for (int f = 0; f < TM * TN; f++) + for (int i = 0; i < kElemsPerFrag; i++) + c_frags[f][i] = 0; + + for (uint k_base = 0; k_base < K; k_base += BK) { + // Cooperative load A[BM, BK] → threadgroup + uint a_elems = BM * BK; + for (uint idx = tid_in_tg; idx < a_elems; idx += NUM_THREADS) { + uint row = idx / BK; + uint col = idx % BK; + uint gm = tid_y * BM + row; + uint gk = k_base + col; + tg_A[row * BK + col] = (gm < M && gk < K) ? A[gm * K + gk] : int8_t(0); + } + // Cooperative load B[BK, BN] → threadgroup + uint b_elems = BK * BN; + for (uint idx = tid_in_tg; idx < b_elems; idx += NUM_THREADS) { + uint row = idx / BN; + uint col = idx % BN; + uint gk = k_base + row; + uint gn = tid_x * BN + col; + tg_B[row * BN + col] = (gk < K && gn < N) ? B[gk * N + gn] : int8_t(0); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (int kk1 = 0; kk1 < BK; kk1 += SK) { + int8_t a_frags[TM][TK][kElemsPerFrag]; + int8_t b_frags[TK][TN][kElemsPerFrag]; + + for (short mm = 0; mm < TM; mm++) + for (short kk = 0; kk < TK; kk++) + nax_frag_load_tg(a_frags[mm][kk], + tg_A + (sg_row * SM) * BK + kk1, + BK, sc, short(mm * 16), short(kk * 16)); + for (short kk = 0; kk < TK; kk++) + for (short nn = 0; nn < TN; nn++) + nax_frag_load_tg(b_frags[kk][nn], + tg_B + kk1 * BN + (sg_col * SN), + BN, sc, short(kk * 16), short(nn * 16)); + + for (short mm = 0; mm < TM; mm++) { + for (short nn = 0; nn < TN; nn += 2) { + for (short kk = 0; kk < TK; kk++) { + for (short i = 0; i < kElemsPerFrag; i++) ct_a[i] = a_frags[mm][kk][i]; + for (short i = 0; i < kElemsPerFrag; i++) { + ct_b[i] = b_frags[kk][nn][i]; + ct_b[kElemsPerFrag + i] = b_frags[kk][nn + 1][i]; + } + short c0 = mm * TN + nn, c1 = c0 + 1; + for (short i = 0; i < kElemsPerFrag; i++) { + ct_c[i] = c_frags[c0][i]; + ct_c[kElemsPerFrag + i] = c_frags[c1][i]; + } + gemm_op.run(ct_a, ct_b, ct_c); + for (short i = 0; i < kElemsPerFrag; i++) { + c_frags[c0][i] = ct_c[i]; + c_frags[c1][i] = ct_c[kElemsPerFrag + i]; + } + } + } + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + + device int32_t *D = C + m_base * N + n_base; + for (short mm = 0; mm < TM; mm++) + for (short nn = 0; nn < TN; nn++) + nax_frag_store_int32(c_frags[mm * TN + nn], D, int(N), sc, + short(mm * 16), short(nn * 16), M, N, m_base, n_base); +} + +// Dequant version +template +void step1_gemm_dequant_impl( + const device int8_t *A, const device int8_t *B, + device half *C, uint M, uint N, uint K, + const device float *scale_a, const device float *scale_w, + uint swizzle_log, uint tiles_m, uint tiles_n, + threadgroup int8_t *tg_A, + threadgroup int8_t *tg_B, + uint2 tgid, uint sgid, uint lid, uint tid_in_tg) { + + constexpr int SM = BM / WM; + constexpr int SN = BN / WN; + constexpr short TM = SM / 16; + constexpr short TN = SN / 16; + constexpr short TK = SK / 16; + constexpr int NUM_THREADS = WM * WN * 32; + + uint tid_y, tid_x; + swizzle_decode(tgid, swizzle_log, tiles_n, tid_y, tid_x); + if (tid_x >= tiles_n || tid_y >= tiles_m) return; + + short2 sc = nax_get_coord(ushort(lid)); + uint sg_row = sgid / WN; + uint sg_col = sgid % WN; + uint m_base = tid_y * BM + sg_row * SM; + uint n_base = tid_x * BN + sg_col * SN; + + constexpr auto desc = mpp::tensor_ops::matmul2d_descriptor( + 16, 32, 16, false, false, true, + mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate); + mpp::tensor_ops::matmul2d gemm_op; + + auto ct_a = gemm_op.get_left_input_cooperative_tensor(); + auto ct_b = gemm_op.get_right_input_cooperative_tensor(); + auto ct_c = gemm_op.get_destination_cooperative_tensor(); + + int32_t c_frags[TM * TN][kElemsPerFrag]; + for (int f = 0; f < TM * TN; f++) + for (int i = 0; i < kElemsPerFrag; i++) + c_frags[f][i] = 0; + + for (uint k_base = 0; k_base < K; k_base += BK) { + uint a_elems = BM * BK; + for (uint idx = tid_in_tg; idx < a_elems; idx += NUM_THREADS) { + uint row = idx / BK; + uint col = idx % BK; + uint gm = tid_y * BM + row; + uint gk = k_base + col; + tg_A[row * BK + col] = (gm < M && gk < K) ? A[gm * K + gk] : int8_t(0); + } + uint b_elems = BK * BN; + for (uint idx = tid_in_tg; idx < b_elems; idx += NUM_THREADS) { + uint row = idx / BN; + uint col = idx % BN; + uint gk = k_base + row; + uint gn = tid_x * BN + col; + tg_B[row * BN + col] = (gk < K && gn < N) ? B[gk * N + gn] : int8_t(0); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (int kk1 = 0; kk1 < BK; kk1 += SK) { + int8_t a_frags[TM][TK][kElemsPerFrag]; + int8_t b_frags[TK][TN][kElemsPerFrag]; + for (short mm = 0; mm < TM; mm++) + for (short kk = 0; kk < TK; kk++) + nax_frag_load_tg(a_frags[mm][kk], + tg_A + (sg_row * SM) * BK + kk1, + BK, sc, short(mm * 16), short(kk * 16)); + for (short kk = 0; kk < TK; kk++) + for (short nn = 0; nn < TN; nn++) + nax_frag_load_tg(b_frags[kk][nn], + tg_B + kk1 * BN + (sg_col * SN), + BN, sc, short(kk * 16), short(nn * 16)); + for (short mm = 0; mm < TM; mm++) { + for (short nn = 0; nn < TN; nn += 2) { + for (short kk = 0; kk < TK; kk++) { + for (short i = 0; i < kElemsPerFrag; i++) ct_a[i] = a_frags[mm][kk][i]; + for (short i = 0; i < kElemsPerFrag; i++) { + ct_b[i] = b_frags[kk][nn][i]; + ct_b[kElemsPerFrag + i] = b_frags[kk][nn + 1][i]; + } + short c0 = mm * TN + nn, c1 = c0 + 1; + for (short i = 0; i < kElemsPerFrag; i++) { + ct_c[i] = c_frags[c0][i]; + ct_c[kElemsPerFrag + i] = c_frags[c1][i]; + } + gemm_op.run(ct_a, ct_b, ct_c); + for (short i = 0; i < kElemsPerFrag; i++) { + c_frags[c0][i] = ct_c[i]; + c_frags[c1][i] = ct_c[kElemsPerFrag + i]; + } + } + } + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + + device half *D = C + m_base * N + n_base; + for (short mm = 0; mm < TM; mm++) + for (short nn = 0; nn < TN; nn++) + nax_frag_store_dequant(c_frags[mm * TN + nn], D, int(N), sc, + short(mm * 16), short(nn * 16), M, N, m_base, n_base, + scale_a, scale_w); +} + +// ============================================================ +// Kernel entry points +// ============================================================ + +kernel void w8a8_matmul_fused_dequant( + const device int8_t *A [[buffer(0)]], + const device int8_t *B [[buffer(1)]], + device half *C [[buffer(2)]], + constant uint &M [[buffer(3)]], + constant uint &N [[buffer(4)]], + constant uint &K [[buffer(5)]], + const device float *scale_a [[buffer(6)]], + const device float *scale_w [[buffer(7)]], + constant uint &swizzle_log [[buffer(8)]], + constant uint &tiles_m [[buffer(9)]], + constant uint &tiles_n [[buffer(10)]], + uint2 tgid [[threadgroup_position_in_grid]], + uint sgid [[simdgroup_index_in_threadgroup]], + uint lid [[thread_index_in_simdgroup]], + uint tid_in_tg [[thread_index_in_threadgroup]]) { + threadgroup int8_t tg_A[128 * 128]; + threadgroup int8_t tg_B[128 * 128]; + step1_gemm_dequant_impl<128, 128, 128, 32, 4, 4>( + A, B, C, M, N, K, scale_a, scale_w, swizzle_log, tiles_m, tiles_n, + tg_A, tg_B, tgid, sgid, lid, tid_in_tg); +} + +kernel void w8a8_matmul_fused_dequant_small( + const device int8_t *A [[buffer(0)]], + const device int8_t *B [[buffer(1)]], + device half *C [[buffer(2)]], + constant uint &M [[buffer(3)]], + constant uint &N [[buffer(4)]], + constant uint &K [[buffer(5)]], + const device float *scale_a [[buffer(6)]], + const device float *scale_w [[buffer(7)]], + constant uint &swizzle_log [[buffer(8)]], + constant uint &tiles_m [[buffer(9)]], + constant uint &tiles_n [[buffer(10)]], + uint2 tgid [[threadgroup_position_in_grid]], + uint sgid [[simdgroup_index_in_threadgroup]], + uint lid [[thread_index_in_simdgroup]], + uint tid_in_tg [[thread_index_in_threadgroup]]) { + // Step 1 has no small tile — fallback to large tile + threadgroup int8_t tg_A[128 * 128]; + threadgroup int8_t tg_B[128 * 128]; + step1_gemm_dequant_impl<128, 128, 128, 32, 4, 4>( + A, B, C, M, N, K, scale_a, scale_w, swizzle_log, tiles_m, tiles_n, + tg_A, tg_B, tgid, sgid, lid, tid_in_tg); +} + +kernel void int8_matmul_int32( + const device int8_t *A [[buffer(0)]], + const device int8_t *B [[buffer(1)]], + device int32_t *C [[buffer(2)]], + constant uint &M [[buffer(3)]], + constant uint &N [[buffer(4)]], + constant uint &K [[buffer(5)]], + constant uint &swizzle_log [[buffer(6)]], + constant uint &tiles_m [[buffer(7)]], + constant uint &tiles_n [[buffer(8)]], + uint2 tgid [[threadgroup_position_in_grid]], + uint sgid [[simdgroup_index_in_threadgroup]], + uint lid [[thread_index_in_simdgroup]], + uint tid_in_tg [[thread_index_in_threadgroup]]) { + threadgroup int8_t tg_A[128 * 128]; + threadgroup int8_t tg_B[128 * 128]; + step1_gemm_int32_impl<128, 128, 128, 32, 4, 4>( + A, B, C, M, N, K, swizzle_log, tiles_m, tiles_n, + tg_A, tg_B, tgid, sgid, lid, tid_in_tg); +} + +kernel void int8_matmul_int32_small( + const device int8_t *A [[buffer(0)]], + const device int8_t *B [[buffer(1)]], + device int32_t *C [[buffer(2)]], + constant uint &M [[buffer(3)]], + constant uint &N [[buffer(4)]], + constant uint &K [[buffer(5)]], + constant uint &swizzle_log [[buffer(6)]], + constant uint &tiles_m [[buffer(7)]], + constant uint &tiles_n [[buffer(8)]], + uint2 tgid [[threadgroup_position_in_grid]], + uint sgid [[simdgroup_index_in_threadgroup]], + uint lid [[thread_index_in_simdgroup]], + uint tid_in_tg [[thread_index_in_threadgroup]]) { + // Step 1 has no small tile — fallback to large tile + threadgroup int8_t tg_A[128 * 128]; + threadgroup int8_t tg_B[128 * 128]; + step1_gemm_int32_impl<128, 128, 128, 32, 4, 4>( + A, B, C, M, N, K, swizzle_log, tiles_m, tiles_n, + tg_A, tg_B, tgid, sgid, lid, tid_in_tg); +} diff --git a/dev/step_kernels/step1/w8a8_quantize.metal b/dev/step_kernels/step1/w8a8_quantize.metal new file mode 100644 index 0000000..f9ae0be --- /dev/null +++ b/dev/step_kernels/step1/w8a8_quantize.metal @@ -0,0 +1,87 @@ +// ============================================================ +// Per-token quantization: FP16 → INT8 + float32 scale +// Target: Apple M5, Metal 4 +// +// Each threadgroup handles one row (one token). +// Threads cooperate to find absmax via simdgroup reduce, +// then quantize in parallel. +// +// Host dispatch: +// threadgroup = (min(256, ceil(K/32)*32), 1, 1) +// grid = (M, 1, 1) +// ============================================================ + +#include +using namespace metal; + +kernel void +quantize_per_token(const device half *X [[buffer(0)]], // [M, K] FP16 input + device int8_t *A [[buffer(1)]], // [M, K] INT8 output + device float *scale [[buffer(2)]], // [M] float32 scale + constant uint &M [[buffer(3)]], + constant uint &K [[buffer(4)]], + uint gid [[threadgroup_position_in_grid]], // row index + uint lid [[thread_index_in_threadgroup]], + uint tg_size [[threads_per_threadgroup]]) { + if (gid >= M) { + return; + } + + const device half *row_in = X + gid * K; + device int8_t *row_out = A + gid * K; + + // Step 1: Find local absmax + float local_max = 0.0f; + for (uint i = lid; i < K; i += tg_size) { + float v = abs(float(row_in[i])); + local_max = max(local_max, v); + } + + // Step 2: Simdgroup reduce max + float sg_max = simd_max(local_max); + + // Step 3: Threadgroup reduce across simdgroups via shared memory + threadgroup float sg_maxes[8]; // up to 8 simdgroups (256/32) + threadgroup float shared_scale; + threadgroup float shared_inv_scale; + uint sg_id = lid / 32; + uint sg_lid = lid % 32; + if (sg_lid == 0) { + sg_maxes[sg_id] = sg_max; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + // Final reduce (first simdgroup only) + if (sg_id == 0) { + float row_max = 0.0f; + uint num_sgs = (tg_size + 31) / 32; + if (sg_lid < num_sgs) { + row_max = sg_maxes[sg_lid]; + } + row_max = simd_max(row_max); + + // Compute and broadcast scale + float s = row_max / 255.0f; + if (s == 0.0f) { + s = 1.0f; + } + + if (sg_lid == 0) { + shared_scale = s; + shared_inv_scale = 1.0f / s; + // Store scale to output + scale[gid] = s; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + // Step 4: All threads read broadcasted scale + float inv_s = shared_inv_scale; + + // Step 5: Quantize + for (uint i = lid; i < K; i += tg_size) { + float v = float(row_in[i]) * inv_s; + v = clamp(round(v), -128.0f, 127.0f); + row_out[i] = int8_t(v); + } +} diff --git a/dev/step_kernels/step2/w8a8_matmul.metal b/dev/step_kernels/step2/w8a8_matmul.metal new file mode 100644 index 0000000..7404bc0 --- /dev/null +++ b/dev/step_kernels/step2/w8a8_matmul.metal @@ -0,0 +1,285 @@ +// ============================================================ +// Step 2: Direct Device Read (No Threadgroup Memory) +// - BM=128, BN=128, BK=128, SK=32 +// - Only large tile (WM=4, WN=4, 512 threads) +// - Direct device memory read (no TG staging) +// - Swizzle decode included (required by C++ dispatch) +// ============================================================ + +#include +#include +using namespace metal; + +constant constexpr short kElemsPerFrag = 8; +constant constexpr short kElemCols = 4; +constant constexpr short kElemRowsJump = 8; + +inline short2 nax_get_coord(ushort lid) { + short qid = short(lid >> 2); + short fm = ((qid & 4) | ((short(lid) >> 1) & 3)); + short fn = ((qid & 2) | (short(lid) & 1)) * 4; + return short2{fn, fm}; +} + +inline void swizzle_decode(uint2 tgid, uint swizzle_log, uint tiles_n, + thread uint &tid_y, thread uint &tid_x) { + uint tile = 1u << swizzle_log; + tid_y = tgid.y * tile + (tgid.x % tile); + tid_x = tgid.x / tile; +} + +template +inline void nax_frag_load(thread T *dst, const device T *src, int ld, short2 sc, + short off_m = 0, short off_n = 0) { + src += (sc.y + off_m) * ld + (sc.x + off_n); + for (short i = 0; i < 2; i++) + for (short j = 0; j < kElemCols; j++) + dst[i * kElemCols + j] = src[(i * kElemRowsJump) * ld + j]; +} + +inline void nax_frag_store_int32(const thread int32_t *src, device int32_t *dst, + int ld, short2 sc, short off_m, short off_n, + uint M, uint N, uint m_base, uint n_base) { + for (short i = 0; i < 2; i++) + for (short j = 0; j < kElemCols; j++) { + uint mi = m_base + sc.y + off_m + i * kElemRowsJump; + uint ni = n_base + sc.x + off_n + j; + if (mi < M && ni < N) + dst[(sc.y + off_m + i * kElemRowsJump) * ld + (sc.x + off_n + j)] = + src[i * kElemCols + j]; + } +} + +inline void nax_frag_store_dequant(const thread int32_t *src, device half *dst, + int ld, short2 sc, short off_m, short off_n, + uint M, uint N, uint m_base, uint n_base, + const device float *scale_a, + const device float *scale_w) { + for (short i = 0; i < 2; i++) + for (short j = 0; j < kElemCols; j++) { + uint mi = m_base + sc.y + off_m + i * kElemRowsJump; + uint ni = n_base + sc.x + off_n + j; + if (mi < M && ni < N) { + float val = float(src[i * kElemCols + j]) * scale_a[mi] * scale_w[ni]; + dst[(sc.y + off_m + i * kElemRowsJump) * ld + (sc.x + off_n + j)] = + half(val); + } + } +} + +// Direct device read, BK=128, large tile only +template +void step2_gemm_int32_impl( + const device int8_t *A, const device int8_t *B, + device int32_t *C, uint M, uint N, uint K, + uint swizzle_log, uint tiles_m, uint tiles_n, + uint2 tgid, uint sgid, uint lid) { + + constexpr int SM = BM / WM; + constexpr int SN = BN / WN; + constexpr short TM = SM / 16; + constexpr short TN = SN / 16; + constexpr short TK = SK / 16; + + uint tid_y, tid_x; + swizzle_decode(tgid, swizzle_log, tiles_n, tid_y, tid_x); + if (tid_x >= tiles_n || tid_y >= tiles_m) return; + + short2 sc = nax_get_coord(ushort(lid)); + uint sg_row = sgid / WN; + uint sg_col = sgid % WN; + uint m_base = tid_y * BM + sg_row * SM; + uint n_base = tid_x * BN + sg_col * SN; + + const device int8_t *sg_A = A + m_base * K; + const device int8_t *sg_B = B + n_base; + + constexpr auto desc = mpp::tensor_ops::matmul2d_descriptor( + 16, 32, 16, false, false, true, + mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate); + mpp::tensor_ops::matmul2d gemm_op; + + auto ct_a = gemm_op.get_left_input_cooperative_tensor(); + auto ct_b = gemm_op.get_right_input_cooperative_tensor(); + auto ct_c = gemm_op.get_destination_cooperative_tensor(); + + int32_t c_frags[TM * TN][kElemsPerFrag]; + for (int f = 0; f < TM * TN; f++) + for (int i = 0; i < kElemsPerFrag; i++) + c_frags[f][i] = 0; + + int gemm_k_iters = int(K) / BK; + for (int kk0 = 0; kk0 < gemm_k_iters; kk0++) { + threadgroup_barrier(mem_flags::mem_none); + for (int kk1 = 0; kk1 < BK; kk1 += SK) { + int8_t a_frags[TM][TK][kElemsPerFrag]; + int8_t b_frags[TK][TN][kElemsPerFrag]; + volatile int compiler_barrier; + + for (short mm = 0; mm < TM; mm++) + for (short kk = 0; kk < TK; kk++) + nax_frag_load(a_frags[mm][kk], sg_A + kk1, int(K), sc, + short(mm * 16), short(kk * 16)); + for (short kk = 0; kk < TK; kk++) + for (short nn = 0; nn < TN; nn++) + nax_frag_load(b_frags[kk][nn], sg_B + kk1 * N, int(N), sc, + short(kk * 16), short(nn * 16)); + + for (short mm = 0; mm < TM; mm++) { + for (short nn = 0; nn < TN; nn += 2) { + for (short kk = 0; kk < TK; kk++) { + for (short i = 0; i < kElemsPerFrag; i++) ct_a[i] = a_frags[mm][kk][i]; + for (short i = 0; i < kElemsPerFrag; i++) { + ct_b[i] = b_frags[kk][nn][i]; + ct_b[kElemsPerFrag + i] = b_frags[kk][nn + 1][i]; + } + short c0 = mm * TN + nn, c1 = c0 + 1; + for (short i = 0; i < kElemsPerFrag; i++) { + ct_c[i] = c_frags[c0][i]; + ct_c[kElemsPerFrag + i] = c_frags[c1][i]; + } + gemm_op.run(ct_a, ct_b, ct_c); + for (short i = 0; i < kElemsPerFrag; i++) { + c_frags[c0][i] = ct_c[i]; + c_frags[c1][i] = ct_c[kElemsPerFrag + i]; + } + } + } + } + (void)compiler_barrier; + } + sg_A += BK; + sg_B += BK * N; + } + + // Handle K remainder + int rem_k = int(K) - gemm_k_iters * BK; + if (rem_k > 0) { + for (int kk1 = 0; kk1 < rem_k; kk1 += SK) { + int actual_sk = min(SK, rem_k - kk1); + int8_t a_frags[TM][TK][kElemsPerFrag]; + int8_t b_frags[TK][TN][kElemsPerFrag]; + + for (short mm = 0; mm < TM; mm++) + for (short kk = 0; kk < TK; kk++) { + const device int8_t *ptr = sg_A + kk1 + (sc.y + mm * 16) * int(K) + (sc.x + kk * 16); + for (short i = 0; i < 2; i++) + for (short j = 0; j < kElemCols; j++) { + int ki = kk * 16 + int(sc.x) + j; + if (ki < actual_sk) { + int row = int(sc.y) + mm * 16 + i * kElemRowsJump; + a_frags[mm][kk][i * kElemCols + j] = sg_A[kk1 + row * int(K) + ki]; + } else { + a_frags[mm][kk][i * kElemCols + j] = 0; + } + } + } + for (short kk = 0; kk < TK; kk++) + for (short nn = 0; nn < TN; nn++) { + for (short i = 0; i < 2; i++) + for (short j = 0; j < kElemCols; j++) { + int ki = kk * 16 + int(sc.y) + i * kElemRowsJump; + if (ki < actual_sk) { + int col = nn * 16 + int(sc.x) + j; + b_frags[kk][nn][i * kElemCols + j] = sg_B[(kk1 + ki) * int(N) + col]; + } else { + b_frags[kk][nn][i * kElemCols + j] = 0; + } + } + } + + for (short mm = 0; mm < TM; mm++) { + for (short nn = 0; nn < TN; nn += 2) { + for (short kk = 0; kk < TK; kk++) { + for (short i = 0; i < kElemsPerFrag; i++) ct_a[i] = a_frags[mm][kk][i]; + for (short i = 0; i < kElemsPerFrag; i++) { + ct_b[i] = b_frags[kk][nn][i]; + ct_b[kElemsPerFrag + i] = b_frags[kk][nn + 1][i]; + } + short c0 = mm * TN + nn, c1 = c0 + 1; + for (short i = 0; i < kElemsPerFrag; i++) { + ct_c[i] = c_frags[c0][i]; + ct_c[kElemsPerFrag + i] = c_frags[c1][i]; + } + gemm_op.run(ct_a, ct_b, ct_c); + for (short i = 0; i < kElemsPerFrag; i++) { + c_frags[c0][i] = ct_c[i]; + c_frags[c1][i] = ct_c[kElemsPerFrag + i]; + } + } + } + } + } + } + + device int32_t *D = C + m_base * N + n_base; + for (short mm = 0; mm < TM; mm++) + for (short nn = 0; nn < TN; nn++) + nax_frag_store_int32(c_frags[mm * TN + nn], D, int(N), sc, + short(mm * 16), short(nn * 16), M, N, m_base, n_base); +} + +// ============================================================ +// Entry points +// ============================================================ + +kernel void w8a8_matmul_fused_dequant( + const device int8_t *A [[buffer(0)]], const device int8_t *B [[buffer(1)]], + device half *C [[buffer(2)]], constant uint &M [[buffer(3)]], + constant uint &N [[buffer(4)]], constant uint &K [[buffer(5)]], + const device float *scale_a [[buffer(6)]], + const device float *scale_w [[buffer(7)]], + constant uint &swizzle_log [[buffer(8)]], + constant uint &tiles_m [[buffer(9)]], constant uint &tiles_n [[buffer(10)]], + uint2 tgid [[threadgroup_position_in_grid]], + uint sgid [[simdgroup_index_in_threadgroup]], + uint lid [[thread_index_in_simdgroup]]) { + step2_gemm_int32_impl<128, 128, 128, 32, 4, 4>( + A, B, (device int32_t *)C, M, N, K, swizzle_log, tiles_m, tiles_n, + tgid, sgid, lid); + // NOTE: for dequant we'd need a separate impl, but for benchmark + // we only test int32 path. This entry exists for compilation only. +} + +kernel void w8a8_matmul_fused_dequant_small( + const device int8_t *A [[buffer(0)]], const device int8_t *B [[buffer(1)]], + device half *C [[buffer(2)]], constant uint &M [[buffer(3)]], + constant uint &N [[buffer(4)]], constant uint &K [[buffer(5)]], + const device float *scale_a [[buffer(6)]], + const device float *scale_w [[buffer(7)]], + constant uint &swizzle_log [[buffer(8)]], + constant uint &tiles_m [[buffer(9)]], constant uint &tiles_n [[buffer(10)]], + uint2 tgid [[threadgroup_position_in_grid]], + uint sgid [[simdgroup_index_in_threadgroup]], + uint lid [[thread_index_in_simdgroup]]) { + step2_gemm_int32_impl<128, 128, 128, 32, 4, 4>( + A, B, (device int32_t *)C, M, N, K, swizzle_log, tiles_m, tiles_n, + tgid, sgid, lid); +} + +kernel void int8_matmul_int32( + const device int8_t *A [[buffer(0)]], const device int8_t *B [[buffer(1)]], + device int32_t *C [[buffer(2)]], constant uint &M [[buffer(3)]], + constant uint &N [[buffer(4)]], constant uint &K [[buffer(5)]], + constant uint &swizzle_log [[buffer(6)]], + constant uint &tiles_m [[buffer(7)]], constant uint &tiles_n [[buffer(8)]], + uint2 tgid [[threadgroup_position_in_grid]], + uint sgid [[simdgroup_index_in_threadgroup]], + uint lid [[thread_index_in_simdgroup]]) { + step2_gemm_int32_impl<128, 128, 128, 32, 4, 4>( + A, B, C, M, N, K, swizzle_log, tiles_m, tiles_n, tgid, sgid, lid); +} + +kernel void int8_matmul_int32_small( + const device int8_t *A [[buffer(0)]], const device int8_t *B [[buffer(1)]], + device int32_t *C [[buffer(2)]], constant uint &M [[buffer(3)]], + constant uint &N [[buffer(4)]], constant uint &K [[buffer(5)]], + constant uint &swizzle_log [[buffer(6)]], + constant uint &tiles_m [[buffer(7)]], constant uint &tiles_n [[buffer(8)]], + uint2 tgid [[threadgroup_position_in_grid]], + uint sgid [[simdgroup_index_in_threadgroup]], + uint lid [[thread_index_in_simdgroup]]) { + // Step 2 has no small tile — fallback to large + step2_gemm_int32_impl<128, 128, 128, 32, 4, 4>( + A, B, C, M, N, K, swizzle_log, tiles_m, tiles_n, tgid, sgid, lid); +} diff --git a/dev/step_kernels/step2/w8a8_quantize.metal b/dev/step_kernels/step2/w8a8_quantize.metal new file mode 100644 index 0000000..f9ae0be --- /dev/null +++ b/dev/step_kernels/step2/w8a8_quantize.metal @@ -0,0 +1,87 @@ +// ============================================================ +// Per-token quantization: FP16 → INT8 + float32 scale +// Target: Apple M5, Metal 4 +// +// Each threadgroup handles one row (one token). +// Threads cooperate to find absmax via simdgroup reduce, +// then quantize in parallel. +// +// Host dispatch: +// threadgroup = (min(256, ceil(K/32)*32), 1, 1) +// grid = (M, 1, 1) +// ============================================================ + +#include +using namespace metal; + +kernel void +quantize_per_token(const device half *X [[buffer(0)]], // [M, K] FP16 input + device int8_t *A [[buffer(1)]], // [M, K] INT8 output + device float *scale [[buffer(2)]], // [M] float32 scale + constant uint &M [[buffer(3)]], + constant uint &K [[buffer(4)]], + uint gid [[threadgroup_position_in_grid]], // row index + uint lid [[thread_index_in_threadgroup]], + uint tg_size [[threads_per_threadgroup]]) { + if (gid >= M) { + return; + } + + const device half *row_in = X + gid * K; + device int8_t *row_out = A + gid * K; + + // Step 1: Find local absmax + float local_max = 0.0f; + for (uint i = lid; i < K; i += tg_size) { + float v = abs(float(row_in[i])); + local_max = max(local_max, v); + } + + // Step 2: Simdgroup reduce max + float sg_max = simd_max(local_max); + + // Step 3: Threadgroup reduce across simdgroups via shared memory + threadgroup float sg_maxes[8]; // up to 8 simdgroups (256/32) + threadgroup float shared_scale; + threadgroup float shared_inv_scale; + uint sg_id = lid / 32; + uint sg_lid = lid % 32; + if (sg_lid == 0) { + sg_maxes[sg_id] = sg_max; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + // Final reduce (first simdgroup only) + if (sg_id == 0) { + float row_max = 0.0f; + uint num_sgs = (tg_size + 31) / 32; + if (sg_lid < num_sgs) { + row_max = sg_maxes[sg_lid]; + } + row_max = simd_max(row_max); + + // Compute and broadcast scale + float s = row_max / 255.0f; + if (s == 0.0f) { + s = 1.0f; + } + + if (sg_lid == 0) { + shared_scale = s; + shared_inv_scale = 1.0f / s; + // Store scale to output + scale[gid] = s; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + // Step 4: All threads read broadcasted scale + float inv_s = shared_inv_scale; + + // Step 5: Quantize + for (uint i = lid; i < K; i += tg_size) { + float v = float(row_in[i]) * inv_s; + v = clamp(round(v), -128.0f, 127.0f); + row_out[i] = int8_t(v); + } +} diff --git a/dev/step_kernels/step3/w8a8_matmul.metal b/dev/step_kernels/step3/w8a8_matmul.metal new file mode 100644 index 0000000..c6fe169 --- /dev/null +++ b/dev/step_kernels/step3/w8a8_matmul.metal @@ -0,0 +1,431 @@ +// ============================================================ +// Step 3: Multi-tile Dispatch (Large BM=128 + Small BM=32) +// - BK=128, SK=32 +// - Direct device read +// - Swizzle decode (required by C++ dispatch) +// - Large: BM=128, BN=128, WM=4, WN=4 (512 threads) +// - Small: BM=32, BN=128, WM=1, WN=4 (128 threads) +// ============================================================ + +#include +#include +using namespace metal; + +constant constexpr short kElemsPerFrag = 8; +constant constexpr short kElemCols = 4; +constant constexpr short kElemRowsJump = 8; + +inline short2 nax_get_coord(ushort lid) { + short qid = short(lid >> 2); + short fm = ((qid & 4) | ((short(lid) >> 1) & 3)); + short fn = ((qid & 2) | (short(lid) & 1)) * 4; + return short2{fn, fm}; +} + +inline void swizzle_decode(uint2 tgid, uint swizzle_log, uint tiles_n, + thread uint &tid_y, thread uint &tid_x) { + uint tile = 1u << swizzle_log; + tid_y = tgid.y * tile + (tgid.x % tile); + tid_x = tgid.x / tile; +} + +template +inline void nax_frag_load(thread T *dst, const device T *src, int ld, short2 sc, + short off_m = 0, short off_n = 0) { + src += (sc.y + off_m) * ld + (sc.x + off_n); + for (short i = 0; i < 2; i++) + for (short j = 0; j < kElemCols; j++) + dst[i * kElemCols + j] = src[(i * kElemRowsJump) * ld + j]; +} + +inline void nax_frag_store_int32(const thread int32_t *src, device int32_t *dst, + int ld, short2 sc, short off_m, short off_n, + uint M, uint N, uint m_base, uint n_base) { + for (short i = 0; i < 2; i++) + for (short j = 0; j < kElemCols; j++) { + uint mi = m_base + sc.y + off_m + i * kElemRowsJump; + uint ni = n_base + sc.x + off_n + j; + if (mi < M && ni < N) + dst[(sc.y + off_m + i * kElemRowsJump) * ld + (sc.x + off_n + j)] = + src[i * kElemCols + j]; + } +} + +inline void nax_frag_store_dequant(const thread int32_t *src, device half *dst, + int ld, short2 sc, short off_m, short off_n, + uint M, uint N, uint m_base, uint n_base, + const device float *scale_a, + const device float *scale_w) { + for (short i = 0; i < 2; i++) + for (short j = 0; j < kElemCols; j++) { + uint mi = m_base + sc.y + off_m + i * kElemRowsJump; + uint ni = n_base + sc.x + off_n + j; + if (mi < M && ni < N) { + float val = float(src[i * kElemCols + j]) * scale_a[mi] * scale_w[ni]; + dst[(sc.y + off_m + i * kElemRowsJump) * ld + (sc.x + off_n + j)] = + half(val); + } + } +} + +// Generic GEMM with swizzle decode +template +void step3_gemm_int32_impl( + const device int8_t *A, const device int8_t *B, + device int32_t *C, uint M, uint N, uint K, + uint swizzle_log, uint tiles_m, uint tiles_n, + uint2 tgid, uint sgid, uint lid) { + + constexpr int SM = BM / WM; + constexpr int SN = BN / WN; + constexpr short TM = SM / 16; + constexpr short TN = SN / 16; + constexpr short TK = SK / 16; + + uint tid_y, tid_x; + swizzle_decode(tgid, swizzle_log, tiles_n, tid_y, tid_x); + if (tid_x >= tiles_n || tid_y >= tiles_m) return; + + short2 sc = nax_get_coord(ushort(lid)); + uint sg_row = sgid / WN; + uint sg_col = sgid % WN; + uint m_base = tid_y * BM + sg_row * SM; + uint n_base = tid_x * BN + sg_col * SN; + + const device int8_t *sg_A = A + m_base * K; + const device int8_t *sg_B = B + n_base; + + constexpr auto desc = mpp::tensor_ops::matmul2d_descriptor( + 16, 32, 16, false, false, true, + mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate); + mpp::tensor_ops::matmul2d gemm_op; + + auto ct_a = gemm_op.get_left_input_cooperative_tensor(); + auto ct_b = gemm_op.get_right_input_cooperative_tensor(); + auto ct_c = gemm_op.get_destination_cooperative_tensor(); + + int32_t c_frags[TM * TN][kElemsPerFrag]; + for (int f = 0; f < TM * TN; f++) + for (int i = 0; i < kElemsPerFrag; i++) + c_frags[f][i] = 0; + + int gemm_k_iters = int(K) / BK; + for (int kk0 = 0; kk0 < gemm_k_iters; kk0++) { + threadgroup_barrier(mem_flags::mem_none); + for (int kk1 = 0; kk1 < BK; kk1 += SK) { + int8_t a_frags[TM][TK][kElemsPerFrag]; + int8_t b_frags[TK][TN][kElemsPerFrag]; + volatile int compiler_barrier; + + for (short mm = 0; mm < TM; mm++) + for (short kk = 0; kk < TK; kk++) + nax_frag_load(a_frags[mm][kk], sg_A + kk1, int(K), sc, + short(mm * 16), short(kk * 16)); + for (short kk = 0; kk < TK; kk++) + for (short nn = 0; nn < TN; nn++) + nax_frag_load(b_frags[kk][nn], sg_B + kk1 * N, int(N), sc, + short(kk * 16), short(nn * 16)); + + for (short mm = 0; mm < TM; mm++) { + for (short nn = 0; nn < TN; nn += 2) { + for (short kk = 0; kk < TK; kk++) { + for (short i = 0; i < kElemsPerFrag; i++) ct_a[i] = a_frags[mm][kk][i]; + for (short i = 0; i < kElemsPerFrag; i++) { + ct_b[i] = b_frags[kk][nn][i]; + ct_b[kElemsPerFrag + i] = b_frags[kk][nn + 1][i]; + } + short c0 = mm * TN + nn, c1 = c0 + 1; + for (short i = 0; i < kElemsPerFrag; i++) { + ct_c[i] = c_frags[c0][i]; + ct_c[kElemsPerFrag + i] = c_frags[c1][i]; + } + gemm_op.run(ct_a, ct_b, ct_c); + for (short i = 0; i < kElemsPerFrag; i++) { + c_frags[c0][i] = ct_c[i]; + c_frags[c1][i] = ct_c[kElemsPerFrag + i]; + } + } + } + } + (void)compiler_barrier; + } + sg_A += BK; + sg_B += BK * N; + } + + // Remainder K + int rem_k = int(K) - gemm_k_iters * BK; + if (rem_k > 0) { + for (int kk1 = 0; kk1 < rem_k; kk1 += SK) { + int actual_sk = min(SK, rem_k - kk1); + int8_t a_frags[TM][TK][kElemsPerFrag]; + int8_t b_frags[TK][TN][kElemsPerFrag]; + + for (short mm = 0; mm < TM; mm++) + for (short kk = 0; kk < TK; kk++) { + for (short i = 0; i < 2; i++) + for (short j = 0; j < kElemCols; j++) { + int ki = kk * 16 + int(sc.x) + j; + if (ki < actual_sk) { + int row = int(sc.y) + mm * 16 + i * kElemRowsJump; + a_frags[mm][kk][i * kElemCols + j] = sg_A[kk1 + row * int(K) + ki]; + } else { + a_frags[mm][kk][i * kElemCols + j] = 0; + } + } + } + for (short kk = 0; kk < TK; kk++) + for (short nn = 0; nn < TN; nn++) { + for (short i = 0; i < 2; i++) + for (short j = 0; j < kElemCols; j++) { + int ki = kk * 16 + int(sc.y) + i * kElemRowsJump; + if (ki < actual_sk) { + int col = nn * 16 + int(sc.x) + j; + b_frags[kk][nn][i * kElemCols + j] = sg_B[(kk1 + ki) * int(N) + col]; + } else { + b_frags[kk][nn][i * kElemCols + j] = 0; + } + } + } + + for (short mm = 0; mm < TM; mm++) { + for (short nn = 0; nn < TN; nn += 2) { + for (short kk = 0; kk < TK; kk++) { + for (short i = 0; i < kElemsPerFrag; i++) ct_a[i] = a_frags[mm][kk][i]; + for (short i = 0; i < kElemsPerFrag; i++) { + ct_b[i] = b_frags[kk][nn][i]; + ct_b[kElemsPerFrag + i] = b_frags[kk][nn + 1][i]; + } + short c0 = mm * TN + nn, c1 = c0 + 1; + for (short i = 0; i < kElemsPerFrag; i++) { + ct_c[i] = c_frags[c0][i]; + ct_c[kElemsPerFrag + i] = c_frags[c1][i]; + } + gemm_op.run(ct_a, ct_b, ct_c); + for (short i = 0; i < kElemsPerFrag; i++) { + c_frags[c0][i] = ct_c[i]; + c_frags[c1][i] = ct_c[kElemsPerFrag + i]; + } + } + } + } + } + } + + device int32_t *D = C + m_base * N + n_base; + for (short mm = 0; mm < TM; mm++) + for (short nn = 0; nn < TN; nn++) + nax_frag_store_int32(c_frags[mm * TN + nn], D, int(N), sc, + short(mm * 16), short(nn * 16), M, N, m_base, n_base); +} + +// Dequant version +template +void step3_gemm_dequant_impl( + const device int8_t *A, const device int8_t *B, + device half *C, uint M, uint N, uint K, + const device float *scale_a, const device float *scale_w, + uint swizzle_log, uint tiles_m, uint tiles_n, + uint2 tgid, uint sgid, uint lid) { + + constexpr int SM = BM / WM; + constexpr int SN = BN / WN; + constexpr short TM = SM / 16; + constexpr short TN = SN / 16; + constexpr short TK = SK / 16; + + uint tid_y, tid_x; + swizzle_decode(tgid, swizzle_log, tiles_n, tid_y, tid_x); + if (tid_x >= tiles_n || tid_y >= tiles_m) return; + + short2 sc = nax_get_coord(ushort(lid)); + uint sg_row = sgid / WN; + uint sg_col = sgid % WN; + uint m_base = tid_y * BM + sg_row * SM; + uint n_base = tid_x * BN + sg_col * SN; + + const device int8_t *sg_A = A + m_base * K; + const device int8_t *sg_B = B + n_base; + + constexpr auto desc = mpp::tensor_ops::matmul2d_descriptor( + 16, 32, 16, false, false, true, + mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate); + mpp::tensor_ops::matmul2d gemm_op; + + auto ct_a = gemm_op.get_left_input_cooperative_tensor(); + auto ct_b = gemm_op.get_right_input_cooperative_tensor(); + auto ct_c = gemm_op.get_destination_cooperative_tensor(); + + int32_t c_frags[TM * TN][kElemsPerFrag]; + for (int f = 0; f < TM * TN; f++) + for (int i = 0; i < kElemsPerFrag; i++) + c_frags[f][i] = 0; + + int gemm_k_iters = int(K) / BK; + for (int kk0 = 0; kk0 < gemm_k_iters; kk0++) { + threadgroup_barrier(mem_flags::mem_none); + for (int kk1 = 0; kk1 < BK; kk1 += SK) { + int8_t a_frags[TM][TK][kElemsPerFrag]; + int8_t b_frags[TK][TN][kElemsPerFrag]; + volatile int compiler_barrier; + for (short mm = 0; mm < TM; mm++) + for (short kk = 0; kk < TK; kk++) + nax_frag_load(a_frags[mm][kk], sg_A + kk1, int(K), sc, + short(mm * 16), short(kk * 16)); + for (short kk = 0; kk < TK; kk++) + for (short nn = 0; nn < TN; nn++) + nax_frag_load(b_frags[kk][nn], sg_B + kk1 * N, int(N), sc, + short(kk * 16), short(nn * 16)); + for (short mm = 0; mm < TM; mm++) { + for (short nn = 0; nn < TN; nn += 2) { + for (short kk = 0; kk < TK; kk++) { + for (short i = 0; i < kElemsPerFrag; i++) ct_a[i] = a_frags[mm][kk][i]; + for (short i = 0; i < kElemsPerFrag; i++) { + ct_b[i] = b_frags[kk][nn][i]; + ct_b[kElemsPerFrag + i] = b_frags[kk][nn + 1][i]; + } + short c0 = mm * TN + nn, c1 = c0 + 1; + for (short i = 0; i < kElemsPerFrag; i++) { + ct_c[i] = c_frags[c0][i]; + ct_c[kElemsPerFrag + i] = c_frags[c1][i]; + } + gemm_op.run(ct_a, ct_b, ct_c); + for (short i = 0; i < kElemsPerFrag; i++) { + c_frags[c0][i] = ct_c[i]; + c_frags[c1][i] = ct_c[kElemsPerFrag + i]; + } + } + } + } + (void)compiler_barrier; + } + sg_A += BK; + sg_B += BK * N; + } + + int rem_k = int(K) - gemm_k_iters * BK; + if (rem_k > 0) { + for (int kk1 = 0; kk1 < rem_k; kk1 += SK) { + int actual_sk = min(SK, rem_k - kk1); + int8_t a_frags[TM][TK][kElemsPerFrag]; + int8_t b_frags[TK][TN][kElemsPerFrag]; + for (short mm = 0; mm < TM; mm++) + for (short kk = 0; kk < TK; kk++) { + for (short i = 0; i < 2; i++) + for (short j = 0; j < kElemCols; j++) { + int ki = kk * 16 + int(sc.x) + j; + if (ki < actual_sk) { + int row = int(sc.y) + mm * 16 + i * kElemRowsJump; + a_frags[mm][kk][i * kElemCols + j] = sg_A[kk1 + row * int(K) + ki]; + } else { + a_frags[mm][kk][i * kElemCols + j] = 0; + } + } + } + for (short kk = 0; kk < TK; kk++) + for (short nn = 0; nn < TN; nn++) { + for (short i = 0; i < 2; i++) + for (short j = 0; j < kElemCols; j++) { + int ki = kk * 16 + int(sc.y) + i * kElemRowsJump; + if (ki < actual_sk) { + int col = nn * 16 + int(sc.x) + j; + b_frags[kk][nn][i * kElemCols + j] = sg_B[(kk1 + ki) * int(N) + col]; + } else { + b_frags[kk][nn][i * kElemCols + j] = 0; + } + } + } + for (short mm = 0; mm < TM; mm++) { + for (short nn = 0; nn < TN; nn += 2) { + for (short kk = 0; kk < TK; kk++) { + for (short i = 0; i < kElemsPerFrag; i++) ct_a[i] = a_frags[mm][kk][i]; + for (short i = 0; i < kElemsPerFrag; i++) { + ct_b[i] = b_frags[kk][nn][i]; + ct_b[kElemsPerFrag + i] = b_frags[kk][nn + 1][i]; + } + short c0 = mm * TN + nn, c1 = c0 + 1; + for (short i = 0; i < kElemsPerFrag; i++) { + ct_c[i] = c_frags[c0][i]; + ct_c[kElemsPerFrag + i] = c_frags[c1][i]; + } + gemm_op.run(ct_a, ct_b, ct_c); + for (short i = 0; i < kElemsPerFrag; i++) { + c_frags[c0][i] = ct_c[i]; + c_frags[c1][i] = ct_c[kElemsPerFrag + i]; + } + } + } + } + } + } + + device half *D = C + m_base * N + n_base; + for (short mm = 0; mm < TM; mm++) + for (short nn = 0; nn < TN; nn++) + nax_frag_store_dequant(c_frags[mm * TN + nn], D, int(N), sc, + short(mm * 16), short(nn * 16), M, N, m_base, n_base, + scale_a, scale_w); +} + +// ============================================================ +// Entry points — Large tile +// ============================================================ +kernel void w8a8_matmul_fused_dequant( + const device int8_t *A [[buffer(0)]], const device int8_t *B [[buffer(1)]], + device half *C [[buffer(2)]], constant uint &M [[buffer(3)]], + constant uint &N [[buffer(4)]], constant uint &K [[buffer(5)]], + const device float *scale_a [[buffer(6)]], + const device float *scale_w [[buffer(7)]], + constant uint &swizzle_log [[buffer(8)]], + constant uint &tiles_m [[buffer(9)]], constant uint &tiles_n [[buffer(10)]], + uint2 tgid [[threadgroup_position_in_grid]], + uint sgid [[simdgroup_index_in_threadgroup]], + uint lid [[thread_index_in_simdgroup]]) { + step3_gemm_dequant_impl<128, 128, 128, 32, 4, 4>( + A, B, C, M, N, K, scale_a, scale_w, swizzle_log, tiles_m, tiles_n, + tgid, sgid, lid); +} + +// Small tile entry point (NEW in Step 3!) +kernel void w8a8_matmul_fused_dequant_small( + const device int8_t *A [[buffer(0)]], const device int8_t *B [[buffer(1)]], + device half *C [[buffer(2)]], constant uint &M [[buffer(3)]], + constant uint &N [[buffer(4)]], constant uint &K [[buffer(5)]], + const device float *scale_a [[buffer(6)]], + const device float *scale_w [[buffer(7)]], + constant uint &swizzle_log [[buffer(8)]], + constant uint &tiles_m [[buffer(9)]], constant uint &tiles_n [[buffer(10)]], + uint2 tgid [[threadgroup_position_in_grid]], + uint sgid [[simdgroup_index_in_threadgroup]], + uint lid [[thread_index_in_simdgroup]]) { + step3_gemm_dequant_impl<32, 128, 128, 32, 1, 4>( + A, B, C, M, N, K, scale_a, scale_w, swizzle_log, tiles_m, tiles_n, + tgid, sgid, lid); +} + +// INT32 entry points +kernel void int8_matmul_int32( + const device int8_t *A [[buffer(0)]], const device int8_t *B [[buffer(1)]], + device int32_t *C [[buffer(2)]], constant uint &M [[buffer(3)]], + constant uint &N [[buffer(4)]], constant uint &K [[buffer(5)]], + constant uint &swizzle_log [[buffer(6)]], + constant uint &tiles_m [[buffer(7)]], constant uint &tiles_n [[buffer(8)]], + uint2 tgid [[threadgroup_position_in_grid]], + uint sgid [[simdgroup_index_in_threadgroup]], + uint lid [[thread_index_in_simdgroup]]) { + step3_gemm_int32_impl<128, 128, 128, 32, 4, 4>( + A, B, C, M, N, K, swizzle_log, tiles_m, tiles_n, tgid, sgid, lid); +} + +kernel void int8_matmul_int32_small( + const device int8_t *A [[buffer(0)]], const device int8_t *B [[buffer(1)]], + device int32_t *C [[buffer(2)]], constant uint &M [[buffer(3)]], + constant uint &N [[buffer(4)]], constant uint &K [[buffer(5)]], + constant uint &swizzle_log [[buffer(6)]], + constant uint &tiles_m [[buffer(7)]], constant uint &tiles_n [[buffer(8)]], + uint2 tgid [[threadgroup_position_in_grid]], + uint sgid [[simdgroup_index_in_threadgroup]], + uint lid [[thread_index_in_simdgroup]]) { + step3_gemm_int32_impl<32, 128, 128, 32, 1, 4>( + A, B, C, M, N, K, swizzle_log, tiles_m, tiles_n, tgid, sgid, lid); +} diff --git a/dev/step_kernels/step3/w8a8_quantize.metal b/dev/step_kernels/step3/w8a8_quantize.metal new file mode 100644 index 0000000..f9ae0be --- /dev/null +++ b/dev/step_kernels/step3/w8a8_quantize.metal @@ -0,0 +1,87 @@ +// ============================================================ +// Per-token quantization: FP16 → INT8 + float32 scale +// Target: Apple M5, Metal 4 +// +// Each threadgroup handles one row (one token). +// Threads cooperate to find absmax via simdgroup reduce, +// then quantize in parallel. +// +// Host dispatch: +// threadgroup = (min(256, ceil(K/32)*32), 1, 1) +// grid = (M, 1, 1) +// ============================================================ + +#include +using namespace metal; + +kernel void +quantize_per_token(const device half *X [[buffer(0)]], // [M, K] FP16 input + device int8_t *A [[buffer(1)]], // [M, K] INT8 output + device float *scale [[buffer(2)]], // [M] float32 scale + constant uint &M [[buffer(3)]], + constant uint &K [[buffer(4)]], + uint gid [[threadgroup_position_in_grid]], // row index + uint lid [[thread_index_in_threadgroup]], + uint tg_size [[threads_per_threadgroup]]) { + if (gid >= M) { + return; + } + + const device half *row_in = X + gid * K; + device int8_t *row_out = A + gid * K; + + // Step 1: Find local absmax + float local_max = 0.0f; + for (uint i = lid; i < K; i += tg_size) { + float v = abs(float(row_in[i])); + local_max = max(local_max, v); + } + + // Step 2: Simdgroup reduce max + float sg_max = simd_max(local_max); + + // Step 3: Threadgroup reduce across simdgroups via shared memory + threadgroup float sg_maxes[8]; // up to 8 simdgroups (256/32) + threadgroup float shared_scale; + threadgroup float shared_inv_scale; + uint sg_id = lid / 32; + uint sg_lid = lid % 32; + if (sg_lid == 0) { + sg_maxes[sg_id] = sg_max; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + // Final reduce (first simdgroup only) + if (sg_id == 0) { + float row_max = 0.0f; + uint num_sgs = (tg_size + 31) / 32; + if (sg_lid < num_sgs) { + row_max = sg_maxes[sg_lid]; + } + row_max = simd_max(row_max); + + // Compute and broadcast scale + float s = row_max / 255.0f; + if (s == 0.0f) { + s = 1.0f; + } + + if (sg_lid == 0) { + shared_scale = s; + shared_inv_scale = 1.0f / s; + // Store scale to output + scale[gid] = s; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + // Step 4: All threads read broadcasted scale + float inv_s = shared_inv_scale; + + // Step 5: Quantize + for (uint i = lid; i < K; i += tg_size) { + float v = float(row_in[i]) * inv_s; + v = clamp(round(v), -128.0f, 127.0f); + row_out[i] = int8_t(v); + } +} diff --git a/dev/step_kernels/step4/w8a8_matmul.metal b/dev/step_kernels/step4/w8a8_matmul.metal new file mode 100644 index 0000000..3e1f987 --- /dev/null +++ b/dev/step_kernels/step4/w8a8_matmul.metal @@ -0,0 +1,431 @@ +// ============================================================ +// Step 4: Deep K-loop (BK=512) Dispatch (Large BM=128 + Small BM=32) +// - BK=512, SK=32 +// - Direct device read +// - Swizzle decode (required by C++ dispatch) +// - Large: BM=128, BN=128, WM=4, WN=4 (512 threads) +// - Small: BM=32, BN=128, WM=1, WN=4 (128 threads) +// ============================================================ + +#include +#include +using namespace metal; + +constant constexpr short kElemsPerFrag = 8; +constant constexpr short kElemCols = 4; +constant constexpr short kElemRowsJump = 8; + +inline short2 nax_get_coord(ushort lid) { + short qid = short(lid >> 2); + short fm = ((qid & 4) | ((short(lid) >> 1) & 3)); + short fn = ((qid & 2) | (short(lid) & 1)) * 4; + return short2{fn, fm}; +} + +inline void swizzle_decode(uint2 tgid, uint swizzle_log, uint tiles_n, + thread uint &tid_y, thread uint &tid_x) { + uint tile = 1u << swizzle_log; + tid_y = tgid.y * tile + (tgid.x % tile); + tid_x = tgid.x / tile; +} + +template +inline void nax_frag_load(thread T *dst, const device T *src, int ld, short2 sc, + short off_m = 0, short off_n = 0) { + src += (sc.y + off_m) * ld + (sc.x + off_n); + for (short i = 0; i < 2; i++) + for (short j = 0; j < kElemCols; j++) + dst[i * kElemCols + j] = src[(i * kElemRowsJump) * ld + j]; +} + +inline void nax_frag_store_int32(const thread int32_t *src, device int32_t *dst, + int ld, short2 sc, short off_m, short off_n, + uint M, uint N, uint m_base, uint n_base) { + for (short i = 0; i < 2; i++) + for (short j = 0; j < kElemCols; j++) { + uint mi = m_base + sc.y + off_m + i * kElemRowsJump; + uint ni = n_base + sc.x + off_n + j; + if (mi < M && ni < N) + dst[(sc.y + off_m + i * kElemRowsJump) * ld + (sc.x + off_n + j)] = + src[i * kElemCols + j]; + } +} + +inline void nax_frag_store_dequant(const thread int32_t *src, device half *dst, + int ld, short2 sc, short off_m, short off_n, + uint M, uint N, uint m_base, uint n_base, + const device float *scale_a, + const device float *scale_w) { + for (short i = 0; i < 2; i++) + for (short j = 0; j < kElemCols; j++) { + uint mi = m_base + sc.y + off_m + i * kElemRowsJump; + uint ni = n_base + sc.x + off_n + j; + if (mi < M && ni < N) { + float val = float(src[i * kElemCols + j]) * scale_a[mi] * scale_w[ni]; + dst[(sc.y + off_m + i * kElemRowsJump) * ld + (sc.x + off_n + j)] = + half(val); + } + } +} + +// Generic GEMM with swizzle decode +template +void step4_gemm_int32_impl( + const device int8_t *A, const device int8_t *B, + device int32_t *C, uint M, uint N, uint K, + uint swizzle_log, uint tiles_m, uint tiles_n, + uint2 tgid, uint sgid, uint lid) { + + constexpr int SM = BM / WM; + constexpr int SN = BN / WN; + constexpr short TM = SM / 16; + constexpr short TN = SN / 16; + constexpr short TK = SK / 16; + + uint tid_y, tid_x; + swizzle_decode(tgid, swizzle_log, tiles_n, tid_y, tid_x); + if (tid_x >= tiles_n || tid_y >= tiles_m) return; + + short2 sc = nax_get_coord(ushort(lid)); + uint sg_row = sgid / WN; + uint sg_col = sgid % WN; + uint m_base = tid_y * BM + sg_row * SM; + uint n_base = tid_x * BN + sg_col * SN; + + const device int8_t *sg_A = A + m_base * K; + const device int8_t *sg_B = B + n_base; + + constexpr auto desc = mpp::tensor_ops::matmul2d_descriptor( + 16, 32, 16, false, false, true, + mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate); + mpp::tensor_ops::matmul2d gemm_op; + + auto ct_a = gemm_op.get_left_input_cooperative_tensor(); + auto ct_b = gemm_op.get_right_input_cooperative_tensor(); + auto ct_c = gemm_op.get_destination_cooperative_tensor(); + + int32_t c_frags[TM * TN][kElemsPerFrag]; + for (int f = 0; f < TM * TN; f++) + for (int i = 0; i < kElemsPerFrag; i++) + c_frags[f][i] = 0; + + int gemm_k_iters = int(K) / BK; + for (int kk0 = 0; kk0 < gemm_k_iters; kk0++) { + threadgroup_barrier(mem_flags::mem_none); + for (int kk1 = 0; kk1 < BK; kk1 += SK) { + int8_t a_frags[TM][TK][kElemsPerFrag]; + int8_t b_frags[TK][TN][kElemsPerFrag]; + volatile int compiler_barrier; + + for (short mm = 0; mm < TM; mm++) + for (short kk = 0; kk < TK; kk++) + nax_frag_load(a_frags[mm][kk], sg_A + kk1, int(K), sc, + short(mm * 16), short(kk * 16)); + for (short kk = 0; kk < TK; kk++) + for (short nn = 0; nn < TN; nn++) + nax_frag_load(b_frags[kk][nn], sg_B + kk1 * N, int(N), sc, + short(kk * 16), short(nn * 16)); + + for (short mm = 0; mm < TM; mm++) { + for (short nn = 0; nn < TN; nn += 2) { + for (short kk = 0; kk < TK; kk++) { + for (short i = 0; i < kElemsPerFrag; i++) ct_a[i] = a_frags[mm][kk][i]; + for (short i = 0; i < kElemsPerFrag; i++) { + ct_b[i] = b_frags[kk][nn][i]; + ct_b[kElemsPerFrag + i] = b_frags[kk][nn + 1][i]; + } + short c0 = mm * TN + nn, c1 = c0 + 1; + for (short i = 0; i < kElemsPerFrag; i++) { + ct_c[i] = c_frags[c0][i]; + ct_c[kElemsPerFrag + i] = c_frags[c1][i]; + } + gemm_op.run(ct_a, ct_b, ct_c); + for (short i = 0; i < kElemsPerFrag; i++) { + c_frags[c0][i] = ct_c[i]; + c_frags[c1][i] = ct_c[kElemsPerFrag + i]; + } + } + } + } + (void)compiler_barrier; + } + sg_A += BK; + sg_B += BK * N; + } + + // Remainder K + int rem_k = int(K) - gemm_k_iters * BK; + if (rem_k > 0) { + for (int kk1 = 0; kk1 < rem_k; kk1 += SK) { + int actual_sk = min(SK, rem_k - kk1); + int8_t a_frags[TM][TK][kElemsPerFrag]; + int8_t b_frags[TK][TN][kElemsPerFrag]; + + for (short mm = 0; mm < TM; mm++) + for (short kk = 0; kk < TK; kk++) { + for (short i = 0; i < 2; i++) + for (short j = 0; j < kElemCols; j++) { + int ki = kk * 16 + int(sc.x) + j; + if (ki < actual_sk) { + int row = int(sc.y) + mm * 16 + i * kElemRowsJump; + a_frags[mm][kk][i * kElemCols + j] = sg_A[kk1 + row * int(K) + ki]; + } else { + a_frags[mm][kk][i * kElemCols + j] = 0; + } + } + } + for (short kk = 0; kk < TK; kk++) + for (short nn = 0; nn < TN; nn++) { + for (short i = 0; i < 2; i++) + for (short j = 0; j < kElemCols; j++) { + int ki = kk * 16 + int(sc.y) + i * kElemRowsJump; + if (ki < actual_sk) { + int col = nn * 16 + int(sc.x) + j; + b_frags[kk][nn][i * kElemCols + j] = sg_B[(kk1 + ki) * int(N) + col]; + } else { + b_frags[kk][nn][i * kElemCols + j] = 0; + } + } + } + + for (short mm = 0; mm < TM; mm++) { + for (short nn = 0; nn < TN; nn += 2) { + for (short kk = 0; kk < TK; kk++) { + for (short i = 0; i < kElemsPerFrag; i++) ct_a[i] = a_frags[mm][kk][i]; + for (short i = 0; i < kElemsPerFrag; i++) { + ct_b[i] = b_frags[kk][nn][i]; + ct_b[kElemsPerFrag + i] = b_frags[kk][nn + 1][i]; + } + short c0 = mm * TN + nn, c1 = c0 + 1; + for (short i = 0; i < kElemsPerFrag; i++) { + ct_c[i] = c_frags[c0][i]; + ct_c[kElemsPerFrag + i] = c_frags[c1][i]; + } + gemm_op.run(ct_a, ct_b, ct_c); + for (short i = 0; i < kElemsPerFrag; i++) { + c_frags[c0][i] = ct_c[i]; + c_frags[c1][i] = ct_c[kElemsPerFrag + i]; + } + } + } + } + } + } + + device int32_t *D = C + m_base * N + n_base; + for (short mm = 0; mm < TM; mm++) + for (short nn = 0; nn < TN; nn++) + nax_frag_store_int32(c_frags[mm * TN + nn], D, int(N), sc, + short(mm * 16), short(nn * 16), M, N, m_base, n_base); +} + +// Dequant version +template +void step4_gemm_dequant_impl( + const device int8_t *A, const device int8_t *B, + device half *C, uint M, uint N, uint K, + const device float *scale_a, const device float *scale_w, + uint swizzle_log, uint tiles_m, uint tiles_n, + uint2 tgid, uint sgid, uint lid) { + + constexpr int SM = BM / WM; + constexpr int SN = BN / WN; + constexpr short TM = SM / 16; + constexpr short TN = SN / 16; + constexpr short TK = SK / 16; + + uint tid_y, tid_x; + swizzle_decode(tgid, swizzle_log, tiles_n, tid_y, tid_x); + if (tid_x >= tiles_n || tid_y >= tiles_m) return; + + short2 sc = nax_get_coord(ushort(lid)); + uint sg_row = sgid / WN; + uint sg_col = sgid % WN; + uint m_base = tid_y * BM + sg_row * SM; + uint n_base = tid_x * BN + sg_col * SN; + + const device int8_t *sg_A = A + m_base * K; + const device int8_t *sg_B = B + n_base; + + constexpr auto desc = mpp::tensor_ops::matmul2d_descriptor( + 16, 32, 16, false, false, true, + mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate); + mpp::tensor_ops::matmul2d gemm_op; + + auto ct_a = gemm_op.get_left_input_cooperative_tensor(); + auto ct_b = gemm_op.get_right_input_cooperative_tensor(); + auto ct_c = gemm_op.get_destination_cooperative_tensor(); + + int32_t c_frags[TM * TN][kElemsPerFrag]; + for (int f = 0; f < TM * TN; f++) + for (int i = 0; i < kElemsPerFrag; i++) + c_frags[f][i] = 0; + + int gemm_k_iters = int(K) / BK; + for (int kk0 = 0; kk0 < gemm_k_iters; kk0++) { + threadgroup_barrier(mem_flags::mem_none); + for (int kk1 = 0; kk1 < BK; kk1 += SK) { + int8_t a_frags[TM][TK][kElemsPerFrag]; + int8_t b_frags[TK][TN][kElemsPerFrag]; + volatile int compiler_barrier; + for (short mm = 0; mm < TM; mm++) + for (short kk = 0; kk < TK; kk++) + nax_frag_load(a_frags[mm][kk], sg_A + kk1, int(K), sc, + short(mm * 16), short(kk * 16)); + for (short kk = 0; kk < TK; kk++) + for (short nn = 0; nn < TN; nn++) + nax_frag_load(b_frags[kk][nn], sg_B + kk1 * N, int(N), sc, + short(kk * 16), short(nn * 16)); + for (short mm = 0; mm < TM; mm++) { + for (short nn = 0; nn < TN; nn += 2) { + for (short kk = 0; kk < TK; kk++) { + for (short i = 0; i < kElemsPerFrag; i++) ct_a[i] = a_frags[mm][kk][i]; + for (short i = 0; i < kElemsPerFrag; i++) { + ct_b[i] = b_frags[kk][nn][i]; + ct_b[kElemsPerFrag + i] = b_frags[kk][nn + 1][i]; + } + short c0 = mm * TN + nn, c1 = c0 + 1; + for (short i = 0; i < kElemsPerFrag; i++) { + ct_c[i] = c_frags[c0][i]; + ct_c[kElemsPerFrag + i] = c_frags[c1][i]; + } + gemm_op.run(ct_a, ct_b, ct_c); + for (short i = 0; i < kElemsPerFrag; i++) { + c_frags[c0][i] = ct_c[i]; + c_frags[c1][i] = ct_c[kElemsPerFrag + i]; + } + } + } + } + (void)compiler_barrier; + } + sg_A += BK; + sg_B += BK * N; + } + + int rem_k = int(K) - gemm_k_iters * BK; + if (rem_k > 0) { + for (int kk1 = 0; kk1 < rem_k; kk1 += SK) { + int actual_sk = min(SK, rem_k - kk1); + int8_t a_frags[TM][TK][kElemsPerFrag]; + int8_t b_frags[TK][TN][kElemsPerFrag]; + for (short mm = 0; mm < TM; mm++) + for (short kk = 0; kk < TK; kk++) { + for (short i = 0; i < 2; i++) + for (short j = 0; j < kElemCols; j++) { + int ki = kk * 16 + int(sc.x) + j; + if (ki < actual_sk) { + int row = int(sc.y) + mm * 16 + i * kElemRowsJump; + a_frags[mm][kk][i * kElemCols + j] = sg_A[kk1 + row * int(K) + ki]; + } else { + a_frags[mm][kk][i * kElemCols + j] = 0; + } + } + } + for (short kk = 0; kk < TK; kk++) + for (short nn = 0; nn < TN; nn++) { + for (short i = 0; i < 2; i++) + for (short j = 0; j < kElemCols; j++) { + int ki = kk * 16 + int(sc.y) + i * kElemRowsJump; + if (ki < actual_sk) { + int col = nn * 16 + int(sc.x) + j; + b_frags[kk][nn][i * kElemCols + j] = sg_B[(kk1 + ki) * int(N) + col]; + } else { + b_frags[kk][nn][i * kElemCols + j] = 0; + } + } + } + for (short mm = 0; mm < TM; mm++) { + for (short nn = 0; nn < TN; nn += 2) { + for (short kk = 0; kk < TK; kk++) { + for (short i = 0; i < kElemsPerFrag; i++) ct_a[i] = a_frags[mm][kk][i]; + for (short i = 0; i < kElemsPerFrag; i++) { + ct_b[i] = b_frags[kk][nn][i]; + ct_b[kElemsPerFrag + i] = b_frags[kk][nn + 1][i]; + } + short c0 = mm * TN + nn, c1 = c0 + 1; + for (short i = 0; i < kElemsPerFrag; i++) { + ct_c[i] = c_frags[c0][i]; + ct_c[kElemsPerFrag + i] = c_frags[c1][i]; + } + gemm_op.run(ct_a, ct_b, ct_c); + for (short i = 0; i < kElemsPerFrag; i++) { + c_frags[c0][i] = ct_c[i]; + c_frags[c1][i] = ct_c[kElemsPerFrag + i]; + } + } + } + } + } + } + + device half *D = C + m_base * N + n_base; + for (short mm = 0; mm < TM; mm++) + for (short nn = 0; nn < TN; nn++) + nax_frag_store_dequant(c_frags[mm * TN + nn], D, int(N), sc, + short(mm * 16), short(nn * 16), M, N, m_base, n_base, + scale_a, scale_w); +} + +// ============================================================ +// Entry points — Large tile +// ============================================================ +kernel void w8a8_matmul_fused_dequant( + const device int8_t *A [[buffer(0)]], const device int8_t *B [[buffer(1)]], + device half *C [[buffer(2)]], constant uint &M [[buffer(3)]], + constant uint &N [[buffer(4)]], constant uint &K [[buffer(5)]], + const device float *scale_a [[buffer(6)]], + const device float *scale_w [[buffer(7)]], + constant uint &swizzle_log [[buffer(8)]], + constant uint &tiles_m [[buffer(9)]], constant uint &tiles_n [[buffer(10)]], + uint2 tgid [[threadgroup_position_in_grid]], + uint sgid [[simdgroup_index_in_threadgroup]], + uint lid [[thread_index_in_simdgroup]]) { + step4_gemm_dequant_impl<128, 128, 512, 32, 4, 4>( + A, B, C, M, N, K, scale_a, scale_w, swizzle_log, tiles_m, tiles_n, + tgid, sgid, lid); +} + +// Small tile entry point (NEW in Step 3!) +kernel void w8a8_matmul_fused_dequant_small( + const device int8_t *A [[buffer(0)]], const device int8_t *B [[buffer(1)]], + device half *C [[buffer(2)]], constant uint &M [[buffer(3)]], + constant uint &N [[buffer(4)]], constant uint &K [[buffer(5)]], + const device float *scale_a [[buffer(6)]], + const device float *scale_w [[buffer(7)]], + constant uint &swizzle_log [[buffer(8)]], + constant uint &tiles_m [[buffer(9)]], constant uint &tiles_n [[buffer(10)]], + uint2 tgid [[threadgroup_position_in_grid]], + uint sgid [[simdgroup_index_in_threadgroup]], + uint lid [[thread_index_in_simdgroup]]) { + step4_gemm_dequant_impl<32, 128, 512, 32, 1, 4>( + A, B, C, M, N, K, scale_a, scale_w, swizzle_log, tiles_m, tiles_n, + tgid, sgid, lid); +} + +// INT32 entry points +kernel void int8_matmul_int32( + const device int8_t *A [[buffer(0)]], const device int8_t *B [[buffer(1)]], + device int32_t *C [[buffer(2)]], constant uint &M [[buffer(3)]], + constant uint &N [[buffer(4)]], constant uint &K [[buffer(5)]], + constant uint &swizzle_log [[buffer(6)]], + constant uint &tiles_m [[buffer(7)]], constant uint &tiles_n [[buffer(8)]], + uint2 tgid [[threadgroup_position_in_grid]], + uint sgid [[simdgroup_index_in_threadgroup]], + uint lid [[thread_index_in_simdgroup]]) { + step4_gemm_int32_impl<128, 128, 512, 32, 4, 4>( + A, B, C, M, N, K, swizzle_log, tiles_m, tiles_n, tgid, sgid, lid); +} + +kernel void int8_matmul_int32_small( + const device int8_t *A [[buffer(0)]], const device int8_t *B [[buffer(1)]], + device int32_t *C [[buffer(2)]], constant uint &M [[buffer(3)]], + constant uint &N [[buffer(4)]], constant uint &K [[buffer(5)]], + constant uint &swizzle_log [[buffer(6)]], + constant uint &tiles_m [[buffer(7)]], constant uint &tiles_n [[buffer(8)]], + uint2 tgid [[threadgroup_position_in_grid]], + uint sgid [[simdgroup_index_in_threadgroup]], + uint lid [[thread_index_in_simdgroup]]) { + step4_gemm_int32_impl<32, 128, 512, 32, 1, 4>( + A, B, C, M, N, K, swizzle_log, tiles_m, tiles_n, tgid, sgid, lid); +} diff --git a/dev/step_kernels/step4/w8a8_quantize.metal b/dev/step_kernels/step4/w8a8_quantize.metal new file mode 100644 index 0000000..f9ae0be --- /dev/null +++ b/dev/step_kernels/step4/w8a8_quantize.metal @@ -0,0 +1,87 @@ +// ============================================================ +// Per-token quantization: FP16 → INT8 + float32 scale +// Target: Apple M5, Metal 4 +// +// Each threadgroup handles one row (one token). +// Threads cooperate to find absmax via simdgroup reduce, +// then quantize in parallel. +// +// Host dispatch: +// threadgroup = (min(256, ceil(K/32)*32), 1, 1) +// grid = (M, 1, 1) +// ============================================================ + +#include +using namespace metal; + +kernel void +quantize_per_token(const device half *X [[buffer(0)]], // [M, K] FP16 input + device int8_t *A [[buffer(1)]], // [M, K] INT8 output + device float *scale [[buffer(2)]], // [M] float32 scale + constant uint &M [[buffer(3)]], + constant uint &K [[buffer(4)]], + uint gid [[threadgroup_position_in_grid]], // row index + uint lid [[thread_index_in_threadgroup]], + uint tg_size [[threads_per_threadgroup]]) { + if (gid >= M) { + return; + } + + const device half *row_in = X + gid * K; + device int8_t *row_out = A + gid * K; + + // Step 1: Find local absmax + float local_max = 0.0f; + for (uint i = lid; i < K; i += tg_size) { + float v = abs(float(row_in[i])); + local_max = max(local_max, v); + } + + // Step 2: Simdgroup reduce max + float sg_max = simd_max(local_max); + + // Step 3: Threadgroup reduce across simdgroups via shared memory + threadgroup float sg_maxes[8]; // up to 8 simdgroups (256/32) + threadgroup float shared_scale; + threadgroup float shared_inv_scale; + uint sg_id = lid / 32; + uint sg_lid = lid % 32; + if (sg_lid == 0) { + sg_maxes[sg_id] = sg_max; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + // Final reduce (first simdgroup only) + if (sg_id == 0) { + float row_max = 0.0f; + uint num_sgs = (tg_size + 31) / 32; + if (sg_lid < num_sgs) { + row_max = sg_maxes[sg_lid]; + } + row_max = simd_max(row_max); + + // Compute and broadcast scale + float s = row_max / 255.0f; + if (s == 0.0f) { + s = 1.0f; + } + + if (sg_lid == 0) { + shared_scale = s; + shared_inv_scale = 1.0f / s; + // Store scale to output + scale[gid] = s; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + // Step 4: All threads read broadcasted scale + float inv_s = shared_inv_scale; + + // Step 5: Quantize + for (uint i = lid; i < K; i += tg_size) { + float v = float(row_in[i]) * inv_s; + v = clamp(round(v), -128.0f, 127.0f); + row_out[i] = int8_t(v); + } +} diff --git a/dev/step_kernels/step5/w8a8_matmul.metal b/dev/step_kernels/step5/w8a8_matmul.metal new file mode 100644 index 0000000..657e3a4 --- /dev/null +++ b/dev/step_kernels/step5/w8a8_matmul.metal @@ -0,0 +1,475 @@ +// ============================================================ +// W8A8 INT8×INT8→INT32 TensorOps GEMM +// Target: Apple M5 (G17G), Metal 4 +// +// Variants: +// - fused dequant: INT8×INT8→FP16, with per-token/per-channel scales +// - raw INT32: INT8×INT8→INT32, no scale (pure integer GEMM) +// Multi-config: large (BM=128) and small (BM=32) tiles +// Swizzle dispatch for L2 cache locality +// +// matmul2d(16,32,16) via MPP cooperative_tensor +// ============================================================ + +#include +#include +using namespace metal; + +// ── NAXFrag layout constants ──────────────────────────────────── +constant constexpr short kElemsPerFrag = 8; +constant constexpr short kElemCols = 4; +constant constexpr short kElemRowsJump = 8; + +// ── NAXFrag coordinate mapping ────────────────────────────────── +inline short2 nax_get_coord(ushort lid) { + short qid = short(lid >> 2); + short fm = ((qid & 4) | ((short(lid) >> 1) & 3)); + short fn = ((qid & 2) | (short(lid) & 1)) * 4; + return short2{fn, fm}; +} + +// ── Fragment load: device → register ──────────────────────────── +template +inline void nax_frag_load(thread T *dst, const device T *src, int ld, short2 sc, + short off_m = 0, short off_n = 0) { + src += (sc.y + off_m) * ld + (sc.x + off_n); + for (short i = 0; i < 2; i++) { + for (short j = 0; j < kElemCols; j++) { + dst[i * kElemCols + j] = src[(i * kElemRowsJump) * ld + j]; + } + } +} + +// ── Fragment store: raw INT32 (no dequant) ─────────────────── +inline void nax_frag_store_int32(const thread int32_t *src, device int32_t *dst, + int ld, short2 sc, short off_m, short off_n, + uint M, uint N, uint m_base, uint n_base) { + for (short i = 0; i < 2; i++) { + for (short j = 0; j < kElemCols; j++) { + uint mi = m_base + sc.y + off_m + i * kElemRowsJump; + uint ni = n_base + sc.x + off_n + j; + if (mi < M && ni < N) { + dst[(sc.y + off_m + i * kElemRowsJump) * ld + (sc.x + off_n + j)] = + src[i * kElemCols + j]; + } + } + } +} + +// ── Fragment store with bounds check and dequant ──────────────── +inline void nax_frag_store_dequant(const thread int32_t *src, device half *dst, + int ld, short2 sc, short off_m, short off_n, + uint M, uint N, uint m_base, uint n_base, + const device float *scale_a, + const device float *scale_w) { + for (short i = 0; i < 2; i++) { + for (short j = 0; j < kElemCols; j++) { + uint mi = m_base + sc.y + off_m + i * kElemRowsJump; + uint ni = n_base + sc.x + off_n + j; + if (mi < M && ni < N) { + float val = float(src[i * kElemCols + j]) * scale_a[mi] * scale_w[ni]; + dst[(sc.y + off_m + i * kElemRowsJump) * ld + (sc.x + off_n + j)] = + half(val); + } + } + } +} + +// ── Generic GEMM kernel ───────────────────────────────────────── +// Template params: BM, BN, BK, SK, WM, WN +// Each SG computes SM×SN = (BM/WM) × (BN/WN) output +// SM and SN must be 32 (2×2 of 16×16 fragments) +// +// swizzle_log: passed via constant buffer +// tid_y = (tgid.y << swizzle_log) + (tgid.x & ((1<> swizzle_log +template +void w8a8_gemm_impl(const device int8_t *A, const device int8_t *B, + device half *C, uint M, uint N, uint K, + const device float *scale_a, const device float *scale_w, + uint swizzle_log, uint tiles_m, uint tiles_n, uint2 tgid, + uint sgid, uint lid) { + constexpr int SM = BM / WM; // 32 + constexpr int SN = BN / WN; // 32 + constexpr short TM = SM / 16; // 2 + constexpr short TN = SN / 16; // 2 + constexpr short TK = SK / 16; // 2 + + // Swizzle decode + uint tid_y = (tgid.y << swizzle_log) + (tgid.x & ((1u << swizzle_log) - 1u)); + uint tid_x = tgid.x >> swizzle_log; + + // Bounds check (swizzle can create out-of-bounds tiles) + if (tid_x >= tiles_n || tid_y >= tiles_m) { + return; + } + + short2 sc = nax_get_coord(ushort(lid)); + uint sg_row = sgid / WN; + uint sg_col = sgid % WN; + uint m_base = tid_y * BM + sg_row * SM; + uint n_base = tid_x * BN + sg_col * SN; + + const device int8_t *sg_A = A + m_base * K; + const device int8_t *sg_B = B + n_base; + + constexpr auto desc = mpp::tensor_ops::matmul2d_descriptor( + 16, 32, 16, false, false, true, + mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate); + mpp::tensor_ops::matmul2d gemm_op; + + auto ct_a = + gemm_op.get_left_input_cooperative_tensor(); + auto ct_b = + gemm_op.get_right_input_cooperative_tensor(); + auto ct_c = + gemm_op.get_destination_cooperative_tensor(); + + int32_t c_frags[TM * TN][kElemsPerFrag]; + for (int f = 0; f < TM * TN; f++) { + for (int i = 0; i < kElemsPerFrag; i++) { + c_frags[f][i] = 0; + } + } + + // ── Main K loop ───────────────────────────────────────────── + int gemm_k_iters = int(K) / BK; + + for (int kk0 = 0; kk0 < gemm_k_iters; kk0++) { + threadgroup_barrier(mem_flags::mem_none); + + for (int kk1 = 0; kk1 < BK; kk1 += SK) { + int8_t a_frags[TM][TK][kElemsPerFrag]; + int8_t b_frags[TK][TN][kElemsPerFrag]; + volatile int compiler_barrier; + + for (short mm = 0; mm < TM; mm++) { + for (short kk = 0; kk < TK; kk++) { + nax_frag_load(a_frags[mm][kk], sg_A + kk1, int(K), sc, short(mm * 16), + short(kk * 16)); + } + } + for (short kk = 0; kk < TK; kk++) { + for (short nn = 0; nn < TN; nn++) { + nax_frag_load(b_frags[kk][nn], sg_B + kk1 * N, int(N), sc, + short(kk * 16), short(nn * 16)); + } + } + for (short mm = 0; mm < TM; mm++) { + for (short nn = 0; nn < TN; nn += 2) { + for (short kk = 0; kk < TK; kk++) { + for (short i = 0; i < kElemsPerFrag; i++) { + ct_a[i] = a_frags[mm][kk][i]; + } + for (short i = 0; i < kElemsPerFrag; i++) { + ct_b[i] = b_frags[kk][nn][i]; + ct_b[kElemsPerFrag + i] = b_frags[kk][nn + 1][i]; + } + short c0 = mm * TN + nn, c1 = c0 + 1; + for (short i = 0; i < kElemsPerFrag; i++) { + ct_c[i] = c_frags[c0][i]; + ct_c[kElemsPerFrag + i] = c_frags[c1][i]; + } + gemm_op.run(ct_a, ct_b, ct_c); + for (short i = 0; i < kElemsPerFrag; i++) { + c_frags[c0][i] = ct_c[i]; + c_frags[c1][i] = ct_c[kElemsPerFrag + i]; + } + } + } + } + (void)compiler_barrier; + } + + sg_A += BK; + sg_B += BK * N; + } + + // ── Remainder K ───────────────────────────────────────────── + int rem_k = int(K) - gemm_k_iters * BK; + for (int kk1 = 0; kk1 < rem_k; kk1 += 16) { + int8_t a_frag[TM][kElemsPerFrag]; + int8_t b_frag[TN][kElemsPerFrag]; + short psk = short(max(0, rem_k - kk1)); + + for (short mm = 0; mm < TM; mm++) { + const device int8_t *ptr = sg_A + kk1 + (sc.y + mm * 16) * K + sc.x; + for (short i = 0; i < 2; i++) { + for (short j = 0; j < kElemCols; j++) { + short ki = short(sc.x + j); + a_frag[mm][i * kElemCols + j] = + (ki < psk) ? ptr[(i * kElemRowsJump) * K + j] : int8_t(0); + } + } + } + + for (short nn = 0; nn < TN; nn++) { + const device int8_t *ptr = sg_B + kk1 * N + nn * 16 + sc.y * N + sc.x; + for (short i = 0; i < 2; i++) { + for (short j = 0; j < kElemCols; j++) { + short ki = short(sc.y + i * kElemRowsJump); + b_frag[nn][i * kElemCols + j] = + (ki < psk) ? ptr[(i * kElemRowsJump) * N + j] : int8_t(0); + } + } + } + + for (short mm = 0; mm < TM; mm++) { + for (short i = 0; i < kElemsPerFrag; i++) { + ct_a[i] = a_frag[mm][i]; + } + for (short i = 0; i < kElemsPerFrag; i++) { + ct_b[i] = b_frag[0][i]; + ct_b[kElemsPerFrag + i] = b_frag[1][i]; + } + short c0 = mm * TN, c1 = c0 + 1; + for (short i = 0; i < kElemsPerFrag; i++) { + ct_c[i] = c_frags[c0][i]; + ct_c[kElemsPerFrag + i] = c_frags[c1][i]; + } + gemm_op.run(ct_a, ct_b, ct_c); + for (short i = 0; i < kElemsPerFrag; i++) { + c_frags[c0][i] = ct_c[i]; + c_frags[c1][i] = ct_c[kElemsPerFrag + i]; + } + } + } + + // ── Store with fused dequant ──────────────────────────────── + device half *D = C + m_base * N + n_base; + for (short mm = 0; mm < TM; mm++) { + for (short nn = 0; nn < TN; nn++) { + nax_frag_store_dequant(c_frags[mm * TN + nn], D, int(N), sc, + short(mm * 16), short(nn * 16), M, N, m_base, + n_base, scale_a, scale_w); + } + } +} + +// ============================================================ +// Kernel entry points +// ============================================================ + +// ── Large tile: BM=128, BN=128, BK=512, WM=4, WN=4 ──────────── +// 16 SG, 512 threads/TG. Best for M≥128. +kernel void w8a8_matmul_fused_dequant( + const device int8_t *A [[buffer(0)]], const device int8_t *B [[buffer(1)]], + device half *C [[buffer(2)]], constant uint &M [[buffer(3)]], + constant uint &N [[buffer(4)]], constant uint &K [[buffer(5)]], + const device float *scale_a [[buffer(6)]], + const device float *scale_w [[buffer(7)]], + constant uint &swizzle_log [[buffer(8)]], + constant uint &tiles_m [[buffer(9)]], constant uint &tiles_n [[buffer(10)]], + uint2 tgid [[threadgroup_position_in_grid]], + uint sgid [[simdgroup_index_in_threadgroup]], + uint lid [[thread_index_in_simdgroup]]) { + w8a8_gemm_impl<128, 128, 512, 32, 4, 4>(A, B, C, M, N, K, scale_a, scale_w, + swizzle_log, tiles_m, tiles_n, tgid, + sgid, lid); +} + +// ── Small tile: BM=32, BN=128, BK=512, WM=1, WN=4 ───────────── +// 4 SG, 128 threads/TG. Optimized for M≤64. +kernel void w8a8_matmul_fused_dequant_small( + const device int8_t *A [[buffer(0)]], const device int8_t *B [[buffer(1)]], + device half *C [[buffer(2)]], constant uint &M [[buffer(3)]], + constant uint &N [[buffer(4)]], constant uint &K [[buffer(5)]], + const device float *scale_a [[buffer(6)]], + const device float *scale_w [[buffer(7)]], + constant uint &swizzle_log [[buffer(8)]], + constant uint &tiles_m [[buffer(9)]], constant uint &tiles_n [[buffer(10)]], + uint2 tgid [[threadgroup_position_in_grid]], + uint sgid [[simdgroup_index_in_threadgroup]], + uint lid [[thread_index_in_simdgroup]]) { + w8a8_gemm_impl<32, 128, 512, 32, 1, 4>(A, B, C, M, N, K, scale_a, scale_w, + swizzle_log, tiles_m, tiles_n, tgid, + sgid, lid); +} + +// ============================================================ +// Raw INT32 GEMM impl (no dequant, no scales) +// ============================================================ + +template +void w8a8_gemm_int32_impl(const device int8_t *A, const device int8_t *B, + device int32_t *C, uint M, uint N, uint K, + uint swizzle_log, uint tiles_m, uint tiles_n, + uint2 tgid, uint sgid, uint lid) { + constexpr int SM = BM / WM; + constexpr int SN = BN / WN; + constexpr short TM = SM / 16; + constexpr short TN = SN / 16; + constexpr short TK = SK / 16; + + uint tid_y = (tgid.y << swizzle_log) + (tgid.x & ((1u << swizzle_log) - 1u)); + uint tid_x = tgid.x >> swizzle_log; + if (tid_x >= tiles_n || tid_y >= tiles_m) { + return; + } + + short2 sc = nax_get_coord(ushort(lid)); + uint sg_row = sgid / WN; + uint sg_col = sgid % WN; + uint m_base = tid_y * BM + sg_row * SM; + uint n_base = tid_x * BN + sg_col * SN; + + const device int8_t *sg_A = A + m_base * K; + const device int8_t *sg_B = B + n_base; + + constexpr auto desc = mpp::tensor_ops::matmul2d_descriptor( + 16, 32, 16, false, false, true, + mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate); + mpp::tensor_ops::matmul2d gemm_op; + + auto ct_a = + gemm_op.get_left_input_cooperative_tensor(); + auto ct_b = + gemm_op.get_right_input_cooperative_tensor(); + auto ct_c = + gemm_op.get_destination_cooperative_tensor(); + + int32_t c_frags[TM * TN][kElemsPerFrag]; + for (int f = 0; f < TM * TN; f++) { + for (int i = 0; i < kElemsPerFrag; i++) { + c_frags[f][i] = 0; + } + } + + int gemm_k_iters = int(K) / BK; + for (int kk0 = 0; kk0 < gemm_k_iters; kk0++) { + threadgroup_barrier(mem_flags::mem_none); + for (int kk1 = 0; kk1 < BK; kk1 += SK) { + int8_t a_frags[TM][TK][kElemsPerFrag]; + int8_t b_frags[TK][TN][kElemsPerFrag]; + volatile int compiler_barrier; + + for (short mm = 0; mm < TM; mm++) { + for (short kk = 0; kk < TK; kk++) { + nax_frag_load(a_frags[mm][kk], sg_A + kk1, int(K), sc, short(mm * 16), + short(kk * 16)); + } + } + for (short kk = 0; kk < TK; kk++) { + for (short nn = 0; nn < TN; nn++) { + nax_frag_load(b_frags[kk][nn], sg_B + kk1 * N, int(N), sc, + short(kk * 16), short(nn * 16)); + } + } + + for (short mm = 0; mm < TM; mm++) { + for (short nn = 0; nn < TN; nn += 2) { + for (short kk = 0; kk < TK; kk++) { + for (short i = 0; i < kElemsPerFrag; i++) { + ct_a[i] = a_frags[mm][kk][i]; + } + for (short i = 0; i < kElemsPerFrag; i++) { + ct_b[i] = b_frags[kk][nn][i]; + ct_b[kElemsPerFrag + i] = b_frags[kk][nn + 1][i]; + } + short c0 = mm * TN + nn, c1 = c0 + 1; + for (short i = 0; i < kElemsPerFrag; i++) { + ct_c[i] = c_frags[c0][i]; + ct_c[kElemsPerFrag + i] = c_frags[c1][i]; + } + gemm_op.run(ct_a, ct_b, ct_c); + for (short i = 0; i < kElemsPerFrag; i++) { + c_frags[c0][i] = ct_c[i]; + c_frags[c1][i] = ct_c[kElemsPerFrag + i]; + } + } + } + } + (void)compiler_barrier; + } + sg_A += BK; + sg_B += BK * N; + } + + // Remainder K + int rem_k = int(K) - gemm_k_iters * BK; + for (int kk1 = 0; kk1 < rem_k; kk1 += 16) { + int8_t a_frag[TM][kElemsPerFrag]; + int8_t b_frag[TN][kElemsPerFrag]; + short psk = short(max(0, rem_k - kk1)); + for (short mm = 0; mm < TM; mm++) { + const device int8_t *ptr = sg_A + kk1 + (sc.y + mm * 16) * K + sc.x; + for (short i = 0; i < 2; i++) { + for (short j = 0; j < kElemCols; j++) { + short ki = short(sc.x + j); + a_frag[mm][i * kElemCols + j] = + (ki < psk) ? ptr[(i * kElemRowsJump) * K + j] : int8_t(0); + } + } + } + for (short nn = 0; nn < TN; nn++) { + const device int8_t *ptr = sg_B + kk1 * N + nn * 16 + sc.y * N + sc.x; + for (short i = 0; i < 2; i++) { + for (short j = 0; j < kElemCols; j++) { + short ki = short(sc.y + i * kElemRowsJump); + b_frag[nn][i * kElemCols + j] = + (ki < psk) ? ptr[(i * kElemRowsJump) * N + j] : int8_t(0); + } + } + } + for (short mm = 0; mm < TM; mm++) { + for (short i = 0; i < kElemsPerFrag; i++) { + ct_a[i] = a_frag[mm][i]; + } + for (short i = 0; i < kElemsPerFrag; i++) { + ct_b[i] = b_frag[0][i]; + ct_b[kElemsPerFrag + i] = b_frag[1][i]; + } + short c0 = mm * TN, c1 = c0 + 1; + for (short i = 0; i < kElemsPerFrag; i++) { + ct_c[i] = c_frags[c0][i]; + ct_c[kElemsPerFrag + i] = c_frags[c1][i]; + } + gemm_op.run(ct_a, ct_b, ct_c); + for (short i = 0; i < kElemsPerFrag; i++) { + c_frags[c0][i] = ct_c[i]; + c_frags[c1][i] = ct_c[kElemsPerFrag + i]; + } + } + } + + // Store raw INT32 + device int32_t *D = C + m_base * N + n_base; + for (short mm = 0; mm < TM; mm++) { + for (short nn = 0; nn < TN; nn++) { + nax_frag_store_int32(c_frags[mm * TN + nn], D, int(N), sc, short(mm * 16), + short(nn * 16), M, N, m_base, n_base); + } + } +} + +// ============================================================ +// Kernel entry points — raw INT32 output +// ============================================================ + +kernel void int8_matmul_int32( + const device int8_t *A [[buffer(0)]], const device int8_t *B [[buffer(1)]], + device int32_t *C [[buffer(2)]], constant uint &M [[buffer(3)]], + constant uint &N [[buffer(4)]], constant uint &K [[buffer(5)]], + constant uint &swizzle_log [[buffer(6)]], + constant uint &tiles_m [[buffer(7)]], constant uint &tiles_n [[buffer(8)]], + uint2 tgid [[threadgroup_position_in_grid]], + uint sgid [[simdgroup_index_in_threadgroup]], + uint lid [[thread_index_in_simdgroup]]) { + w8a8_gemm_int32_impl<128, 128, 512, 32, 4, 4>( + A, B, C, M, N, K, swizzle_log, tiles_m, tiles_n, tgid, sgid, lid); +} + +kernel void int8_matmul_int32_small( + const device int8_t *A [[buffer(0)]], const device int8_t *B [[buffer(1)]], + device int32_t *C [[buffer(2)]], constant uint &M [[buffer(3)]], + constant uint &N [[buffer(4)]], constant uint &K [[buffer(5)]], + constant uint &swizzle_log [[buffer(6)]], + constant uint &tiles_m [[buffer(7)]], constant uint &tiles_n [[buffer(8)]], + uint2 tgid [[threadgroup_position_in_grid]], + uint sgid [[simdgroup_index_in_threadgroup]], + uint lid [[thread_index_in_simdgroup]]) { + w8a8_gemm_int32_impl<32, 128, 512, 32, 1, 4>( + A, B, C, M, N, K, swizzle_log, tiles_m, tiles_n, tgid, sgid, lid); +} diff --git a/dev/step_kernels/step5/w8a8_quantize.metal b/dev/step_kernels/step5/w8a8_quantize.metal new file mode 100644 index 0000000..f9ae0be --- /dev/null +++ b/dev/step_kernels/step5/w8a8_quantize.metal @@ -0,0 +1,87 @@ +// ============================================================ +// Per-token quantization: FP16 → INT8 + float32 scale +// Target: Apple M5, Metal 4 +// +// Each threadgroup handles one row (one token). +// Threads cooperate to find absmax via simdgroup reduce, +// then quantize in parallel. +// +// Host dispatch: +// threadgroup = (min(256, ceil(K/32)*32), 1, 1) +// grid = (M, 1, 1) +// ============================================================ + +#include +using namespace metal; + +kernel void +quantize_per_token(const device half *X [[buffer(0)]], // [M, K] FP16 input + device int8_t *A [[buffer(1)]], // [M, K] INT8 output + device float *scale [[buffer(2)]], // [M] float32 scale + constant uint &M [[buffer(3)]], + constant uint &K [[buffer(4)]], + uint gid [[threadgroup_position_in_grid]], // row index + uint lid [[thread_index_in_threadgroup]], + uint tg_size [[threads_per_threadgroup]]) { + if (gid >= M) { + return; + } + + const device half *row_in = X + gid * K; + device int8_t *row_out = A + gid * K; + + // Step 1: Find local absmax + float local_max = 0.0f; + for (uint i = lid; i < K; i += tg_size) { + float v = abs(float(row_in[i])); + local_max = max(local_max, v); + } + + // Step 2: Simdgroup reduce max + float sg_max = simd_max(local_max); + + // Step 3: Threadgroup reduce across simdgroups via shared memory + threadgroup float sg_maxes[8]; // up to 8 simdgroups (256/32) + threadgroup float shared_scale; + threadgroup float shared_inv_scale; + uint sg_id = lid / 32; + uint sg_lid = lid % 32; + if (sg_lid == 0) { + sg_maxes[sg_id] = sg_max; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + // Final reduce (first simdgroup only) + if (sg_id == 0) { + float row_max = 0.0f; + uint num_sgs = (tg_size + 31) / 32; + if (sg_lid < num_sgs) { + row_max = sg_maxes[sg_lid]; + } + row_max = simd_max(row_max); + + // Compute and broadcast scale + float s = row_max / 255.0f; + if (s == 0.0f) { + s = 1.0f; + } + + if (sg_lid == 0) { + shared_scale = s; + shared_inv_scale = 1.0f / s; + // Store scale to output + scale[gid] = s; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + // Step 4: All threads read broadcasted scale + float inv_s = shared_inv_scale; + + // Step 5: Quantize + for (uint i = lid; i < K; i += tg_size) { + float v = float(row_in[i]) * inv_s; + v = clamp(round(v), -128.0f, 127.0f); + row_out[i] = int8_t(v); + } +} diff --git a/examples/basic_usage.py b/examples/basic_usage.py new file mode 100644 index 0000000..1f8d242 --- /dev/null +++ b/examples/basic_usage.py @@ -0,0 +1,44 @@ +#!/usr/bin/env python3 +"""Basic usage example for Cider. + +Demonstrates W8A8 and W4A8 quantized linear layers. +All operations return lazy mx.array — evaluated on mx.eval(). +""" +import numpy as np +import mlx.core as mx +from cider import W8A8Linear, W4A8Linear, is_available + +# Check hardware +if not is_available(): + print("INT8 TensorOps not available (requires Apple M5+)") + exit(1) +print("INT8 TensorOps: available") + +# ── Prepare weights ────────────────────────────────────────────── +K, N = 4096, 4096 +W_fp16 = np.random.randn(K, N).astype(np.float16) + +# ── W8A8 ───────────────────────────────────────────────────────── +w8a8 = W8A8Linear.from_weights(W_fp16) +print(f"W8A8: [{w8a8.input_dims}, {w8a8.output_dims}]") + +x = mx.random.normal((32, K)).astype(mx.float16) +y = w8a8(x) # lazy — not computed yet +mx.eval(y) # now the GPU runs +print(f"W8A8 output: {y.shape}, dtype={y.dtype}") + +# ── W4A8 ───────────────────────────────────────────────────────── +w4a8 = W4A8Linear.from_weights(W_fp16) +print(f"W4A8: [{w4a8.input_dims}, {w4a8.output_dims}]") +print(f"W4A8 weight storage: {w4a8.packed_weight.nbytes / 1024:.0f} KB " + f"(vs W8A8: {w8a8.weight.nbytes / 1024:.0f} KB)") + +y4 = w4a8(x) +mx.eval(y4) +print(f"W4A8 output: {y4.shape}, dtype={y4.dtype}") + +# ── Lazy composition ───────────────────────────────────────────── +# Multiple ops compose into a single MLX graph, evaluated together +z = w8a8(x) + w4a8(x) # builds graph, no GPU work yet +mx.eval(z) # single evaluation +print(f"Composed output: {z.shape}") diff --git a/experimental/README.md b/experimental/README.md new file mode 100644 index 0000000..fa41305 --- /dev/null +++ b/experimental/README.md @@ -0,0 +1,94 @@ +# ANE Split — ANE+GPU Tensor Parallelism for Prefill Acceleration +*Status: experimental research prototype. Uses private ANE APIs and is not part of Cider's stable inference path.* + +Split each linear layer's GEMM along output channels: **ANE** and **GPU** computes running concurrently. This explores whether the otherwise idle Apple Neural Engine can help accelerate prefill, with minimal degradation observed in our current tests. + + +> **Platform:** Apple M4 (tested). M5 ANE API changes may cause failures — not yet validated. + +## How It Works + +``` +Input x ──┬── ANE (r output channels, FP32, private API) ──┐ + │ ├── concat → output + └── GPU (1 - r output channels, FP16, MLX matmul) ──┘ +``` + +1. **SplitLinear** wraps each `nn.Linear` / `nn.QuantizedLinear` +2. In **prefill mode** (seq ≥ 192): split path with ANE+GPU concurrency +3. In **decode mode**: falls back to original `nn.Linear` on GPU (zero overhead) +4. Same-input projections (Q/K/V, Gate/Up) share input preparation via `_InputGroup` + +### Automatic Layer Routing + +| Layer Type | Routing | Reason | +|------------|---------|--------| +| Q, K, V, O projections | ANE+GPU split | Expand: IC → OC, ANE efficient | +| Gate, Up projections | ANE+GPU split | Expand: IC → OC, ANE efficient | +| Down projection | GPU only | Narrow: IC > 2×OC, ANE inefficient | +| Short sequences (< 192) | GPU only | Split overhead > benefit | + +## Performance (Apple M4, Qwen3-VL-2B) + +| seq | W8A16 GPU | SplitLinear | Speedup vs W8A16 | +|-----|----------|-----------|-------------| +| 512 | 639.9 ms | **615.9 ms** | **1.039×** | +| 1024 | 1348.6 ms | **1156.9 ms** | **1.17×** | + +- Accuracy: cos ≈ 1.0, top-1 match = 100% +- 168 layers split (28 layers × 6 projections), 28 GPU-only (down_proj) + + +## Quick Start + +```python +from split_linear import patch_model, SplitLinear + +# Load any MLX VLM model +from mlx_vlm.utils import load as vlm_load +model, processor = vlm_load("path/to/model") + +# Patch all linear layers (one-liner) +bridge = patch_model(model, seq=512) + +# Enable split for prefill, disable for decode +SplitLinear.set_prefill(True) # prefill: ANE+GPU parallel +# ... run prefill ... +SplitLinear.set_prefill(False) # decode: original GPU path +``` + +## Benchmark + +```bash +# Default: seq=512 +python3 bench.py + +# Custom seq length +python3 bench.py 1024 + +# Custom model path +MODEL_PATH=/path/to/model python3 bench.py 512 +``` + +## Building the ANE Bridge + +To build libane_bridge_v6.dylib from the source: + +```bash +clang -shared -O2 -framework Foundation -framework CoreML \ + -framework Accelerate -o libane_bridge_v6.dylib libane_bridge_v6.m +``` + +> Requires macOS with ANE private frameworks. Uses undocumented `_ANEClient` API. + +## Limitations + +- **M4 only** — M5 ANE internal changes may break the private API bridge +- **Fixed sequence length** — ANE models are compiled for a specific seq; re-patch needed for different lengths +- **FP32 on ANE** — ANE operates in FP32 (no INT8/FP16 GEMM); benefit comes from parallelism, not precision +- **Memory overhead** — ANE models consume additional system memory (~200MB for 2B model) +- **No decode benefit** — Decode is single-token, falls back to GPU (no split overhead) +- **No E2E benefit** - MLX natively employs lazy evaluation to reduce synchronization overhead. In end-to-end testing, our hybrid approach currently shows no advantage because we haven't yet implemented it with MLX's lazy evaluation. +## License + +MIT diff --git a/experimental/bench.py b/experimental/bench.py new file mode 100644 index 0000000..8c7fdb4 --- /dev/null +++ b/experimental/bench.py @@ -0,0 +1,108 @@ +#!/usr/bin/env python3 +""" +bench_w8a16_splitlinear.py — W8A16 vs SplitLinear prefill benchmark. + +Usage: + python3 bench_w8a16_splitlinear.py [seq_len] # default 512 + +Requires: + - mlx_vlm with Qwen3-VL-2B-Instruct-8bit model + - split_linear.py (V4) + libane_bridge_v6.dylib in same directory +""" +import sys, time, os +import numpy as np +import mlx.core as mx +from mlx_vlm.utils import load as vlm_load +from mlx_vlm.models.cache import KVCache + +# ─── Config ─── +SEQ = int(sys.argv[1]) if len(sys.argv) > 1 else 512 +W8A16_MODEL = '~/Downloads/sft_baseline_v2_w8a16' +N_WARMUP = 3 +N_BENCH = 8 + +def bench_forward(lang, ids, pos, n_layers, reps): + ts = [] + for _ in range(reps): + cache = [KVCache() for _ in range(n_layers)] + t0 = time.perf_counter() + mx.eval(lang(ids, cache=cache, position_ids=pos).logits) + ts.append((time.perf_counter() - t0) * 1000) + return ts + +def main(): + print(f"\n{'='*60}") + print(f" W8A16 vs SplitLinear Prefill Benchmark (seq={SEQ})") + print(f"{'='*60}") + print(f" Model: {W8A16_MODEL}") + print(f" Warmup: {N_WARMUP}, Bench: {N_BENCH}\n") + + # Load model + print("[1/4] Loading W8A16 model...") + model, _ = vlm_load(W8A16_MODEL) + lang = model.language_model + N = lang.args.num_hidden_layers + print(f" {N} layers loaded\n") + + # Prepare inputs + ids = mx.ones((1, SEQ), dtype=mx.int32) + pos = mx.broadcast_to( + mx.arange(SEQ).reshape(1, SEQ)[None, :, :], + (3, 1, SEQ) + ) + mx.eval(ids, pos) + + # ─── W8A16 Baseline ─── + print("[2/4] W8A16 GPU baseline") + bench_forward(lang, ids, pos, N, N_WARMUP) # warmup + ts_w8 = bench_forward(lang, ids, pos, N, N_BENCH) + for i, t in enumerate(ts_w8): + print(f" Run {i+1}: {t:.1f}ms") + med_w8 = float(np.median(ts_w8)) + print(f" Median: {med_w8:.1f}ms\n") + + # ─── Reference logits ─── + print("[3/4] Computing reference logits for accuracy check...") + c_ref = [KVCache() for _ in range(N)] + ref = np.array(lang(ids, cache=c_ref, position_ids=pos).logits.astype(mx.float32)) + + # ─── SplitLinear ─── + print("[4/4] Patch with SplitLinear + benchmark") + from split_linear import patch_model, SplitLinear + bridge = patch_model(model, SEQ) + SplitLinear.set_prefill(True) + + # Warmup + bench_forward(lang, ids, pos, N, N_WARMUP) + + # Accuracy + c_hyb = [KVCache() for _ in range(N)] + hyb = np.array(lang(ids, cache=c_hyb, position_ids=pos).logits.astype(mx.float32)) + cos = float(np.dot(ref.flatten(), hyb.flatten()) / + (np.linalg.norm(ref.flatten()) * np.linalg.norm(hyb.flatten()) + 1e-12)) + top1 = float((ref.argmax(-1) == hyb.argmax(-1)).mean() * 100) + print(f" Accuracy: cos={cos:.6f}, top1={top1:.1f}%") + + # Benchmark + ts_sp = bench_forward(lang, ids, pos, N, N_BENCH) + for i, t in enumerate(ts_sp): + print(f" Run {i+1}: {t:.1f}ms") + med_sp = float(np.median(ts_sp)) + + # ─── Summary ─── + speedup = med_w8 / med_sp + delta = med_sp - med_w8 + print(f"\n{'='*60}") + print(f" W8A16 GPU: {med_w8:.1f}ms") + print(f" SplitLinear: {med_sp:.1f}ms ({speedup:.3f}x)") + print(f" Delta: {delta:+.1f}ms") + print(f" Accuracy: cos={cos:.6f} top1={top1:.1f}%") + print(f"{'='*60}") + + if speedup >= 1.0: + print(f" ✅ SplitLinear faster by {(speedup-1)*100:.1f}%") + else: + print(f" ⚠️ SplitLinear slower by {(1-speedup)*100:.1f}%") + +if __name__ == '__main__': + main() diff --git a/experimental/libane_bridge_v6.m b/experimental/libane_bridge_v6.m new file mode 100644 index 0000000..2cfcdd6 --- /dev/null +++ b/experimental/libane_bridge_v6.m @@ -0,0 +1,675 @@ +// libane_bridge_v6.m — ANE bridge with manual retain (no ARC), IOSurface +// isolation +// Compile: clang -shared -o libane_bridge_v6.dylib libane_bridge_v6.m \ +// -framework Foundation -framework IOSurface -framework Accelerate -lobjc -O2 +// -fno-objc-arc +#import +#import +#import +#import +#import +#import +#import + +static Class g_D, g_I, g_AR, g_AIO; +static int g_init = 0; +static dispatch_queue_t g_ane_queue; +static dispatch_group_t g_ane_group; + +typedef struct { + id model; // retained manually + int ic, oc, seq; + int ioInIdx, ioOutIdx; +} ANEModel; + +#define MAX_MODELS 256 +#define MAX_SURFACES 32 +static ANEModel g_models[MAX_MODELS]; +static id g_requests[MAX_MODELS]; // retained manually +static int g_count = 0; + +typedef struct { + IOSurfaceRef surface; + id aneObj; // retained manually + size_t allocSize; +} IOSEntry; +static IOSEntry g_surfaces[MAX_SURFACES]; +static int g_surfCount = 0; + +static int find_or_create_surface(size_t needed, int exclude_idx) { + needed = (needed + 0xFFFF) & ~(size_t)0xFFFF; + if (needed < 65536) { + needed = 65536; + } + for (int i = 0; i < g_surfCount; i++) { + if (g_surfaces[i].allocSize == needed && i != exclude_idx) { + return i; + } + } + if (g_surfCount >= MAX_SURFACES) { + return -1; + } + IOSurfaceRef s = IOSurfaceCreate((__bridge CFDictionaryRef) @{ + (id)kIOSurfaceWidth : @(needed), + (id)kIOSurfaceHeight : @1, + (id)kIOSurfaceBytesPerElement : @1, + (id)kIOSurfaceBytesPerRow : @(needed), + (id)kIOSurfaceAllocSize : @(needed), + (id)kIOSurfacePixelFormat : @0 + }); + if (!s) { + return -1; + } + id obj = ((id (*)(Class, SEL, IOSurfaceRef))objc_msgSend)( + g_AIO, @selector(objectWithIOSurface:), s); + [obj retain]; + int idx = g_surfCount++; + g_surfaces[idx] = (IOSEntry){s, obj, needed}; + return idx; +} + +static NSData *build_blob(const float *w, int oc, int ic) { + size_t n = (size_t)oc * ic, wsize = n * 2, total = 128 + wsize; + uint8_t *buf = (uint8_t *)calloc(total, 1); + buf[0] = 0x01; + buf[4] = 0x02; + uint8_t *chunk = buf + 64; + chunk[0] = 0xEF; + chunk[1] = 0xBE; + chunk[2] = 0xAD; + chunk[3] = 0xDE; + chunk[4] = 0x01; + *(uint32_t *)(chunk + 8) = (uint32_t)wsize; + *(uint32_t *)(chunk + 16) = 128; + _Float16 *fp16 = (_Float16 *)(buf + 128); + for (size_t i = 0; i < n; i++) { + fp16[i] = (_Float16)w[i]; + } + return [NSData dataWithBytesNoCopy:buf length:total freeWhenDone:YES]; +} + +static NSString *gen_conv_mil(int ic, int oc, int seq) { + return [NSString + stringWithFormat: + @"program(1.3)\n[buildInfo = dict({{\"coremlc-component-MIL\", \"3510.2.1\"}, " + "{\"coremlc-version\", \"3505.4.1\"}, " + "{\"coremltools-component-milinternal\", \"\"}, " + "{\"coremltools-version\", \"9.0\"}})]\n{\n" + " func main(tensor x) {\n" + " string to_fp16 = const()[name = string(\"to_fp16\"), val = " + "string(\"fp16\")];\n" + " tensor x16 = cast(dtype = to_fp16, x " + "= x)[name = string(\"cast_in\")];\n" + " tensor W = const()[name = " + "string(\"W\"), val = tensor" + "(BLOBFILE(path = string(\"@model_path/weights/weight.bin\"), " + "offset = uint64(64)))];\n" + " string c_pad_type = const()[name = string(\"c_pad_type\"), " + "val = string(\"valid\")];\n" + " tensor c_strides = const()[name = " + "string(\"c_strides\"), val = tensor([1, 1])];\n" + " tensor c_pad = const()[name = " + "string(\"c_pad\"), val = tensor([0, 0, 0, 0])];\n" + " tensor c_dilations = const()[name = " + "string(\"c_dilations\"), val = tensor([1, 1])];\n" + " int32 c_groups = const()[name = string(\"c_groups\"), val " + "= int32(1)];\n" + " tensor y16 = conv(dilations = " + "c_dilations, groups = c_groups, pad = c_pad, " + "pad_type = c_pad_type, strides = c_strides, weight = W, x = " + "x16)[name = string(\"conv\")];\n" + " string to_fp32 = const()[name = string(\"to_fp32\"), val = " + "string(\"fp32\")];\n" + " tensor y = cast(dtype = to_fp32, x = " + "y16)[name = string(\"cast_out\")];\n" + " } -> (y);\n}\n", + ic, seq, ic, seq, oc, ic, oc, ic, oc, seq, oc, seq]; +} + +// INT8 MIL text: uses constexpr_blockwise_shift_scale for INT8 weight +// quantization +static NSString *gen_conv_mil_int8(int ic, int oc, int seq) { + return [NSString + stringWithFormat: + @"program(1.3)\n[buildInfo = dict({{\"coremlc-component-MIL\", \"3510.2.1\"}, " + "{\"coremlc-version\", \"3505.4.1\"}, " + "{\"coremltools-component-milinternal\", \"\"}, " + "{\"coremltools-version\", \"9.0\"}})]\ +{\n" + " func main(tensor x) {\n" + " string to_fp16 = const()[name = string(\"to_fp16\"), val = " + "string(\"fp16\")];\n" + " tensor x16 = cast(dtype = to_fp16, x " + "= x)[name = string(\"cast_in\")];\n" + " tensor W_data = const()[name = " + "string(\"W_data\"), " + "val = tensor(BLOBFILE(path = " + "string(\"@model_path/weights/weight_data.bin\"), offset = " + "uint64(64)))];\n" + " tensor W_scale = const()[name = " + "string(\"W_scale\"), " + "val = tensor(BLOBFILE(path = " + "string(\"@model_path/weights/weight_scale.bin\"), offset = " + "uint64(64)))];\n" + " tensor W_offset = const()[name = " + "string(\"W_offset\"), " + "val = tensor(BLOBFILE(path = " + "string(\"@model_path/weights/weight_offset.bin\"), offset = " + "uint64(64)))];\n" + " tensor W = " + "constexpr_blockwise_shift_scale(" + "data = W_data, scale = W_scale, offset = W_offset)[name = " + "string(\"dequant\")];\n" + " string c_pad_type = const()[name = string(\"c_pad_type\"), " + "val = string(\"valid\")];\n" + " tensor c_strides = const()[name = " + "string(\"c_strides\"), val = tensor([1, 1])];\n" + " tensor c_pad = const()[name = " + "string(\"c_pad\"), val = tensor([0, 0, 0, 0])];\n" + " tensor c_dilations = const()[name = " + "string(\"c_dilations\"), val = tensor([1, 1])];\n" + " int32 c_groups = const()[name = string(\"c_groups\"), val " + "= int32(1)];\n" + " tensor y16 = conv(dilations = " + "c_dilations, groups = c_groups, pad = c_pad, " + "pad_type = c_pad_type, strides = c_strides, weight = W, x = " + "x16)[name = string(\"conv\")];\n" + " string to_fp32 = const()[name = string(\"to_fp32\"), val = " + "string(\"fp32\")];\n" + " tensor y = cast(dtype = to_fp32, x = " + "y16)[name = string(\"cast_out\")];\n" + " } -> (y);\n}\n", + ic, seq, ic, seq, oc, ic, oc, ic, // W_data shape + oc, oc, // W_scale shape + oc, oc, // W_offset shape + oc, ic, // W dequantized shape + oc, seq, // y16 shape + oc, seq]; // y shape +} + +// Build INT8 weight blob: header(64) + magic(64) + data +static NSData *build_blob_int8(const int8_t *data, int count) { + size_t wsize = (size_t)count, total = 128 + wsize; + uint8_t *buf = (uint8_t *)calloc(total, 1); + buf[0] = 0x01; + buf[4] = 0x02; + uint8_t *chunk = buf + 64; + chunk[0] = 0xEF; + chunk[1] = 0xBE; + chunk[2] = 0xAD; + chunk[3] = 0xDE; + chunk[4] = 0x01; + *(uint32_t *)(chunk + 8) = (uint32_t)wsize; + *(uint32_t *)(chunk + 16) = 128; + memcpy(buf + 128, data, wsize); + return [NSData dataWithBytesNoCopy:buf length:total freeWhenDone:YES]; +} + +static NSData *build_blob_fp16(const _Float16 *data, int count) { + size_t wsize = (size_t)count * 2, total = 128 + wsize; + uint8_t *buf = (uint8_t *)calloc(total, 1); + buf[0] = 0x01; + buf[4] = 0x02; + uint8_t *chunk = buf + 64; + chunk[0] = 0xEF; + chunk[1] = 0xBE; + chunk[2] = 0xAD; + chunk[3] = 0xDE; + chunk[4] = 0x01; + *(uint32_t *)(chunk + 8) = (uint32_t)wsize; + *(uint32_t *)(chunk + 16) = 128; + memcpy(buf + 128, data, wsize); + return [NSData dataWithBytesNoCopy:buf length:total freeWhenDone:YES]; +} + +static int setup_model(id mdl, int ic, int oc, int seq) { + size_t inB = (size_t)ic * seq * 4; + size_t outB = (size_t)oc * seq * 4; + int ioInIdx = find_or_create_surface(inB, -1); + int ioOutIdx = find_or_create_surface(outB, ioInIdx); // exclude input surface + if (ioInIdx < 0 || ioOutIdx < 0) { + return -1; + } + id wI = g_surfaces[ioInIdx].aneObj; + id wO = g_surfaces[ioOutIdx].aneObj; + id req = ((id (*)(Class, SEL, id, id, id, id, id, id, id))objc_msgSend)( + g_AR, + @selector(requestWithInputs:inputIndices:outputs:outputIndices: + weightsBuffer:perfStats:procedureIndex:), + @[ wI ], @[ @0 ], @[ wO ], @[ @0 ], nil, nil, @0); + [req retain]; + [mdl retain]; + + int handle = g_count; + g_models[handle] = (ANEModel){mdl, ic, oc, seq, ioInIdx, ioOutIdx}; + g_requests[handle] = req; + g_count++; + return handle; +} + +// ===== Public C API ===== + +int ane_init(void) { + if (g_init) { + return 0; + } + dlopen("/System/Library/PrivateFrameworks/AppleNeuralEngine.framework/" + "AppleNeuralEngine", + RTLD_NOW); + g_D = NSClassFromString(@"_ANEInMemoryModelDescriptor"); + g_I = NSClassFromString(@"_ANEInMemoryModel"); + g_AR = NSClassFromString(@"_ANERequest"); + g_AIO = NSClassFromString(@"_ANEIOSurfaceObject"); + if (!g_D || !g_I || !g_AR || !g_AIO) + return -1; + g_ane_queue = dispatch_queue_create("ane.eval", DISPATCH_QUEUE_SERIAL); + g_ane_group = dispatch_group_create(); + g_init = 1; + return 0; +} + +int ane_load_model(int ic, int oc, int seq, const float *w) { + if (!g_init || g_count >= MAX_MODELS) { + return -1; + } + NSAutoreleasePool *pool = [[NSAutoreleasePool alloc] init]; + + NSData *mil = + [gen_conv_mil(ic, oc, seq) dataUsingEncoding:NSUTF8StringEncoding]; + NSData *blob = build_blob(w, oc, ic); + NSDictionary *wd = + @{@"@model_path/weights/weight.bin" : @{@"offset" : @0, @"data" : blob}}; + NSError *e = nil; + + id desc = ((id (*)(Class, SEL, id, id, id))objc_msgSend)( + g_D, @selector(modelWithMILText:weights:optionsPlist:), mil, wd, nil); + if (!desc) { + [pool drain]; + return -1; + } + id mdl = ((id (*)(Class, SEL, id))objc_msgSend)( + g_I, @selector(inMemoryModelWithDescriptor:), desc); + if (!mdl) { + [pool drain]; + return -1; + } + + id hx = ((id (*)(id, SEL))objc_msgSend)(mdl, @selector(hexStringIdentifier)); + NSString *td = [NSTemporaryDirectory() stringByAppendingPathComponent:hx]; + [[NSFileManager defaultManager] + createDirectoryAtPath:[td stringByAppendingPathComponent:@"weights"] + withIntermediateDirectories:YES + attributes:nil + error:nil]; + [mil writeToFile:[td stringByAppendingPathComponent:@"model.mil"] + atomically:YES]; + [blob writeToFile:[td stringByAppendingPathComponent:@"weights/weight.bin"] + atomically:YES]; + + BOOL ok; + ok = ((BOOL (*)(id, SEL, unsigned int, id, NSError **))objc_msgSend)( + mdl, @selector(compileWithQoS:options:error:), 21, @{}, &e); + if (!ok) { + NSLog(@"compile failed %d->%d: %@", ic, oc, e); + [pool drain]; + return -1; + } + ok = ((BOOL (*)(id, SEL, unsigned int, id, NSError **))objc_msgSend)( + mdl, @selector(loadWithQoS:options:error:), 21, @{}, &e); + if (!ok) { + NSLog(@"load failed %d->%d: %@", ic, oc, e); + [pool drain]; + return -1; + } + + int result = setup_model(mdl, ic, oc, seq); + [pool drain]; + return result; +} + +// Load INT8 quantized conv model onto ANE +// w_int8: [oc, ic] int8 row-major, scale: [oc] fp32, offset: [oc] fp32 +int ane_load_model_int8(int ic, int oc, int seq, const int8_t *w_int8, + const float *scale_fp32, const float *offset_fp32) { + if (!g_init || g_count >= MAX_MODELS) { + return -1; + } + NSAutoreleasePool *pool = [[NSAutoreleasePool alloc] init]; + + NSData *mil = + [gen_conv_mil_int8(ic, oc, seq) dataUsingEncoding:NSUTF8StringEncoding]; + + // Build weight blobs: INT8 data, FP16 scale, FP16 offset + NSData *blob_data = build_blob_int8(w_int8, oc * ic); + + _Float16 *scale_fp16 = (_Float16 *)calloc(oc, sizeof(_Float16)); + _Float16 *offset_fp16 = (_Float16 *)calloc(oc, sizeof(_Float16)); + for (int i = 0; i < oc; i++) { + scale_fp16[i] = (_Float16)scale_fp32[i]; + offset_fp16[i] = (_Float16)offset_fp32[i]; + } + NSData *blob_scale = build_blob_fp16(scale_fp16, oc); + NSData *blob_offset = build_blob_fp16(offset_fp16, oc); + free(scale_fp16); + free(offset_fp16); + + NSDictionary *wd = @{ + @"@model_path/weights/weight_data.bin" : + @{@"offset" : @0, @"data" : blob_data}, + @"@model_path/weights/weight_scale.bin" : + @{@"offset" : @0, @"data" : blob_scale}, + @"@model_path/weights/weight_offset.bin" : + @{@"offset" : @0, @"data" : blob_offset} + }; + NSError *e = nil; + + id desc = ((id (*)(Class, SEL, id, id, id))objc_msgSend)( + g_D, @selector(modelWithMILText:weights:optionsPlist:), mil, wd, nil); + if (!desc) { + NSLog(@"INT8 desc failed %d->%d", ic, oc); + [pool drain]; + return -1; + } + id mdl = ((id (*)(Class, SEL, id))objc_msgSend)( + g_I, @selector(inMemoryModelWithDescriptor:), desc); + if (!mdl) { + NSLog(@"INT8 model failed %d->%d", ic, oc); + [pool drain]; + return -1; + } + + id hx = ((id (*)(id, SEL))objc_msgSend)(mdl, @selector(hexStringIdentifier)); + NSString *td = [NSTemporaryDirectory() stringByAppendingPathComponent:hx]; + [[NSFileManager defaultManager] + createDirectoryAtPath:[td stringByAppendingPathComponent:@"weights"] + withIntermediateDirectories:YES + attributes:nil + error:nil]; + [mil writeToFile:[td stringByAppendingPathComponent:@"model.mil"] + atomically:YES]; + [blob_data + writeToFile:[td stringByAppendingPathComponent:@"weights/weight_data.bin"] + atomically:YES]; + [blob_scale writeToFile:[td stringByAppendingPathComponent: + @"weights/weight_scale.bin"] + atomically:YES]; + [blob_offset writeToFile:[td stringByAppendingPathComponent: + @"weights/weight_offset.bin"] + atomically:YES]; + + BOOL ok; + ok = ((BOOL (*)(id, SEL, unsigned int, id, NSError **))objc_msgSend)( + mdl, @selector(compileWithQoS:options:error:), 21, @{}, &e); + if (!ok) { + NSLog(@"INT8 compile failed %d->%d: %@", ic, oc, e); + [pool drain]; + return -1; + } + ok = ((BOOL (*)(id, SEL, unsigned int, id, NSError **))objc_msgSend)( + mdl, @selector(loadWithQoS:options:error:), 21, @{}, &e); + if (!ok) { + NSLog(@"INT8 load failed %d->%d: %@", ic, oc, e); + [pool drain]; + return -1; + } + + int result = setup_model(mdl, ic, oc, seq); + [pool drain]; + return result; +} + +int ane_eval(int handle) { + if (handle < 0 || handle >= g_count) { + return -1; + } + NSError *e = nil; + return ((BOOL (*)(id, SEL, unsigned int, id, id, NSError **))objc_msgSend)( + g_models[handle].model, + @selector(evaluateWithQoS:options:request:error:), 21, @{}, + g_requests[handle], &e) + ? 0 + : -2; +} + +int ane_run(int handle, const float *input, float *output) { + if (handle < 0 || handle >= g_count) { + return -1; + } + ANEModel *m = &g_models[handle]; + IOSurfaceRef ioI = g_surfaces[m->ioInIdx].surface; + IOSurfaceRef ioO = g_surfaces[m->ioOutIdx].surface; + IOSurfaceLock(ioI, 0, NULL); + memcpy(IOSurfaceGetBaseAddress(ioI), input, + (size_t)m->ic * m->seq * sizeof(float)); + IOSurfaceUnlock(ioI, 0, NULL); + int rc = ane_eval(handle); + if (rc != 0) { + return rc; + } + IOSurfaceLock(ioO, kIOSurfaceLockReadOnly, NULL); + memcpy(output, IOSurfaceGetBaseAddress(ioO), + (size_t)m->oc * m->seq * sizeof(float)); + IOSurfaceUnlock(ioO, kIOSurfaceLockReadOnly, NULL); + return 0; +} + +int ane_run_async(int handle, const float *input) { + if (handle < 0 || handle >= g_count) { + return -1; + } + ANEModel *m = &g_models[handle]; + IOSurfaceRef ioI = g_surfaces[m->ioInIdx].surface; + IOSurfaceLock(ioI, 0, NULL); + memcpy(IOSurfaceGetBaseAddress(ioI), input, + (size_t)m->ic * m->seq * sizeof(float)); + IOSurfaceUnlock(ioI, 0, NULL); + dispatch_group_async(g_ane_group, g_ane_queue, ^{ + NSError *e = nil; + ((BOOL (*)(id, SEL, unsigned int, id, id, NSError **))objc_msgSend)( + g_models[handle].model, + @selector(evaluateWithQoS:options:request:error:), 21, @{}, + g_requests[handle], &e); + }); + return 0; +} + +int ane_wait_read(int handle, float *output) { + if (handle < 0 || handle >= g_count) { + return -1; + } + dispatch_group_wait(g_ane_group, DISPATCH_TIME_FOREVER); + ANEModel *m = &g_models[handle]; + IOSurfaceRef ioO = g_surfaces[m->ioOutIdx].surface; + IOSurfaceLock(ioO, kIOSurfaceLockReadOnly, NULL); + memcpy(output, IOSurfaceGetBaseAddress(ioO), + (size_t)m->oc * m->seq * sizeof(float)); + IOSurfaceUnlock(ioO, kIOSurfaceLockReadOnly, NULL); + return 0; +} + +void ane_wait(void) { dispatch_group_wait(g_ane_group, DISPATCH_TIME_FOREVER); } +int ane_model_count(void) { return g_count; } +int ane_surface_count(void) { return g_surfCount; } + +void *ane_get_input_ptr(int handle) { + if (handle < 0 || handle >= g_count) { + return NULL; + } + return IOSurfaceGetBaseAddress(g_surfaces[g_models[handle].ioInIdx].surface); +} +void *ane_get_output_ptr(int handle) { + if (handle < 0 || handle >= g_count) + return NULL; + return IOSurfaceGetBaseAddress(g_surfaces[g_models[handle].ioOutIdx].surface); +} +int ane_get_ic(int h) { return (h >= 0 && h < g_count) ? g_models[h].ic : -1; } +int ane_get_oc(int h) { return (h >= 0 && h < g_count) ? g_models[h].oc : -1; } +int ane_get_seq(int h) { + return (h >= 0 && h < g_count) ? g_models[h].seq : -1; +} + +// Write input directly into IOSurface (caller provides [IC, SEQ] C-contiguous +// data) +void ane_write_input(int handle, const float *data) { + if (handle < 0 || handle >= g_count) { + return; + } + ANEModel *m = &g_models[handle]; + IOSurfaceRef ioI = g_surfaces[m->ioInIdx].surface; + IOSurfaceLock(ioI, 0, NULL); + memcpy(IOSurfaceGetBaseAddress(ioI), data, + (size_t)m->ic * m->seq * sizeof(float)); + IOSurfaceUnlock(ioI, 0, NULL); +} + +// Read output directly from IOSurface into [OC, SEQ] C-contiguous buffer +void ane_read_output(int handle, float *data) { + if (handle < 0 || handle >= g_count) { + return; + } + ANEModel *m = &g_models[handle]; + IOSurfaceRef ioO = g_surfaces[m->ioOutIdx].surface; + IOSurfaceLock(ioO, kIOSurfaceLockReadOnly, NULL); + memcpy(data, IOSurfaceGetBaseAddress(ioO), + (size_t)m->oc * m->seq * sizeof(float)); + IOSurfaceUnlock(ioO, kIOSurfaceLockReadOnly, NULL); +} + +// Tile transpose: src[rows][cols] → dst[cols][rows], cache-friendly 16×16 +// blocks +static void tile_transpose(const float *__restrict__ src, + float *__restrict__ dst, int rows, int cols) { + const int T = 16; + for (int r = 0; r < rows; r += T) { + int rend = r + T < rows ? r + T : rows; + for (int c = 0; c < cols; c += T) { + int cend = c + T < cols ? c + T : cols; + for (int ri = r; ri < rend; ri++) + for (int ci = c; ci < cend; ci++) + dst[ci * rows + ri] = src[ri * cols + ci]; + } + } +} + +// ane_run_T: input [SEQ, IC] row-major, output [SEQ, OC] row-major +// Tile-transpose input to [IC, SEQ] for IOSurface, eval, transpose output back +int ane_run_T(int handle, const float *input_seq_ic, float *output_seq_oc) { + if (handle < 0 || handle >= g_count) { + return -1; + } + ANEModel *m = &g_models[handle]; + IOSurfaceRef ioI = g_surfaces[m->ioInIdx].surface; + IOSurfaceRef ioO = g_surfaces[m->ioOutIdx].surface; + IOSurfaceLock(ioI, 0, NULL); + tile_transpose(input_seq_ic, (float *)IOSurfaceGetBaseAddress(ioI), m->seq, + m->ic); + IOSurfaceUnlock(ioI, 0, NULL); + NSError *e = nil; + BOOL ok = ((BOOL (*)(id, SEL, unsigned int, id, id, NSError **))objc_msgSend)( + m->model, @selector(evaluateWithQoS:options:request:error:), 21, @{}, + g_requests[handle], &e); + if (!ok) { + return -2; + } + IOSurfaceLock(ioO, kIOSurfaceLockReadOnly, NULL); + tile_transpose((const float *)IOSurfaceGetBaseAddress(ioO), output_seq_oc, + m->oc, m->seq); + IOSurfaceUnlock(ioO, kIOSurfaceLockReadOnly, NULL); + return 0; +} + +// ane_eval_direct: write input, eval, read output — single call, minimal +// overhead +int ane_eval_direct(int handle, const float *input, float *output) { + if (handle < 0 || handle >= g_count) { + return -1; + } + ANEModel *m = &g_models[handle]; + IOSurfaceRef ioI = g_surfaces[m->ioInIdx].surface; + IOSurfaceRef ioO = g_surfaces[m->ioOutIdx].surface; + // Write input + IOSurfaceLock(ioI, 0, NULL); + memcpy(IOSurfaceGetBaseAddress(ioI), input, + (size_t)m->ic * m->seq * sizeof(float)); + IOSurfaceUnlock(ioI, 0, NULL); + // Eval + NSError *e = nil; + BOOL ok = ((BOOL (*)(id, SEL, unsigned int, id, id, NSError **))objc_msgSend)( + m->model, @selector(evaluateWithQoS:options:request:error:), 21, @{}, + g_requests[handle], &e); + if (!ok) { + return -2; + } + // Read output + IOSurfaceLock(ioO, kIOSurfaceLockReadOnly, NULL); + memcpy(output, IOSurfaceGetBaseAddress(ioO), + (size_t)m->oc * m->seq * sizeof(float)); + IOSurfaceUnlock(ioO, kIOSurfaceLockReadOnly, NULL); + return 0; +} + +// ane_run_rowmajor: input [L, IC] row-major, output [L, OC] row-major +// Uses vDSP_mtrans for fast transpose (Apple Accelerate) +// When L==seq: pure vDSP path (fastest). L= g_count) { + return -1; + } + ANEModel *m = &g_models[handle]; + if (L > m->seq) { + return -3; + } + IOSurfaceRef ioI = g_surfaces[m->ioInIdx].surface; + IOSurfaceRef ioO = g_surfaces[m->ioOutIdx].surface; + + // === Transpose input [L, IC] -> IOSurface [IC, seq] === + IOSurfaceLock(ioI, 0, NULL); + float *ioBuf = (float *)IOSurfaceGetBaseAddress(ioI); + if (L == m->seq) { + // Fast path: vDSP_mtrans [L][IC] -> [IC][L=seq] + vDSP_mtrans(input_rm, 1, ioBuf, 1, m->ic, L); + } else { + // L < seq: transpose to temp, then scatter with zero-pad + // For each channel ch: ioBuf[ch*seq+0..L-1] = transposed, [L..seq-1] = 0 + // Use vDSP_mtrans to a contiguous temp, then scatter + float *tmp = (float *)malloc((size_t)m->ic * L * sizeof(float)); + vDSP_mtrans(input_rm, 1, tmp, 1, m->ic, L); + for (int ch = 0; ch < m->ic; ch++) { + memcpy(&ioBuf[ch * m->seq], &tmp[ch * L], L * sizeof(float)); + if (L < m->seq) { + memset(&ioBuf[ch * m->seq + L], 0, (m->seq - L) * sizeof(float)); + } + } + free(tmp); + } + IOSurfaceUnlock(ioI, 0, NULL); + + // === ANE eval === + NSError *e = nil; + BOOL ok = ((BOOL (*)(id, SEL, unsigned int, id, id, NSError **))objc_msgSend)( + m->model, @selector(evaluateWithQoS:options:request:error:), 21, @{}, + g_requests[handle], &e); + if (!ok) { + return -2; + } + // === Transpose output IOSurface [OC, seq] -> [L, OC] === + IOSurfaceLock(ioO, kIOSurfaceLockReadOnly, NULL); + const float *outBuf = (const float *)IOSurfaceGetBaseAddress(ioO); + if (L == m->seq) { + // Fast path: vDSP_mtrans [OC][seq] -> [seq][OC] + vDSP_mtrans(outBuf, 1, output_rm, 1, L, m->oc); + } else { + // Gather first L elements from each of OC rows, then transpose + float *tmp = (float *)malloc((size_t)m->oc * L * sizeof(float)); + for (int ch = 0; ch < m->oc; ch++) { + memcpy(&tmp[ch * L], &outBuf[ch * m->seq], L * sizeof(float)); + } + vDSP_mtrans(tmp, 1, output_rm, 1, L, m->oc); + free(tmp); + } + IOSurfaceUnlock(ioO, kIOSurfaceLockReadOnly, NULL); + return 0; +} diff --git a/experimental/split_linear.py b/experimental/split_linear.py new file mode 100644 index 0000000..b8a4e55 --- /dev/null +++ b/experimental/split_linear.py @@ -0,0 +1,467 @@ +#!/usr/bin/env python3 +""" +SplitLinear: Drop-in replacement for nn.Linear / nn.QuantizedLinear. +Splits GEMM along output channels — ANE does ~65%, GPU does ~35%, concurrent. + +Key optimization: same-input projections (Q/K/V, Gate/Up) share a single +input preparation via _InputGroup, eliminating redundant transpose+eval+numpy. + +Usage: + from split_linear import patch_model + bridge = patch_model(model, seq=512) + +API: + patch_model(model, seq) → high-level one-liner + SplitLinear(layer, bridge, seq) → single layer replacement + ANEBridge() → ANE private API wrapper +""" +import os, ctypes +from concurrent.futures import ThreadPoolExecutor +import numpy as np +import mlx.core as mx +import mlx.nn as nn + +LIB_DIR = os.path.dirname(os.path.abspath(__file__)) +SPLIT_ALIGN = 64 +MIN_SEQ_FOR_SPLIT = 192 # Below this, split overhead > benefit + + +# ─── ANE Bridge ─── + +class ANEBridge: + """Thin wrapper around ANE private API.""" + _instance = None + + def __init__(self): + lib = ctypes.CDLL(os.path.join(LIB_DIR, 'libane_bridge_v6.dylib')) + lib.ane_init.restype = ctypes.c_int + lib.ane_load_model.restype = ctypes.c_int + lib.ane_load_model.argtypes = [ctypes.c_int, ctypes.c_int, ctypes.c_int, + ctypes.POINTER(ctypes.c_float)] + lib.ane_run.restype = ctypes.c_int + lib.ane_run.argtypes = [ctypes.c_int, ctypes.POINTER(ctypes.c_float), + ctypes.POINTER(ctypes.c_float)] + lib.ane_run_rowmajor.restype = ctypes.c_int + lib.ane_run_rowmajor.argtypes = [ctypes.c_int, ctypes.POINTER(ctypes.c_float), + ctypes.c_int, ctypes.POINTER(ctypes.c_float)] + lib.ane_model_count.restype = ctypes.c_int + lib.ane_surface_count.restype = ctypes.c_int + assert lib.ane_init() == 0, "ANE init failed" + self.lib = lib + + def load(self, ic, oc, seq, w_fp32): + w = np.ascontiguousarray(w_fp32, dtype=np.float32) + h = self.lib.ane_load_model(ic, oc, seq, + w.ctypes.data_as(ctypes.POINTER(ctypes.c_float))) + assert h >= 0, f"ANE load failed: {ic}→{oc}, seq={seq}" + return h + + def run(self, h, inp, out): + self.lib.ane_run(h, + inp.ctypes.data_as(ctypes.POINTER(ctypes.c_float)), + out.ctypes.data_as(ctypes.POINTER(ctypes.c_float))) + + def run_rowmajor(self, h, inp_rm, L, out_rm): + """ANE compute with row-major I/O. Uses vDSP for transpose. + inp_rm: [L, IC] row-major float32 + out_rm: [L, OC] row-major float32 (pre-allocated) + """ + self.lib.ane_run_rowmajor(h, + inp_rm.ctypes.data_as(ctypes.POINTER(ctypes.c_float)), + L, + out_rm.ctypes.data_as(ctypes.POINTER(ctypes.c_float))) + + @property + def model_count(self): + return self.lib.ane_model_count() + + @classmethod + def shared(cls): + if cls._instance is None: + cls._instance = cls() + return cls._instance + + +# Single ANE worker thread +_ane_pool = ThreadPoolExecutor(max_workers=1, thread_name_prefix='ane') + + +# ─── Helpers ─── + +def _extract_weight(layer): + """Extract FP32 weight [OC, IC] from nn.Linear or nn.QuantizedLinear.""" + if isinstance(layer, nn.QuantizedLinear): + mx.eval(layer.weight, layer.scales) + b = getattr(layer, 'biases', None) + if b is not None: + mx.eval(b) + w = mx.dequantize(layer.weight, layer.scales, b, + layer.group_size, layer.bits) + else: + w = layer.weight + mx.eval(w) + return np.array(w, dtype=np.float32) + + +class _InputGroup: + """ + Shared input preparation for same-input projections (Q/K/V or Gate/Up). + + Zero-copy path: eval x as FP32 row-major, share numpy view across group. + No GPU transpose needed — ane_run_rowmajor handles transpose in C via vDSP. + """ + __slots__ = ('_np', '_x2d', '_L', '_seq') + + def __init__(self, ic=0, seq=0): + self._np = None + self._x2d = None + self._L = 0 + self._seq = seq + + def prepare(self, x): + """Prepare input. Called by first proj in group. + Returns (inp_np_rowmajor [L, IC], x_2d, L).""" + x_2d = x.reshape(-1, x.shape[-1]) if x.ndim == 3 else x + L = x_2d.shape[0] + + # Cast to FP32 contiguous row-major, eval, then zero-copy numpy view + x_f32 = x_2d.astype(mx.float32) if x_2d.dtype != mx.float32 else x_2d + x_f32 = mx.contiguous(x_f32) + mx.eval(x_f32) + self._np = np.array(x_f32, copy=False) # [L, IC] row-major, zero-copy + + self._L = L + self._x2d = x_2d + return self._np, x_2d, L + + def get(self): + """Get cached input. Called by subsequent projs. + Returns (inp_np_rowmajor [L, IC], x_2d, L).""" + return self._np, self._x2d, self._L + + +# ─── SplitLinear ─── + +class SplitLinear: + """ + Drop-in nn.Linear replacement with ANE+GPU tensor parallelism. + + Two modes controlled externally via set_prefill(True/False): + - prefill: ANE ~65% + GPU ~35% concurrent + - decode: original nn.Linear on GPU → zero overhead + """ + _prefill_mode = False # class-level flag + + def __init__(self, layer, bridge, seq, ane_frac=None, name="", + input_group=None, is_first=True): + self.name = name + self._orig = layer # always keep original for decode + w_np = _extract_weight(layer) + oc, ic = w_np.shape + self.ic = ic + self.oc = oc + self.seq = seq + self._is_first = is_first + + # Auto-detect split fraction + if ane_frac is None: + if seq < MIN_SEQ_FOR_SPLIT: + ane_frac = 0.0 + elif ic > oc * 2: + ane_frac = 0.0 # Wide→narrow: ANE inefficient (down_proj) + elif oc < SPLIT_ALIGN * 2: + ane_frac = 0.0 + else: + ane_frac = 0.65 + + self.ane_frac = ane_frac + + if ane_frac <= 0: + self.mode = 'gpu' + return + + self.mode = 'split' + self.ane = bridge + + self.ane_oc = (int(oc * ane_frac) // SPLIT_ALIGN) * SPLIT_ALIGN + self.gpu_oc = oc - self.ane_oc + + if self.ane_oc < SPLIT_ALIGN or self.gpu_oc < 1: + self.mode = 'gpu' + return + + self.h_ane = bridge.load(ic, self.ane_oc, seq, w_np[:self.ane_oc, :]) + self.buf_ane = np.empty((seq, self.ane_oc), dtype=np.float32) # ANE FP32 output + self.buf_ane_f16 = np.empty((seq, self.ane_oc), dtype=np.float16) # FP16 cast buffer + + # GPU weight for split path + self._is_quantized = isinstance(layer, nn.QuantizedLinear) + if self._is_quantized: + # Re-quantize GPU portion as QuantizedLinear → native quantized_matmul + w_gpu_fp32 = mx.array(w_np[self.ane_oc:, :]) # [gpu_oc, ic] float32 + w_q, scales, biases = mx.quantize(w_gpu_fp32, + group_size=layer.group_size, + bits=layer.bits) + self._gpu_layer = nn.QuantizedLinear( + ic, self.gpu_oc, bias=False, + group_size=layer.group_size, bits=layer.bits) + self._gpu_layer.weight = w_q + self._gpu_layer.scales = scales + self._gpu_layer.biases = biases + mx.eval(self._gpu_layer.parameters()) + self._w_gpu = None # not used for quantized path + else: + self._gpu_layer = None + self._w_gpu = None + + if input_group is not None: + self._grp = input_group + else: + self._grp = _InputGroup(ic, seq) + self._is_first = True + + @classmethod + def set_prefill(cls, enabled): + cls._prefill_mode = enabled + + def __call__(self, x): + # decode / gpu-only / short seq: use original nn.Linear + if self.mode == 'gpu' or not SplitLinear._prefill_mode: + return self._orig(x) + + L = x.shape[-2] if x.ndim == 3 else x.shape[0] + if L < MIN_SEQ_FOR_SPLIT: + return self._orig(x) + + # ── prefill split path (zero-copy + concurrent ANE/GPU) ── + orig_shape = x.shape + + # Input preparation (shared or fresh) — row-major [L, IC] FP32 + if self._is_first: + inp_np, x_2d, L = self._grp.prepare(x) + else: + inp_np, x_2d, L = self._grp.get() + + # ANE: launch in worker thread (row-major path, vDSP transpose in C) + out_buf = self.buf_ane[:L] # [L, ane_oc] view + fut = _ane_pool.submit(self.ane.run_rowmajor, self.h_ane, inp_np, L, out_buf) + + # GPU: matmul with GPU portion weight + if self._gpu_layer is not None: + # Quantized: use native quantized_matmul via QuantizedLinear + gpu_out = self._gpu_layer(x_2d) + elif self._w_gpu is not None: + gpu_out = x_2d @ self._w_gpu.T + else: + w_gpu = self._orig.weight[self.ane_oc:, :] + gpu_out = x_2d @ w_gpu.T + mx.eval(gpu_out) # sync GPU — enables concurrent ANE execution + + # Wait for ANE + fut.result() + + # Merge: numpy FP32→FP16 cast + lazy concat + np.copyto(self.buf_ane_f16[:L], out_buf, casting='same_kind') + ane_out = mx.array(self.buf_ane_f16[:L]) # [L, ane_oc] FP16 + merged = mx.concatenate([ane_out, gpu_out], axis=-1) + + if len(orig_shape) == 3: + merged = merged.reshape(orig_shape[0], orig_shape[1], -1) + return merged + + def __repr__(self): + if self.mode == 'gpu': + return f"SplitLinear({self.name}, gpu_only, {self.ic}→{self.oc})" + return (f"SplitLinear({self.name}, {self.ane_oc}ane+{self.gpu_oc}gpu, " + f"{self.ic}→{self.oc}, frac={self.ane_frac:.0%})") + + +# ─── Tree Walk ─── + +def _find_linears(module, prefix=""): + """Walk MLX model tree, yield (parent, key, full_name, linear).""" + if not hasattr(module, 'children'): + return + for attr_name, child in module.children().items(): + full_name = f"{prefix}.{attr_name}" if prefix else attr_name + if isinstance(child, (nn.Linear, nn.QuantizedLinear)): + yield (module, attr_name, full_name, child) + elif isinstance(child, nn.Module): + yield from _find_linears(child, full_name) + elif isinstance(child, list): + for i, v in enumerate(child): + fname = f"{full_name}.{i}" + if isinstance(v, (nn.Linear, nn.QuantizedLinear)): + yield (child, i, fname, v) + elif hasattr(v, 'children'): + yield from _find_linears(v, fname) + elif isinstance(child, dict): + for k, v in child.items(): + fname = f"{full_name}.{k}" + if isinstance(v, (nn.Linear, nn.QuantizedLinear)): + yield (module, k, fname, v) + elif hasattr(v, 'children'): + yield from _find_linears(v, fname) + elif hasattr(child, 'children'): + yield from _find_linears(child, full_name) + + +def _get_input_dims(layer): + """Get true input dimensions for nn.Linear or nn.QuantizedLinear.""" + if isinstance(layer, nn.QuantizedLinear): + return layer.weight.shape[-1] * 32 // layer.bits + return layer.weight.shape[-1] + + +def patch_model(model, seq, bridge=None, verbose=True): + """ + Patch all linear layers with SplitLinear. + + Same-input projections share an InputGroup: + - Q/K/V share one (same hidden_state input) + - Gate/Up share one (same hidden_state after attn) + - O and Down each get their own (unique inputs) + + Returns: bridge instance. + """ + if bridge is None: + bridge = ANEBridge.shared() + + lang = model.language_model + lm = lang.model + N = len(lm.layers) + n_split = n_gpu = 0 + + for li in range(N): + la = lm.layers[li] + attn = la.self_attn + mlp = la.mlp + + ic_attn = _get_input_dims(attn.q_proj) + ic_mlp = _get_input_dims(mlp.gate_proj) + + # Q/K/V share input group + qkv_grp = _InputGroup(ic_attn, seq) + for i, name in enumerate(('q_proj', 'k_proj', 'v_proj')): + orig = getattr(attn, name) + sl = SplitLinear(orig, bridge, seq, + name=f"layer.{li}.{name}", + input_group=qkv_grp, + is_first=(i == 0)) + setattr(attn, name, sl) + n_split += 1 if sl.mode == 'split' else 0 + n_gpu += 1 if sl.mode == 'gpu' else 0 + + # O proj — own input + sl = SplitLinear(attn.o_proj, bridge, seq, + name=f"layer.{li}.o_proj") + attn.o_proj = sl + n_split += 1 if sl.mode == 'split' else 0 + n_gpu += 1 if sl.mode == 'gpu' else 0 + + # Gate/Up share input group + gu_grp = _InputGroup(ic_mlp, seq) + for i, name in enumerate(('gate_proj', 'up_proj')): + orig = getattr(mlp, name) + sl = SplitLinear(orig, bridge, seq, + name=f"layer.{li}.{name}", + input_group=gu_grp, + is_first=(i == 0)) + setattr(mlp, name, sl) + n_split += 1 if sl.mode == 'split' else 0 + n_gpu += 1 if sl.mode == 'gpu' else 0 + + # Down proj — own input (auto GPU-only due to IC>OC*2) + sl = SplitLinear(mlp.down_proj, bridge, seq, + name=f"layer.{li}.down_proj") + mlp.down_proj = sl + n_split += 1 if sl.mode == 'split' else 0 + n_gpu += 1 if sl.mode == 'gpu' else 0 + + if verbose: + print(f"[SplitLinear] {N} layers: {n_split} split, {n_gpu} gpu-only") + print(f"[SplitLinear] ANE models: {bridge.model_count}") + + return bridge + + +# ─── Self-test ─── +if __name__ == '__main__': + import sys, time + from mlx_vlm.utils import load as vlm_load + from mlx_vlm.models.cache import KVCache + from mlx.utils import tree_flatten + + SEQ = int(sys.argv[1]) if len(sys.argv) > 1 else 512 + N_W = 3; N_B = 10 + FP16_MODEL = '~/Downloads/weights/mlx/Qwen3-VL-2B-Instruct-16bit' + + print(f"\n{'='*60}") + print(f" SplitLinear Self-Test (seq={SEQ})") + print(f"{'='*60}\n") + + model, _ = vlm_load(FP16_MODEL) + # Cast to true FP16 + flat = tree_flatten(model.trainable_parameters()) + fp16 = [(k, v.astype(mx.float16)) for k, v in flat] + model.load_weights(fp16) + mx.eval(model.parameters()) + + lang = model.language_model + N = lang.args.num_hidden_layers + ids = mx.ones((1, SEQ), dtype=mx.int32) + pos = mx.broadcast_to(mx.arange(SEQ).reshape(1, SEQ)[None, :, :], (3, 1, SEQ)) + mx.eval(ids, pos) + + # GPU baseline + print("[1/3] GPU FP16 baseline") + for _ in range(N_W): + c = [KVCache() for _ in range(N)] + mx.eval(lang(ids, cache=c, position_ids=pos).logits) + ts = [] + for _ in range(N_B): + c = [KVCache() for _ in range(N)] + t0 = time.perf_counter() + mx.eval(lang(ids, cache=c, position_ids=pos).logits) + ts.append((time.perf_counter()-t0)*1000) + bl = float(np.median(ts)) + print(f" {bl:.1f}ms\n") + + # Reference logits (FP32 for cos_sim) + c_ref = [KVCache() for _ in range(N)] + ref = np.array(lang(ids, cache=c_ref, position_ids=pos).logits.astype(mx.float32)) + + # Patch + print("[2/3] Patch + benchmark") + bridge = patch_model(model, SEQ) + SplitLinear.set_prefill(True) + + for _ in range(N_W): + c = [KVCache() for _ in range(N)] + mx.eval(lang(ids, cache=c, position_ids=pos).logits) + + # Accuracy + c_hyb = [KVCache() for _ in range(N)] + hyb = np.array(lang(ids, cache=c_hyb, position_ids=pos).logits.astype(mx.float32)) + cos = float(np.dot(ref.flatten(), hyb.flatten()) / + (np.linalg.norm(ref.flatten()) * np.linalg.norm(hyb.flatten()) + 1e-12)) + top1 = float((ref.argmax(-1) == hyb.argmax(-1)).mean() * 100) + print(f" cos={cos:.6f}, top1={top1:.1f}%") + + # Benchmark + ts = [] + for i in range(N_B): + c = [KVCache() for _ in range(N)] + t0 = time.perf_counter() + mx.eval(lang(ids, cache=c, position_ids=pos).logits) + t = (time.perf_counter()-t0)*1000 + ts.append(t) + print(f" Run {i+1}: {t:.1f}ms") + med = float(np.median(ts)) + + print(f"\n{'='*60}") + print(f" GPU FP16: {bl:.1f}ms") + print(f" SplitLinear: {med:.1f}ms ({bl/med:.3f}x)") + print(f" cos={cos:.6f} top1={top1:.1f}%") + print(f" delta: {med - bl:+.1f}ms") + print(f"{'='*60}") diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..652904b --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,39 @@ +[build-system] +requires = ["setuptools>=68", "wheel", "nanobind>=2.12", "mlx>=0.31", "cmake"] +build-backend = "setuptools.build_meta" + +[project] +name = "cider" +version = "0.8.0" +description = "INT8 TensorOps quantized matmul + optimized SDPA for Apple M5+ (W8A8, W4A8)" +readme = "README.md" +license = "MIT" +requires-python = ">=3.12" +authors = [ + {name = "Wang Shuo"}, +] +dependencies = [ + "mlx>=0.31", + "numpy", +] +classifiers = [ + "Development Status :: 3 - Alpha", + "Intended Audience :: Developers", + "Intended Audience :: Science/Research", + "Operating System :: MacOS", + "Programming Language :: Python :: 3", + "Programming Language :: C++", + "Topic :: Scientific/Engineering :: Artificial Intelligence", +] + +[project.urls] +Homepage = "https://github.com/Mininglamp-AI/cider" + +[tool.setuptools.packages.find] +include = ["cider*", "vlm_service*"] + +[tool.setuptools.package-data] +"cider" = ["kernels/*.metal", "kernels/*.h", "lib/*.so", "lib/*.dylib"] + +[project.optional-dependencies] +serve = ["fastapi", "uvicorn", "pyyaml"] diff --git a/session_data/replay_prompt.py b/session_data/replay_prompt.py new file mode 100644 index 0000000..2650dc9 --- /dev/null +++ b/session_data/replay_prompt.py @@ -0,0 +1,188 @@ +#!/usr/bin/env python3 +"""Replay the prompt that LocalAgent would have built at a given step. + +Usage: + python replay_prompt.py + python replay_prompt.py ~/mano-trajectory/result/sess-20260415-153428-1fe23a68 3 + +Output: JSON to stdout with keys: + task, step, prompt, images (base64 list), system_prompt + +Step is 0-indexed. Step 0 = first prediction (1 image: current screenshot). +""" + +import argparse +import base64 +import io +import json +import os +import re +import sys + +from PIL import Image + +SCREENSHOT_WIDTH = 1280 +HISTORY_IMAGE_COUNT = 2 + +SYSTEM_PROMPT = "You are a helpful assistant." + +ACTION_SPACE = """\ +click(start_box='<|box_start|>(x1,y1)<|box_end|>') +doubleclick(start_box='<|box_start|>(x1,y1)<|box_end|>') +select(start_box='<|box_start|>(x1,y1)<|box_end|>') +drag(start_box='<|box_start|>(x1,y1)<|box_end|>', end_box='<|box_start|>(x3,y3)<|box_end|>') +hotkey(key='') +wait() +call_user() +type(content='', start_box='<|box_start|>(x1,y1)<|box_end|>') +stop(reason='') +scroll(direction='down or up or right or left') +scrollmenu(start_box='<|box_start|>(x1,y1),(x2,y2)<|box_end|>', direction='down or up or right or left') +finish()""" + +INSTRUCTION_TEMPLATE = """\ +You are a GUI agent. You are given a task and your action history, with screenshots. You need to perform the next action to complete the task. When outputting, output the thought process for the next action between the and tags, its description between the and tags, and the action itself between the and tags. + +## Output Format +think process +next action description +next action + +## Action Space +{action_space} + +## User Instruction +### task: {task} +### action history: {history} +{current_screenshot}""" + + +def load_session(session_dir: str): + """Load result.json and return task + history_resps.""" + result_path = os.path.join(session_dir, "result.json") + with open(result_path, "r", encoding="utf-8") as f: + data = json.load(f) + return data["task"], data["history_resps"] + + +def parse_action_desc(resp_str: str) -> str: + """Extract action description from history_resps entry. + Prefer (matches runtime behavior), fallback to parsing . + """ + # Prefer action_desp (what LocalAgent actually stores in prompt_history) + m = re.search(r"(.*?)", resp_str, re.DOTALL) + if m and m.group(1).strip(): + return m.group(1).strip() + + # Fallback: parse from tag + m = re.search(r"(.*?)", resp_str, re.DOTALL) + if not m: + return "unknown" + action_str = m.group(1).strip() + + coord_match = re.search(r"'action':\s*'(\w+)'", action_str) + action_name = coord_match.group(1) if coord_match else "unknown" + + coord_match2 = re.search(r"'coordinate':\s*\[(\d+),\s*(\d+)\]", action_str) + if coord_match2: + return f"{action_name} ({coord_match2.group(1)}, {coord_match2.group(2)})" + + text_match = re.search(r"'text':\s*'([^']*)'", action_str) + if text_match: + return f"{action_name}: {text_match.group(1)[:30]}" + + dir_match = re.search(r"'scroll_direction':\s*'(\w+)'", action_str) + if dir_match: + return f"{action_name} {dir_match.group(1)}" + + if "DONE" in action_str: + return "DONE" + if "FAIL" in action_str: + return "FAIL" + + return action_name + + +def img_to_b64(img_path: str) -> str: + """Load image, resize to SCREENSHOT_WIDTH, return base64.""" + img = Image.open(img_path) + if img.width != SCREENSHOT_WIDTH: + ratio = SCREENSHOT_WIDTH / img.width + new_h = int(img.height * ratio) + img = img.resize((SCREENSHOT_WIDTH, new_h), Image.LANCZOS) + buf = io.BytesIO() + img.save(buf, format="PNG") + return base64.b64encode(buf.getvalue()).decode("utf-8") + + +def build_prompt_at_step(session_dir: str, step: int): + """Reconstruct the prompt that LocalAgent would build at the given step.""" + task, history_resps = load_session(session_dir) + trajectory_dir = os.path.join(session_dir, "trajectory") + + if step < 0 or step >= len(history_resps): + sys.exit(json.dumps({"error": f"step {step} out of range (0-{len(history_resps)-1})"})) + + # Build prompt_history for steps 0..step-1 (previous steps) + prompt_history = [] + for i in range(step): + desc = parse_action_desc(history_resps[i]) + screenshot_path = os.path.join(trajectory_dir, f"{i}.png") + screenshot_b64 = img_to_b64(screenshot_path) if os.path.exists(screenshot_path) else "" + prompt_history.append({"desc": desc, "screenshot_b64": screenshot_b64}) + + # Current screenshot = step.png + current_path = os.path.join(trajectory_dir, f"{step}.png") + if not os.path.exists(current_path): + sys.exit(json.dumps({"error": f"{current_path} not found"})) + current_b64 = img_to_b64(current_path) + + # Build prompt text + images list (same logic as LocalAgent._build_prompt) + images = [] + recent = prompt_history[-(HISTORY_IMAGE_COUNT + 1):] + + history_parts = [] + for i, h in enumerate(prompt_history): + step_num = i + 1 + desc = h["desc"] + if h in recent and h.get("screenshot_b64"): + images.append(h["screenshot_b64"]) + history_parts.append(f"第{step_num}步：{desc}，对应的截图为") + else: + history_parts.append(f"第{step_num}步：{desc}") + + history_text = "\n".join(history_parts) if history_parts else "无" + images.append(current_b64) + + prompt_text = INSTRUCTION_TEMPLATE.format( + action_space=ACTION_SPACE, + task=task, + history=history_text, + current_screenshot="当前截图为", + ) + + return { + "task": task, + "step": step, + "system_prompt": SYSTEM_PROMPT, + "prompt": prompt_text, + "images": images, + "image_count": len(images), + "history": [h["desc"] for h in prompt_history], + } + + +def main(): + parser = argparse.ArgumentParser(description="Replay LocalAgent prompt at a given step") + parser.add_argument("session_dir", help="Path to session directory") + parser.add_argument("step", type=int, help="Step number (0-indexed)") + args = parser.parse_args() + + session_dir = os.path.expanduser(args.session_dir) + result = build_prompt_at_step(session_dir, args.step) + with open("prompt", "w") as fo: + json.dump(result, fo, ensure_ascii=False) + + +if __name__ == "__main__": + main() diff --git a/session_data/result.json b/session_data/result.json new file mode 100644 index 0000000..2c236e6 --- /dev/null +++ b/session_data/result.json @@ -0,0 +1,15 @@ +{ + "task": "Compare available plans for the AeroAPI on Flightaware.", + "agent_type": "local", + "history_resps": [ + "当前页面为Flightaware的首页，主要展示了搜索航班的目的地和一些导航栏选项。由于需要比较AeroAPI在Flightaware上的可用套餐，通常这类信息会在“Products”或相关菜单下展示。点击“Products”可以进入产品列表页面，进一步查找和比较AeroAPI的套餐选项。因此，点击“Products”是合理的下一步操作。\n[{'name': 'computer', 'input': {'action': 'left_click', 'coordinate': [196, 150]}, 'id': '65b7dd0d-06eb-4825-b122-11d861f2a586', 'action_type': 'tool_use'}]", + "当前页面已经进入了Flightaware的产品（Products）板块，左侧栏展示了多个产品选项，包括AeroAPI。为了比较AeroAPI在不同计划下的可用性，需要先进入AeroAPI的详细页面，查看其具体信息和可选计划。因此，接下来应点击“AeroAPI”选项，以获取相关计划详情。\n[{'name': 'computer', 'input': {'action': 'left_click', 'coordinate': [515, 349]}, 'id': 'ae2ac4c6-e686-492a-8e32-270231588479', 'action_type': 'tool_use'}]", + "当前页面显示了AeroAPI的详细信息，包括数据产品、应用和接收器等选项。为了比较不同计划，需要进一步查看页面上关于价格和订阅的详细信息，因此需要向下滚动页面以找到相关计划详情。\n[{'name': 'computer', 'input': {'action': 'scroll', 'scroll_direction': 'down', 'coordinate': [640, 360], 'scroll_amount': 30}, 'id': '1699e398-5cc4-4363-a126-09b34871e984', 'action_type': 'tool_use'}]", + "当前页面主要展示了Flightaware的产品和服务介绍，但没有看到具体的计划对比内容。为了找到AeroAPI的可用计划对比信息，需要继续向下滑动页面，查看是否有相关的计划列表或对比内容。页面应向下滑动，以便浏览更多内容，寻找目标信息。\n[{'name': 'computer', 'input': {'action': 'scroll', 'scroll_direction': 'down', 'coordinate': [640, 360], 'scroll_amount': 30}, 'id': '69f7ecd2-f82f-4702-80ba-b15fda7722cc', 'action_type': 'tool_use'}]", + "当前页面已经显示了AeroAPI的介绍和相关链接，但没有看到关于可用计划的详细信息。为了查找AeroAPI的可用计划，需要继续向下滚动页面，查看是否有相关的计划对比或详细介绍内容。页面应该向下滑动，以便查看更多信息。\n[{'name': 'computer', 'input': {'action': 'scroll', 'scroll_direction': 'down', 'coordinate': [640, 360], 'scroll_amount': 30}, 'id': 'a020e194-9bcf-441e-a2f4-a8393b85e9ec', 'action_type': 'tool_use'}]", + "当前页面主要展示了Flightaware的商业产品和解决方案，但没有直接显示AeroAPI的具体价格和套餐选项。页面上有“Request pricing”按钮，这通常是用于获取免费报价或套餐信息的入口。点击该按钮可以进入下一步，查看AeroAPI的不同套餐和价格详情。因此，接下来应该点击“Request pricing”按钮以获取更多关于AeroAPI套餐的信息。\n[{'name': 'computer', 'input': {'action': 'left_click', 'coordinate': [505, 590]}, 'id': 'b3048205-f135-479d-bd97-ec9236a78efb', 'action_type': 'tool_use'}]", + "已经按照操作步骤进入了相关的页面，并且所有必要的点击和滚动操作都已完成。页面上显示了所需的信息和按钮，没有进一步需要操作的地方。因此，任务已经执行完毕，可以结束操作。\n[{'action_type': 'DONE'}]" + ], + "elapsed_time_sec": 67.93, + "step_count": 7 +} \ No newline at end of file diff --git a/session_data/trajectory/0.png b/session_data/trajectory/0.png new file mode 100644 index 0000000..f5f9ec4 Binary files /dev/null and b/session_data/trajectory/0.png differ diff --git a/session_data/trajectory/1.png b/session_data/trajectory/1.png new file mode 100644 index 0000000..a84c868 Binary files /dev/null and b/session_data/trajectory/1.png differ diff --git a/session_data/trajectory/2.png b/session_data/trajectory/2.png new file mode 100644 index 0000000..3763eeb Binary files /dev/null and b/session_data/trajectory/2.png differ diff --git a/session_data/trajectory/3.png b/session_data/trajectory/3.png new file mode 100644 index 0000000..666eed3 Binary files /dev/null and b/session_data/trajectory/3.png differ diff --git a/session_data/trajectory/4.png b/session_data/trajectory/4.png new file mode 100644 index 0000000..b54bb89 Binary files /dev/null and b/session_data/trajectory/4.png differ diff --git a/session_data/trajectory/5.png b/session_data/trajectory/5.png new file mode 100644 index 0000000..691fd25 Binary files /dev/null and b/session_data/trajectory/5.png differ diff --git a/session_data/trajectory/6.png b/session_data/trajectory/6.png new file mode 100644 index 0000000..189e089 Binary files /dev/null and b/session_data/trajectory/6.png differ diff --git a/session_data/trajectory_visual/0.png b/session_data/trajectory_visual/0.png new file mode 100644 index 0000000..3993e0b Binary files /dev/null and b/session_data/trajectory_visual/0.png differ diff --git a/session_data/trajectory_visual/1.png b/session_data/trajectory_visual/1.png new file mode 100644 index 0000000..ad43e94 Binary files /dev/null and b/session_data/trajectory_visual/1.png differ diff --git a/session_data/trajectory_visual/2.png b/session_data/trajectory_visual/2.png new file mode 100644 index 0000000..18579e7 Binary files /dev/null and b/session_data/trajectory_visual/2.png differ diff --git a/session_data/trajectory_visual/3.png b/session_data/trajectory_visual/3.png new file mode 100644 index 0000000..9476d5b Binary files /dev/null and b/session_data/trajectory_visual/3.png differ diff --git a/session_data/trajectory_visual/4.png b/session_data/trajectory_visual/4.png new file mode 100644 index 0000000..0c97974 Binary files /dev/null and b/session_data/trajectory_visual/4.png differ diff --git a/session_data/trajectory_visual/5.png b/session_data/trajectory_visual/5.png new file mode 100644 index 0000000..00bf96d Binary files /dev/null and b/session_data/trajectory_visual/5.png differ diff --git a/session_data/trajectory_visual/6.png b/session_data/trajectory_visual/6.png new file mode 100644 index 0000000..98a81b5 Binary files /dev/null and b/session_data/trajectory_visual/6.png differ diff --git a/setup.py b/setup.py new file mode 100644 index 0000000..adaa562 --- /dev/null +++ b/setup.py @@ -0,0 +1,128 @@ +"""Build script for Cider. + +Usage: + pip install -e . # Editable install (dev) + pip install . # Regular install + python setup.py build_ext # Build C++ extension only + +The C++ extension is built via CMake and installed into cider/lib/. +On Apple M4 and below, the C++ extension is skipped (INT8 TensorOps +require M5+). The pure-Python package still installs and provides +graceful fallback via is_available() → False. +""" + +import os +import platform +import re +import subprocess +import sys +from pathlib import Path + +from setuptools import setup, Extension +from setuptools.command.build_ext import build_ext + + +def _detect_apple_chip() -> int: + """Detect Apple Silicon generation. Returns chip number (4, 5, ...) or 0 if unknown.""" + if platform.system() != "Darwin": + return 0 + try: + brand = subprocess.run( + ["sysctl", "-n", "machdep.cpu.brand_string"], + capture_output=True, text=True, timeout=5 + ).stdout.strip() + m = re.match(r"Apple M(\d+)", brand) + return int(m.group(1)) if m else 0 + except Exception: + return 0 + + +def _should_build_ext() -> bool: + """Determine if C++ extension should be built.""" + # Allow forcing via env var + force = os.environ.get("CIDER_FORCE_BUILD", "").lower() + if force in ("1", "true", "yes"): + return True + if force in ("0", "false", "no"): + return False + # Auto-detect: only build on M5+ + chip = _detect_apple_chip() + if chip >= 5: + return True + if chip > 0: + print(f"[cider] Detected Apple M{chip} — skipping C++ extension (requires M5+)") + return False + # Unknown chip (cross-compile, CI, etc.) — try to build + return True + + +class CMakeBuild(build_ext): + def build_extension(self, ext): + src_dir = Path(__file__).parent.resolve() + + # Resolve paths that cmake can't find in pip's isolated build env + import sysconfig + python_include = sysconfig.get_path("include") + + try: + import nanobind + nanobind_dir = nanobind.cmake_dir() + except ImportError: + nanobind_dir = "" + + cmake_args = [ + f"-DPython_EXECUTABLE={sys.executable}", + f"-DPython_INCLUDE_DIR={python_include}", + "-DCMAKE_BUILD_TYPE=Release", + ] + if nanobind_dir: + cmake_args.append(f"-Dnanobind_DIR={nanobind_dir}") + + build_args = ["--config", "Release", "-j"] + + # In-source build (cmake outputs .so to src_dir) + subprocess.check_call( + ["cmake", "."] + cmake_args, + cwd=src_dir, + ) + subprocess.check_call( + ["cmake", "--build", "."] + build_args, + cwd=src_dir, + ) + + import shutil + + # Copy .so to build_lib root (where setuptools expects the extension) + build_lib = Path(self.build_lib) + build_lib.mkdir(parents=True, exist_ok=True) + for f in src_dir.glob("_cider_prim*.so"): + shutil.copy2(f, build_lib) + + # Also copy to cider/lib/ for runtime loading + output_lib_dir = build_lib / "cider" / "lib" + output_lib_dir.mkdir(parents=True, exist_ok=True) + for f in src_dir.glob("*.so"): + shutil.copy2(f, output_lib_dir) + for f in src_dir.glob("*.dylib"): + shutil.copy2(f, output_lib_dir) + + # Write to source tree cider/lib/ so editable installs work + src_lib_dir = src_dir / "cider" / "lib" + src_lib_dir.mkdir(exist_ok=True) + for f in src_dir.glob("*.so"): + shutil.copy2(f, src_lib_dir) + for f in src_dir.glob("*.dylib"): + shutil.copy2(f, src_lib_dir) + + +if _should_build_ext(): + ext_modules = [Extension("_cider_prim", sources=[])] + cmdclass = {"build_ext": CMakeBuild} +else: + ext_modules = [] + cmdclass = {} + +setup( + ext_modules=ext_modules, + cmdclass=cmdclass, +) diff --git a/tools/convert_compressed_tensors_to_mlx.py b/tools/convert_compressed_tensors_to_mlx.py new file mode 100644 index 0000000..2aebf12 --- /dev/null +++ b/tools/convert_compressed_tensors_to_mlx.py @@ -0,0 +1,275 @@ +#!/usr/bin/env python3 +""" +Convert LLMCompressor compressed-tensors quantized model → MLX-native format. + +Supports any bit-width (W4A16, W8A16, etc.) produced by LLMCompressor's +pack-quantized format, including per-channel (strategy=channel) and per-group +(strategy=group) quantization. Handles cross-shard layers automatically. + +Usage: + python convert_compressed_tensors_to_mlx.py /path/to/src /path/to/dst [--verify] +""" +import argparse +import json +import os +import shutil +import glob +import numpy as np +from pathlib import Path + +_MLX_MAX_GROUP_SIZE = {4: 128, 8: 128} +_MLX_DEFAULT_GROUP_SIZE = 128 + + +def _effective_group_size(num_bits, config_group_size, in_features): + if config_group_size is not None and config_group_size > 0: + return config_group_size + max_gs = _MLX_MAX_GROUP_SIZE.get(num_bits, _MLX_DEFAULT_GROUP_SIZE) + return min(max_gs, in_features) + + +def convert(src_dir: str, dst_dir: str) -> dict: + from safetensors.torch import load_file, save_file + import torch + + src_dir = Path(src_dir) + dst_dir = Path(dst_dir) + dst_dir.mkdir(parents=True, exist_ok=True) + + with open(src_dir / "config.json") as f: + config = json.load(f) + + qconfig = config.get("quantization_config", {}) + group_cfg = qconfig.get("config_groups", {}).get("group_0", {}).get("weights", {}) + config_group_size = group_cfg.get("group_size", None) + num_bits = group_cfg.get("num_bits", 4) + symmetric = group_cfg.get("symmetric", True) + strategy = group_cfg.get("strategy", "group") + ignore = set(qconfig.get("ignore", [])) + + pack_factor = 32 // num_bits + offset = 1 << (num_bits - 1) + per_channel = (strategy == "channel") or (config_group_size is None) + + print(f"Config: {num_bits}-bit, strategy={strategy}, " + f"config_group_size={config_group_size}, " + f"symmetric={symmetric}, offset={offset}, ignore={ignore}") + + # Copy auxiliary files + skip_ext = {".safetensors"} + skip_names = {"model.safetensors.index.json", "config.json"} + for f in src_dir.iterdir(): + if f.suffix in skip_ext or f.name in skip_names: + continue + if f.is_file(): + shutil.copy2(f, dst_dir / f.name) + + # ── Load ALL shards, track which shard each key came from ──── + st_files = sorted(glob.glob(str(src_dir / "model*.safetensors"))) + if not st_files: + st_files = sorted(glob.glob(str(src_dir / "*.safetensors"))) + + # key → (tensor, source_filename) + all_tensors = {} + for st_file in st_files: + fname = os.path.basename(st_file) + print(f" Loading {fname}...") + for key, tensor in load_file(st_file).items(): + all_tensors[key] = (tensor, fname) + + # ── Classify: quantized vs plain ───────────────────────────── + quant_layers = {} # base → {packed, scale, shape} + plain = {} # key → (tensor, fname) + + for key, (tensor, fname) in all_tensors.items(): + matched = False + for suffix, part in [(".weight_packed", "packed"), + (".weight_scale", "scale"), + (".weight_shape", "shape")]: + if key.endswith(suffix): + base = key.removesuffix(suffix) + quant_layers.setdefault(base, {})[part] = (tensor, fname) + matched = True + break + if not matched: + plain[key] = (tensor, fname) + + # ── Convert ────────────────────────────────────────────────── + # Output: per-shard tensors + import torch + shard_out = {} # fname → {key: tensor} + weight_map = {} + total_converted = 0 + actual_group_size = None + + for base, parts in sorted(quant_layers.items()): + if "packed" not in parts: + print(f" ⚠ Skipping {base}: no weight_packed") + continue + if "scale" not in parts: + print(f" ⚠ Skipping {base}: no weight_scale") + continue + + packed, packed_fname = parts["packed"] + scale, _ = parts["scale"] + shape_tensor = parts.get("shape", (None, None))[0] + + if shape_tensor is not None: + in_features = int(shape_tensor[1].item()) + out_features = int(shape_tensor[0].item()) + else: + in_features = packed.shape[1] * pack_factor + out_features = packed.shape[0] + + gs = _effective_group_size(num_bits, config_group_size, in_features) + if actual_group_size is None: + actual_group_size = gs + print(f" Using group_size={gs} " + f"({'broadcast from per-channel' if per_channel else 'from config'})") + + # packed → uint32 as-is + mlx_weight = torch.from_numpy(packed.numpy().view(np.uint32).copy()) + + # scale → fp16, broadcast if per-channel + scale_fp16 = scale.to(torch.float16) + n_groups = in_features // gs + if scale_fp16.shape[1] < n_groups: + scale_fp16 = scale_fp16.expand(out_features, n_groups).contiguous() + + mlx_scales = scale_fp16 + mlx_biases = -float(offset) * mlx_scales + + # Write to the shard that owns weight_packed + shard_out.setdefault(packed_fname, {})[f"{base}.weight"] = mlx_weight + shard_out[packed_fname][f"{base}.scales"] = mlx_scales + shard_out[packed_fname][f"{base}.biases"] = mlx_biases + total_converted += 1 + + # Plain tensors + total_passthrough = 0 + for key, (tensor, fname) in plain.items(): + if tensor.dtype == torch.bfloat16: + tensor = tensor.to(torch.float16) + shard_out.setdefault(fname, {})[key] = tensor + total_passthrough += 1 + + # ── Save shards ────────────────────────────────────────────── + for fname, tensors in sorted(shard_out.items()): + out_path = dst_dir / fname + save_file(tensors, str(out_path)) + print(f" Saved {fname}: {len(tensors)} tensors") + for key in tensors: + weight_map[key] = fname + + # Index + with open(dst_dir / "model.safetensors.index.json", "w") as f: + json.dump({"metadata": {"total_size": 0}, "weight_map": weight_map}, f, indent=2) + + # Config + new_config = dict(config) + new_config.pop("quantization_config", None) + new_config["quantization"] = { + "group_size": actual_group_size or _MLX_DEFAULT_GROUP_SIZE, + "bits": num_bits, + } + with open(dst_dir / "config.json", "w") as f: + json.dump(new_config, f, indent=2) + + stats = {"bits": num_bits, "group_size": actual_group_size, + "strategy": strategy, "converted": total_converted, + "passthrough": total_passthrough} + print(f"\n✅ Done — {total_converted} quantized layers converted.") + print(f" Output: {dst_dir}") + return stats + + +def verify(src_dir: str, dst_dir: str, n_check: int = 3): + from safetensors.torch import load_file + import torch + + try: + import mlx.core as mx + except ImportError: + print("⚠ mlx not available — skipping verification") + return + + src_dir, dst_dir = Path(src_dir), Path(dst_dir) + + with open(src_dir / "config.json") as f: + src_cfg = json.load(f) + qconfig = src_cfg.get("quantization_config", {}) + group_cfg = qconfig.get("config_groups", {}).get("group_0", {}).get("weights", {}) + config_group_size = group_cfg.get("group_size", None) + num_bits = group_cfg.get("num_bits", 4) + + with open(dst_dir / "config.json") as f: + dst_cfg = json.load(f).get("quantization", {}) + mlx_group_size = dst_cfg.get("group_size", 128) + bits = dst_cfg.get("bits", num_bits) + offset = 1 << (bits - 1) + pack_factor = 32 // bits + + # Load all shards from both dirs + src_all = {} + for f in sorted(glob.glob(str(src_dir / "model*.safetensors"))): + src_all.update(load_file(f)) + dst_all = {} + for f in sorted(glob.glob(str(dst_dir / "model*.safetensors"))): + dst_all.update(load_file(f)) + + checked = 0 + for key in sorted(src_all): + if not key.endswith(".weight_packed") or checked >= n_check: + continue + base = key.removesuffix(".weight_packed") + + packed = src_all[f"{base}.weight_packed"].numpy().view(np.uint32) + scale = src_all[f"{base}.weight_scale"].to(torch.float32).numpy() + shape_t = src_all.get(f"{base}.weight_shape") + if shape_t is not None: + out_feat, in_feat = int(shape_t[0].item()), int(shape_t[1].item()) + else: + in_feat = packed.shape[1] * pack_factor + out_feat = packed.shape[0] + + ref = np.zeros((out_feat, in_feat), dtype=np.float32) + for i in range(pack_factor): + vals = (packed >> (bits * i)) & ((1 << bits) - 1) + signed = vals.astype(np.int32) - offset + cols = np.arange(packed.shape[1]) * pack_factor + i + cols = cols[cols < in_feat] + if scale.shape[1] == 1: + ref[:, cols] = signed[:, :len(cols)] * scale[:, 0:1] + else: + groups = cols // config_group_size + ref[:, cols] = signed[:, :len(cols)] * scale[:, groups] + + w = mx.array(dst_all[f"{base}.weight"].numpy().view(np.uint32)) + s = mx.array(dst_all[f"{base}.scales"].to(torch.float32).numpy()) + b = mx.array(dst_all[f"{base}.biases"].to(torch.float32).numpy()) + mlx_deq = np.array(mx.dequantize(w, s, b, group_size=mlx_group_size, bits=bits)) + + max_diff = np.max(np.abs(ref - mlx_deq)) + cos = (np.dot(ref.flatten(), mlx_deq.flatten()) + / (np.linalg.norm(ref.flatten()) * np.linalg.norm(mlx_deq.flatten()) + 1e-12)) + status = "✅" if cos > 0.9999 else "❌" + print(f" {status} {base}: max_diff={max_diff:.2e}, cos={cos:.8f}") + checked += 1 + + print(f"\n Verified {checked} layers") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="Convert LLMCompressor compressed-tensors → MLX format" + ) + parser.add_argument("src", help="Source compressed-tensors model directory") + parser.add_argument("dst", help="Destination MLX model directory") + parser.add_argument("--verify", action="store_true", + help="Spot-check dequantized values after conversion") + args = parser.parse_args() + + convert(args.src, args.dst) + if args.verify: + print("\n--- Verification ---") + verify(args.src, args.dst) diff --git a/tools/eval_ppl_all.py b/tools/eval_ppl_all.py new file mode 100644 index 0000000..496f4ab --- /dev/null +++ b/tools/eval_ppl_all.py @@ -0,0 +1,143 @@ +#!/usr/bin/env python3 +"""Unified PPL evaluation: FP16, W8A16, Cider per-channel, per-group gs=64, per-group gs=128.""" +import argparse, math, sys, time, gc, os +import mlx.core as mx +import mlx.nn as nn +import numpy as np +import tqdm +from mlx_lm.utils import get_total_parameters, load +from cider.nn import CiderLinear + +home_dir = os.path.expanduser("~") +WIKITEXT_LOCAL = os.path.join(home_dir, "Downloads/wikitext/wikitext-2-raw-v1") + +FP16_MODEL = os.path.join(home_dir, "Downloads/Meta-Llama-3-8B") +W8A16_MODEL = os.path.join(home_dir, "Downloads/Llama3-8B-GPTQ-W8A16-mlx") + +def load_wikitext_local(tokenizer, split="test"): + from datasets import load_dataset + split_map = { + "test": f"{WIKITEXT_LOCAL}/test-00000-of-00001.parquet", + } + data = load_dataset("parquet", data_files={split: split_map[split]}, split=split) + encodings = tokenizer.encode("\n\n".join(data["text"]), return_tensors="pt") + return encodings.numpy() + +def convert_model_with_gs(model, target_group_size): + counter = [0] + _TARGET = (nn.Linear, nn.QuantizedLinear) + def _walk(module): + for name, child in module.children().items(): + if isinstance(child, _TARGET): + setattr(module, name, CiderLinear.from_float(child, target_group_size=target_group_size)) + counter[0] += 1 + if counter[0] % 28 == 0: gc.collect() + elif isinstance(child, list): + for i, item in enumerate(child): + if isinstance(item, _TARGET): + child[i] = CiderLinear.from_float(item, target_group_size=target_group_size) + counter[0] += 1 + if counter[0] % 28 == 0: gc.collect() + elif isinstance(item, nn.Module): + _walk(item) + elif isinstance(child, nn.Module): + _walk(child) + _walk(model) + return counter[0] + +def eval_ppl(model, data, n_samples, seq_length): + all_losses = [] + total_chunks = data.shape[1] // seq_length + n = n_samples if n_samples > 0 else total_chunks + for i in tqdm.tqdm(range(n), desc="Evaluating PPL"): + batch = mx.array(data[:, i * seq_length: (i + 1) * seq_length]) + logits = model(batch[:, :-1]).astype(mx.float32) + losses = nn.losses.cross_entropy(logits, batch[:, 1:], reduction="none") + mx.eval(losses) + all_losses.append(losses.flatten()) + all_losses = mx.concatenate(all_losses) + mean_loss = all_losses.mean().item() + ppl = math.exp(mean_loss) + std_dev = mx.sqrt(mx.var(all_losses, ddof=1)).item() + se_ppl = ppl * (std_dev / math.sqrt(all_losses.size)) + return ppl, se_ppl + +def run_config(config_name, model_path, cider_gs=None, data=None, n_samples=50, seq_length=2048): + """Run a single config and return results dict.""" + print(f"\n{'='*60}") + print(f" {config_name}") + print(f"{'='*60}") + + print(f" Loading {model_path}...") + model, tokenizer = load(model_path, tokenizer_config={"trust_remote_code": True}) + + if cider_gs is not None: + gs_str = "per-channel" if cider_gs == 0 else f"per-group(gs={cider_gs})" + print(f" Converting to CiderLinear ({gs_str})...") + t0 = time.perf_counter() + n = convert_model_with_gs(model, cider_gs) + print(f" {n} layers converted in {time.perf_counter()-t0:.1f}s") + + # Warmup + dummy = mx.array([[1, 2, 3, 4, 5, 6, 7, 8]]) + _ = model(dummy); mx.eval(_) + + if data is None: + data = load_wikitext_local(tokenizer, "test") + + mx.reset_peak_memory() + start = time.time() + ppl, se = eval_ppl(model, data, n_samples, seq_length) + elapsed = time.time() - start + peak_mem = mx.get_peak_memory() / 1e9 + + result = {"name": config_name, "ppl": ppl, "se": se, "time": elapsed, "mem": peak_mem} + print(f" PPL = {ppl:.3f} ± {se:.3f} | {elapsed:.1f}s | {peak_mem:.2f} GB") + + # Free memory + del model + gc.collect() + mx.metal.clear_cache() + + return result, data + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--num-samples", type=int, default=-1) + parser.add_argument("--sequence-length", type=int, default=2048) + args = parser.parse_args() + + np.random.seed(123) + mx.random.seed(123) + + configs = [ + # ("Baseline FP16", FP16_MODEL, None), + # ("Baseline W8A16 (MLX native)", W8A16_MODEL, None), + # ("Cider W8A8 per-channel", W8A16_MODEL, 0), + ("Cider W8A8 per-group(gs=64)", W8A16_MODEL, 64), + ("Cider W8A8 per-group(gs=128)", W8A16_MODEL, 128), + ] + + results = [] + data = None + + for name, model_path, cider_gs in configs: + np.random.seed(123) + mx.random.seed(123) + r, data = run_config(name, model_path, cider_gs, data, args.num_samples, args.sequence_length) + results.append(r) + + # Summary table + print(f"\n\n{'='*70}") + print(f" PPL COMPARISON SUMMARY (Qwen3-8B, wikitext-2, seq={args.sequence_length}, n={args.num_samples})") + print(f"{'='*70}") + print(f" {'Config':<30} {'PPL':>10} {'±SE':>8} {'Time':>8} {'Mem':>8}") + print(f" {'-'*30} {'-'*10} {'-'*8} {'-'*8} {'-'*8}") + + baseline_ppl = results[0]["ppl"] if results else None + for r in results: + delta = f"(+{(r['ppl']-baseline_ppl)/baseline_ppl*100:.2f}%)" if baseline_ppl and r['ppl'] != baseline_ppl else "" + print(f" {r['name']:<30} {r['ppl']:>8.3f} ±{r['se']:.3f} {r['time']:>6.1f}s {r['mem']:>6.2f}GB {delta}") + +if __name__ == "__main__": + main() diff --git a/tools/smoothquant.py b/tools/smoothquant.py new file mode 100644 index 0000000..82b3358 --- /dev/null +++ b/tools/smoothquant.py @@ -0,0 +1,118 @@ +import base64 +from io import BytesIO +import json +from datasets import Dataset +import torch +from compressed_tensors.offload import dispatch_model +from datasets import load_dataset + +# Note: this is an optional utility for processing vision inputs for qwen. +# This can be installed via the "qwen" extra +from qwen_vl_utils import process_vision_info +from transformers import AutoProcessor, Qwen3VLForConditionalGeneration + +from llmcompressor import oneshot +from llmcompressor.modifiers.quantization import QuantizationModifier +from llmcompressor.modifiers.awq import AWQModifier +import random +# Load model. +model_id = "/mnt/data/ws/project/ms-swift/output/v3_sft_4b/v0-20260421-095015/checkpoint-50349" +model = Qwen3VLForConditionalGeneration.from_pretrained(model_id, torch_dtype=torch.bfloat16, device_map="auto") +processor = AutoProcessor.from_pretrained(model_id) + +# Oneshot arguments +DATASET_ID = "/mnt/data/ws/project/ms-swift/train_v4.jsonl" +DATASET_SPLIT = "test[:512]" +NUM_CALIBRATION_SAMPLES = 512 +MAX_SEQUENCE_LENGTH = 40000 + +# Load dataset and preprocess. +def reservoir_sampling(file_path, k): + """ + 从大文件中均匀随机采样 k 个样本 + + 时间复杂度: O(n) + 空间复杂度: O(k) + """ + reservoir = [] + + with open(file_path, 'r', encoding='utf-8') as f: + # 先填充前 k 个样本 + for i, line in enumerate(f): + l = line.strip() + if len(l) == 0: + continue + if i < k: + reservoir.append(l) + else: + # 对于第 i 个样本，以 k/i 的概率替换蓄水池中的某个样本 + j = random.randint(0, i) + if j < k: + reservoir[j] = l + + return reservoir + + +# Apply chat template and tokenize inputs. +calib_data_raw = reservoir_sampling(DATASET_ID, NUM_CALIBRATION_SAMPLES) + + +def preprocess_and_tokenize(example): + # preprocess + + text = processor.apply_chat_template( + example["messages"], tokenize=False, add_generation_prompt=True + ) + if "" not in text: + return None + text = text.replace("", "<|vision_start|><|image_pad|><|vision_end|>") + + res = processor( + text=[text], + images=example["images"], + padding=False, + max_length=MAX_SEQUENCE_LENGTH, + truncation=True, + ) + # tokenize + return res + +calib_data = [] +for i, item in enumerate(calib_data_raw): + print("load %d/%d items..." % (i, len(calib_data_raw))) + u = preprocess_and_tokenize(json.loads(item)) + if u is None: + continue + calib_data.append(preprocess_and_tokenize(json.loads(item))) + +calib_dataset = Dataset.from_list(calib_data) +# Define a oneshot data collator for multimodal inputs. +def data_collator(batch): + assert len(batch) == 1 + return {key: torch.tensor(value) for key, value in batch[0].items()} + + +# Recipe +recipe = [ + AWQModifier(duo_scaling=False), + QuantizationModifier( + scheme="W8A16", + ignore=["re:.*lm_head", "re:.*visual.*"], + ), +] +# Perform oneshot +oneshot( + model=model, + tokenizer=model_id, + dataset=calib_dataset, + recipe=recipe, + max_seq_length=MAX_SEQUENCE_LENGTH, + num_calibration_samples=NUM_CALIBRATION_SAMPLES, + data_collator=data_collator, + sequential_targets=["Qwen3VLTextDecoderLayer"], +) + +# Save to disk compressed. +SAVE_DIR = model_id.rstrip("/").split("/")[-1] + "-W8A16" +model.save_pretrained(SAVE_DIR, save_compressed=True) +processor.save_pretrained(SAVE_DIR) \ No newline at end of file diff --git a/tutorial/how_to_write_efficient_int_gemm_m5_en.md b/tutorial/how_to_write_efficient_int_gemm_m5_en.md new file mode 100644 index 0000000..99ec0cc --- /dev/null +++ b/tutorial/how_to_write_efficient_int_gemm_m5_en.md @@ -0,0 +1,428 @@ +# How to Write an Efficient INT8 GEMM Kernel on Apple M5 + +> A step-by-step optimization guide: from naive threadgroup staging to 2.78x speedup using Metal 4 TensorOps +> +> All performance numbers in this tutorial are **real measurements** on Apple M5 Pro, bit-exact verified. + +## Table of Contents + +1. [Background & Motivation](#1-background--motivation) +2. [Step 1: First INT8 TensorOps Kernel — Naive Threadgroup Staging](#2-step-1-first-int8-tensorops-kernel--naive-threadgroup-staging) +3. [Step 2: Eliminate Threadgroup Memory — Direct Device Reads](#3-step-2-eliminate-threadgroup-memory--direct-device-reads) +4. [Step 3: Multi-Tile Dispatch — Small Batch Occupancy Fix](#4-step-3-multi-tile-dispatch--small-batch-occupancy-fix) +5. [Step 4: Deep K-Loop (BK=512) — Amortize Loop Overhead](#5-step-4-deep-k-loop-bk512--amortize-loop-overhead) +6. [Step 5: Swizzle Dispatch — L2 Cache Locality](#6-step-5-swizzle-dispatch--l2-cache-locality) +7. [Key Technique: NAXFrag Register Layout](#7-key-technique-naxfrag-register-layout) +8. [Pitfalls & Lessons Learned](#8-pitfalls--lessons-learned) +9. [Final Results](#9-final-results) + +--- + +## 1. Background & Motivation + +Apple's M5 chip introduced **Metal 4** with **TensorOps** (`MetalPerformancePrimitives`). This provides a hardware `matmul2d` instruction operating on `cooperative_tensor` types — analogous to NVIDIA's Tensor Cores. + +The key operation: + +```metal +mpp::tensor_ops::matmul2d(ct_a, ct_b, ct_c); +// INT8 × INT8 → INT32, 16×32×16 tile +``` + +**Why INT8?** Modern LLM inference quantizes both weights and activations to INT8 (W8A8), halving memory bandwidth vs FP16 while the TensorOps hardware executes INT8 matmul at higher throughput. + +**The problem:** MLX's GEMM kernel (`nax.h`) dequantizes quantized weights back to FP16 before computing — it never uses INT8×INT8→INT32 TensorOps. We write one from scratch. + +### Hardware Specs (Apple M5 Pro) + +| Feature | Value | +|---------|-------| +| GPU Cores | 20 | +| TensorOps tile | 16×32×16 (M×K×N) | +| Supported types | INT8×INT8→INT32, FP16×FP16→FP32, BF16×BF16→FP32 | +| SIMD width | 32 threads/simdgroup | + +### Benchmark Setup + +All benchmarks use **INT8×INT8→INT32 matmul** (pure integer, no dequant), verified **bit-exact** against NumPy int64 reference. Median of 50 runs after 5 warmup iterations. + +Primary benchmark shape: **(128, 4096, 4096)** — representative of LLM inference prefill. + +--- + +## 2. Step 1: First INT8 TensorOps Kernel — Naive Threadgroup Staging + +Our first kernel follows the classic GPU GEMM pattern: load tiles from device memory to threadgroup (shared) memory, then from threadgroup to registers for TensorOps. + +**Config:** BM=128, BN=128, BK=128, SK=32, WM=4, WN=4 (512 threads) + +```metal +// Cooperative load: device → threadgroup +threadgroup int8_t tg_A[128 * 128]; // BM × BK +threadgroup int8_t tg_B[128 * 128]; // BK × BN + +for (uint idx = tid_in_tg; idx < BM*BK; idx += NUM_THREADS) { + tg_A[row * BK + col] = A[global_row * K + global_col]; +} +// ... same for B ... +threadgroup_barrier(mem_flags::mem_threadgroup); + +// Load from threadgroup → registers +nax_frag_load_tg(a_frags, &tg_A[...], BK, sc, ...); +nax_frag_load_tg(b_frags, &tg_B[...], BN, sc, ...); + +// TensorOps +gemm_op.run(ct_a, ct_b, ct_c); +``` + +### Step 1 Performance (Real M5 Pro Data) + +| Shape (M,K,N) | Step 1 (ms) | +|----------------|------------| +| (128, 4096, 4096) | **1.002** | +| (256, 4096, 4096) | **1.179** | +| (128, 3584, 18944) | **1.817** | +| (256, 3584, 18944) | **3.120** | + +This is our baseline. Now we optimize. + +--- + +## 3. Step 2: Eliminate Threadgroup Memory — Direct Device Reads + +**Key insight:** Apple Silicon uses **unified memory**. Unlike discrete GPUs where device→shared is a separate memory hierarchy, on Apple GPU the "device" and "threadgroup" address spaces share the same physical memory. The threadgroup staging adds latency without benefit. + +```metal +// BEFORE (Step 1): device → threadgroup → register +nax_frag_load_tg(dst, &tg_A[offset], BK, sc, ...); + +// AFTER (Step 2): device → register directly +template +inline void nax_frag_load(thread T *dst, const device T *src, + int ld, short2 sc, + short off_m = 0, short off_n = 0) { + src += (sc.y + off_m) * ld + (sc.x + off_n); + for (short i = 0; i < 2; i++) + for (short j = 0; j < kElemCols; j++) + dst[i * kElemCols + j] = src[(i * kElemRowsJump) * ld + j]; +} +``` + +We eliminate the `threadgroup` buffers, the cooperative load loop, and the `threadgroup_barrier`. Each simdgroup loads its own data directly from device memory. + +### Step 2 Performance + +| Shape (M,K,N) | Step 1 (ms) | Step 2 (ms) | Speedup | +|----------------|------------|------------|---------| +| (128, 4096, 4096) | 1.002 | **0.438** | **2.29x** | +| (256, 4096, 4096) | 1.179 | **0.429** | **2.75x** | +| (128, 3584, 18944) | 1.817 | **0.815** | **2.23x** | +| (256, 3584, 18944) | 3.120 | **1.125** | **2.77x** | + +**The single biggest optimization: 2.3-2.8x speedup.** Removing threadgroup staging is critical on unified memory architectures. + +--- + +## 4. Step 3: Multi-Tile Dispatch — Small Batch Occupancy Fix + +Step 1 and 2 use a single tile config: BM=128, BN=128, 512 threads. When M is small (e.g., M=16), the large tile wastes compute — most simdgroups process padding zeros. + +**Solution:** Add a small-tile kernel variant: + +| Config | BM | BN | Simdgroups | Threads | Best for | +|--------|----|----|-----------|---------|----------| +| Large | 128 | 128 | 16 (4×4) | 512 | M > 64 | +| Small | 32 | 128 | 4 (1×4) | 128 | M ≤ 64 | + +```metal +// Same template, different instantiation +template +void gemm_int32_impl(...) { + constexpr int SM = BM / WM; // Sub-tile per simdgroup + constexpr int SN = BN / WN; + // ... +} + +// Large: BM=128, WM=4, WN=4 → SM=32, SN=32 +kernel void int8_matmul_int32(...) { + gemm_int32_impl<128, 128, 128, 32, 4, 4>(...); +} + +// Small: BM=32, WM=1, WN=4 → SM=32, SN=32 +kernel void int8_matmul_int32_small(...) { + gemm_int32_impl<32, 128, 128, 32, 1, 4>(...); +} +``` + +Host-side selection: + +```cpp +bool use_small = (M <= 64); +``` + +### Step 3 Performance + +**Large tile (M≥128):** + +| Shape (M,K,N) | Step 2 (ms) | Step 3 (ms) | Speedup | +|----------------|------------|------------|---------| +| (128, 4096, 4096) | 0.438 | **0.408** | **1.07x** | +| (256, 4096, 4096) | 0.429 | **0.410** | **1.05x** | +| (128, 3584, 18944) | 0.815 | **0.797** | **1.02x** | +| (256, 3584, 18944) | 1.125 | **1.130** | 1.00x | + +**Small tile (M<128, Step 3 vs Step 2 not applicable — Step 2 has no small tile):** + +| Shape (M,K,N) | Step 3 (ms) | +|----------------|------------| +| (1, 4096, 4096) | **0.239** | +| (16, 4096, 4096) | **0.241** | +| (32, 4096, 4096) | **0.240** | +| (64, 4096, 4096) | **0.247** | + +For large shapes the gain is marginal (~5%), but **small-tile support is essential** — without it, M<128 would require padding to 128 rows, wasting compute and memory. + +--- + +## 5. Step 4: Deep K-Loop (BK=512) — Amortize Loop Overhead + +Increasing BK from 128 to 512 means each outer K-loop iteration processes 4x more data, reducing loop overhead and barrier frequency: + +``` +BK=128: K=4096 → 32 outer iterations, 32 barriers +BK=512: K=4096 → 8 outer iterations, 8 barriers (4x fewer) +``` + +The inner loop processes BK/SK = 512/32 = 16 sub-iterations per outer step: + +```metal +constexpr int BK = 512; // was 128 + +for (int kk0 = 0; kk0 < K/BK; kk0++) { + threadgroup_barrier(mem_flags::mem_none); + for (int kk1 = 0; kk1 < BK; kk1 += SK) { + // 16 sub-iterations of load + compute + int8_t a_frags[TM][TK][kElemsPerFrag]; + int8_t b_frags[TK][TN][kElemsPerFrag]; + volatile int compiler_barrier; // prevent register spilling + + // Load fragments from device + for (short mm = 0; mm < TM; mm++) + for (short kk = 0; kk < TK; kk++) + nax_frag_load(a_frags[mm][kk], sg_A + kk1, K, sc, ...); + + // TensorOps compute + for (...) gemm_op.run(ct_a, ct_b, ct_c); + (void)compiler_barrier; + } + sg_A += BK; + sg_B += BK * N; +} +``` + +**Why `volatile int compiler_barrier`?** The Metal compiler may hoist all loads above all computes within the inner loop, causing register spill to device memory. The volatile variable forces a scheduling boundary. + +### Step 4 Performance + +| Shape (M,K,N) | Step 3 (ms) | Step 4 (ms) | Speedup | +|----------------|------------|------------|---------| +| (128, 4096, 4096) | 0.408 | **0.363** | **1.12x** | +| (256, 4096, 4096) | 0.410 | **0.399** | **1.03x** | +| (128, 3584, 18944) | 0.797 | **0.781** | **1.02x** | +| (256, 3584, 18944) | 1.130 | **1.082** | **1.04x** | + +The (128,4096,4096) shape benefits most — BK=512 amortizes the 4096-element K dimension more efficiently. + +--- + +## 6. Step 5: Swizzle Dispatch — L2 Cache Locality + +Default threadgroup indexing assigns adjacent threadgroups to adjacent tile columns. When N is large, this causes L2 cache thrashing: + +``` +Default mapping: + TG(0,0)→tile(0,0) TG(1,0)→tile(0,1) TG(2,0)→tile(0,2) ... + → All TGs in a row touch different B columns → B evicted from L2 cache +``` + +**Swizzle** interleaves the mapping so nearby TGs share B tile columns: + +```metal +inline void swizzle_decode(uint2 tgid, uint swizzle_log, uint tiles_n, + thread uint &tid_y, thread uint &tid_x) { + uint tile = 1u << swizzle_log; + tid_y = tgid.y * tile + (tgid.x % tile); + tid_x = tgid.x / tile; +} +``` + +Host computes `swizzle_log`: + +```cpp +int swizzle_log = 0; +for (int s = 1; s <= 4; s++) { + if ((tiles_n / (1 << s)) >= 1) swizzle_log = s; +} +``` + +### Step 5 Performance (= Production Kernel) + +| Shape (M,K,N) | Step 4 (ms) | Step 5 (ms) | Speedup | +|----------------|------------|------------|---------| +| (128, 4096, 4096) | 0.363 | **0.361** | **1.01x** | +| (256, 4096, 4096) | 0.399 | **0.397** | **1.01x** | +| (128, 3584, 18944) | 0.781 | **0.794** | 0.98x | +| (256, 3584, 18944) | 1.082 | **1.086** | 1.00x | + +At these benchmark sizes, swizzle has minimal effect (~1%). The benefit grows with larger N and more GPU cores competing for L2. On shapes with N>18944 or on M5 Ultra (40 cores), swizzle would show larger gains. + +--- + +## 7. Key Technique: NAXFrag Register Layout + +Understanding the `cooperative_tensor` fragment layout is essential for writing correct TensorOps kernels. + +The `matmul2d(16,32,16)` instruction operates on a 16×16 logical tile, distributed across 32 threads (one simdgroup). Each thread holds **8 elements** arranged as 2 rows × 4 columns: + +``` +Thread layout in a 16×16 tile: + - 32 threads, each contributes 8 elements + - Elements span rows [fm, fm+8] and columns [fn..fn+3] + +short2 nax_get_coord(ushort lid) { + short qid = short(lid >> 2); // quad group (0-7) + short fm = ((qid & 4) | ((short(lid) >> 1) & 3)); // fragment row + short fn = ((qid & 2) | (short(lid) & 1)) * 4; // fragment col + return short2{fn, fm}; // (col, row) +} +``` + +**Fragment load** reads 8 scattered elements from a contiguous matrix: + +```metal +// Each thread reads 2 rows (stride 8) × 4 columns (contiguous) +template +inline void nax_frag_load(thread T *dst, const device T *src, + int ld, short2 sc, short off_m, short off_n) { + src += (sc.y + off_m) * ld + (sc.x + off_n); + for (short i = 0; i < 2; i++) + for (short j = 0; j < 4; j++) + dst[i * 4 + j] = src[(i * 8) * ld + j]; +} +``` + +**The `matmul2d` instruction then consumes two fragments:** +- **ct_a**: 8 elements of A (left input) +- **ct_b**: 16 elements of B (right input, spans 2 logical N tiles) +- **ct_c**: 16 elements of C (destination, matches ct_b layout) + +```metal +// The instruction signature: +// matmul2d(16, 32, 16): A[16×16] × B[16×32] → C[16×32] +// ct_a: 8 elements (one 16×16 fragment) +// ct_b: 16 elements (two 16×16 fragments packed) +// ct_c: 16 elements (two 16×16 fragments packed) + +for (short i = 0; i < 8; i++) ct_a[i] = a_frags[mm][kk][i]; +for (short i = 0; i < 8; i++) { + ct_b[i] = b_frags[kk][nn][i]; + ct_b[8 + i] = b_frags[kk][nn+1][i]; +} +gemm_op.run(ct_a, ct_b, ct_c); +``` + +--- + +## 8. Pitfalls & Lessons Learned + +### Pitfall 1: Threadgroup Memory on Unified Memory + +The most common mistake when porting GPU kernels to Apple Silicon. On discrete GPUs (NVIDIA, AMD), threadgroup/shared memory is physically separate SRAM — staging data there is essential. On Apple Silicon, threadgroup memory is carved from the same unified RAM, so staging adds latency without benefit. + +**Rule:** On Apple Silicon, prefer direct device reads unless you need inter-thread data sharing within a threadgroup. + +### Pitfall 2: Pipeline Cache Staleness + +MLX's `PipelineCache` uses directory path equality to decide whether to recompile: + +```cpp +if (kernel_dir_ == kernel_dir) return; // Stale! +``` + +If you modify `.metal` source without changing the directory path, you get stale pipelines — correct compilation but wrong runtime results. + +**Workaround:** Copy kernels to a fresh directory for each variant. + +### Pitfall 3: Compiler Register Spilling + +Without the `volatile int compiler_barrier`, the Metal compiler may hoist all fragment loads above all TensorOps calls, creating huge register pressure and spilling to device memory. The volatile acts as a compiler scheduling fence. + +### Pitfall 4: `mx.fast.metal_kernel` Incompatibility + +MLX's `mx.fast.metal_kernel` API auto-generates a wrapper function around your kernel body. This wrapper is incompatible with TensorOps `cooperative_tensor` operations, which require **all simdgroup threads to participate uniformly**. The auto-generated wrapper may insert non-uniform control flow that breaks cooperative semantics. + +**Solution:** Use the C++ primitive API (`mx::Primitive`) with `PipelineCache` for TensorOps kernels. + +--- + +## 9. Final Results + +### Cumulative Optimization Summary + +All measurements on Apple M5 Pro, shape **(128, 4096, 4096)**, INT8×INT8→INT32, bit-exact verified. + +| Step | Technique | Time (ms) | vs Previous | vs Step 1 | +|------|-----------|-----------|-------------|-----------| +| 1 | Naive TG staging | 1.002 | — | 1.00x | +| 2 | Direct device read | 0.438 | **2.29x** | **2.29x** | +| 3 | Multi-tile dispatch | 0.408 | 1.07x | 2.46x | +| 4 | Deep K-loop (BK=512) | 0.363 | 1.12x | 2.76x | +| 5 | Swizzle dispatch | 0.361 | 1.01x | **2.78x** | + +### Cross-Shape Performance (Step 5 = Production) + +| Shape (M,K,N) | Step 1 (ms) | Step 5 (ms) | Total Speedup | +|----------------|------------|------------|---------------| +| (128, 4096, 4096) | 1.002 | 0.361 | **2.78x** | +| (256, 4096, 4096) | 1.179 | 0.397 | **2.97x** | +| (128, 3584, 18944) | 1.817 | 0.794 | **2.29x** | +| (256, 3584, 18944) | 3.120 | 1.086 | **2.87x** | + +### Small Tile Performance (Step 3+) + +| Shape (M,K,N) | Step 3 (ms) | Step 5 (ms) | +|----------------|------------|------------| +| (1, 4096, 4096) | 0.239 | 0.214 | +| (16, 4096, 4096) | 0.241 | 0.237 | +| (32, 4096, 4096) | 0.240 | 0.234 | +| (64, 4096, 4096) | 0.247 | 0.250 | + +### Key Takeaways + +1. **Eliminating threadgroup staging** is the single biggest win (2.3x) on unified memory. +2. **Multi-tile dispatch** is essential for small-M occupancy, marginal for large M. +3. **BK depth** matters — BK=512 vs BK=128 gives ~12% from amortizing loop overhead. +4. **Swizzle** has minimal effect at moderate N but grows with larger grids. +5. **Every number in this tutorial is real** — all kernel variants are in `dev/step_kernels/` and independently verifiable. + +--- + +## Appendix: Reproducing These Results + +All step kernel variants are available in the `dev/step_kernels/` directory: + +``` +dev/step_kernels/ +├── step1/w8a8_matmul.metal # Naive TG staging +├── step2/w8a8_matmul.metal # Direct device read +├── step3/w8a8_matmul.metal # Multi-tile (BM=128 + BM=32) +├── step4/w8a8_matmul.metal # Deep K-loop (BK=512) +├── step5/w8a8_matmul.metal # Swizzle (= production) +└── bench_final.py # Benchmark script +``` + +Run: `python dev/step_kernels/bench_final.py` + +## License + +MIT. See [Cider] for the complete SDK. diff --git a/tutorial/how_to_write_efficient_int_gemm_m5_zh.md b/tutorial/how_to_write_efficient_int_gemm_m5_zh.md new file mode 100644 index 0000000..e9c8f71 --- /dev/null +++ b/tutorial/how_to_write_efficient_int_gemm_m5_zh.md @@ -0,0 +1,367 @@ +# 如何在 Apple M5 芯片上编写高效的 INT8 GEMM Kernel + +> 从朴素实现到 2.78 倍加速的逐步优化指南 —— Metal 4 TensorOps 实战 +> +> 本文所有性能数据均为 **Apple M5 Pro 真机实测**，bit-exact 精度验证通过。 + +## 目录 + +1. [背景与动机](#1-背景与动机) +2. [Step 1：朴素 INT8 TensorOps Kernel —— Threadgroup 暂存](#2-step-1朴素-int8-tensorops-kernel--threadgroup-暂存) +3. [Step 2：消除 Threadgroup —— 直读 Device Memory](#3-step-2消除-threadgroup--直读-device-memory) +4. [Step 3：多 Tile 调度 —— 小 Batch 占用率优化](#4-step-3多-tile-调度--小-batch-占用率优化) +5. [Step 4：深 K 循环 (BK=512) —— 摊薄循环开销](#5-step-4深-k-循环-bk512--摊薄循环开销) +6. [Step 5：Swizzle 调度 —— L2 Cache 局部性](#6-step-5swizzle-调度--l2-cache-局部性) +7. [核心技术：NAXFrag 寄存器布局](#7-核心技术naxfrag-寄存器布局) +8. [踩坑与教训](#8-踩坑与教训) +9. [最终结果](#9-最终结果) + +--- + +## 1. 背景与动机 + +Apple M5 芯片引入了 **Metal 4** 的 **TensorOps** (`MetalPerformancePrimitives`)，提供硬件级 `matmul2d` 指令，操作 `cooperative_tensor` 类型——类似 NVIDIA 的 Tensor Cores。 + +核心操作： + +```metal +mpp::tensor_ops::matmul2d(ct_a, ct_b, ct_c); +// INT8 × INT8 → INT32, 16×32×16 tile +``` + +**为什么用 INT8？** LLM 推理中，W8A8 方案将权重和激活同时量化为 INT8，内存带宽减半，且硬件 INT8 matmul 吞吐量高于 FP16。 + +**问题：** MLX 的 GEMM kernel (`nax.h`) 会先将量化权重反量化为 FP16 再计算——从未使用 INT8×INT8→INT32 TensorOps 路径。我们需要从头写一个。 + +### 硬件参数 (Apple M5 Pro) + +| 特性 | 值 | +|------|-----| +| GPU 核心数 | 20 | +| TensorOps Tile | 16×32×16 (M×K×N) | +| 支持类型 | INT8×INT8→INT32, FP16×FP16→FP32, BF16×BF16→FP32 | +| SIMD 宽度 | 32 线程/simdgroup | + +### Benchmark 配置 + +所有 benchmark 使用 **INT8×INT8→INT32**（纯整数 matmul），与 NumPy int64 参考值 **bit-exact** 验证。每个 shape 跑 50 次取中位数，5 次 warmup。 + +主要基准 shape：**(128, 4096, 4096)**——代表 LLM 推理 prefill 场景。 + +--- + +## 2. Step 1：朴素 INT8 TensorOps Kernel —— Threadgroup 暂存 + +第一版 kernel 采用经典 GPU GEMM 模式：先将数据从 device memory 加载到 threadgroup (shared) memory，再从 threadgroup 读到寄存器执行 TensorOps。 + +**配置：** BM=128, BN=128, BK=128, SK=32, WM=4, WN=4 (512 线程) + +```metal +// 协作加载：device → threadgroup +threadgroup int8_t tg_A[128 * 128]; // BM × BK +threadgroup int8_t tg_B[128 * 128]; // BK × BN + +for (uint idx = tid_in_tg; idx < BM*BK; idx += NUM_THREADS) { + tg_A[row * BK + col] = A[global_row * K + global_col]; +} +threadgroup_barrier(mem_flags::mem_threadgroup); + +// 从 threadgroup 读到寄存器 +nax_frag_load_tg(a_frags, &tg_A[...], BK, sc, ...); +nax_frag_load_tg(b_frags, &tg_B[...], BN, sc, ...); + +// 执行 TensorOps +gemm_op.run(ct_a, ct_b, ct_c); +``` + +### Step 1 实测数据 + +| Shape (M,K,N) | Step 1 (ms) | +|----------------|------------| +| (128, 4096, 4096) | **1.002** | +| (256, 4096, 4096) | **1.179** | +| (128, 3584, 18944) | **1.817** | +| (256, 3584, 18944) | **3.120** | + +这是我们的起点。 + +--- + +## 3. Step 2：消除 Threadgroup —— 直读 Device Memory + +**核心发现：** Apple Silicon 使用**统一内存**架构。跟独立 GPU (NVIDIA, AMD) 的 shared memory 是物理独立 SRAM 不同，Apple GPU 的 threadgroup memory 从同一块统一内存划分。Threadgroup 暂存只是增加了延迟，没有任何收益。 + +```metal +// 之前 (Step 1): device → threadgroup → register +nax_frag_load_tg(dst, &tg_A[offset], BK, sc, ...); + +// 之后 (Step 2): device → register 直读 +template +inline void nax_frag_load(thread T *dst, const device T *src, + int ld, short2 sc, + short off_m = 0, short off_n = 0) { + src += (sc.y + off_m) * ld + (sc.x + off_n); + for (short i = 0; i < 2; i++) + for (short j = 0; j < kElemCols; j++) + dst[i * kElemCols + j] = src[(i * kElemRowsJump) * ld + j]; +} +``` + +去掉 `threadgroup` buffer、协作加载循环和 `threadgroup_barrier`。每个 simdgroup 独立从 device memory 读数据。 + +### Step 2 实测数据 + +| Shape (M,K,N) | Step 1 (ms) | Step 2 (ms) | 加速比 | +|----------------|------------|------------|--------| +| (128, 4096, 4096) | 1.002 | **0.438** | **2.29x** | +| (256, 4096, 4096) | 1.179 | **0.429** | **2.75x** | +| (128, 3584, 18944) | 1.817 | **0.815** | **2.23x** | +| (256, 3584, 18944) | 3.120 | **1.125** | **2.77x** | + +**全篇最大优化：2.3-2.8 倍加速。** 统一内存架构下去掉 threadgroup 暂存是第一要务。 + +--- + +## 4. Step 3：多 Tile 调度 —— 小 Batch 占用率优化 + +Step 1/2 只有一种 tile 配置：BM=128, 512 线程。当 M 很小时（如 M=16），大 tile 浪费算力——大部分 simdgroup 在算 padding 的零。 + +**方案：** 增加小 tile kernel 变体： + +| 配置 | BM | BN | Simdgroups | 线程数 | 适用 | +|------|----|----|-----------|--------|------| +| 大 Tile | 128 | 128 | 16 (4×4) | 512 | M > 64 | +| 小 Tile | 32 | 128 | 4 (1×4) | 128 | M ≤ 64 | + +Host 侧选择： + +```cpp +bool use_small = (M <= 64); +``` + +### Step 3 实测数据 + +**大 tile (M≥128)：** + +| Shape (M,K,N) | Step 2 (ms) | Step 3 (ms) | 加速比 | +|----------------|------------|------------|--------| +| (128, 4096, 4096) | 0.438 | **0.408** | **1.07x** | +| (256, 4096, 4096) | 0.429 | **0.410** | **1.05x** | +| (128, 3584, 18944) | 0.815 | **0.797** | **1.02x** | +| (256, 3584, 18944) | 1.125 | **1.130** | 1.00x | + +**小 tile (M<128，Step 3 新增能力)：** + +| Shape (M,K,N) | Step 3 (ms) | +|----------------|------------| +| (1, 4096, 4096) | **0.239** | +| (16, 4096, 4096) | **0.241** | +| (32, 4096, 4096) | **0.240** | +| (64, 4096, 4096) | **0.247** | + +大 shape 上增益不大 (~5%)，但**小 tile 支持是刚需**——没有它，M<128 只能 pad 到 128 行，浪费算力和内存。 + +--- + +## 5. Step 4：深 K 循环 (BK=512) —— 摊薄循环开销 + +BK 从 128 增加到 512，外层 K 循环迭代次数减少 4 倍： + +``` +BK=128: K=4096 → 32 次外层迭代, 32 次 barrier +BK=512: K=4096 → 8 次外层迭代, 8 次 barrier (少 4 倍) +``` + +```metal +constexpr int BK = 512; // 从 128 改为 512 + +for (int kk0 = 0; kk0 < K/BK; kk0++) { + threadgroup_barrier(mem_flags::mem_none); + for (int kk1 = 0; kk1 < BK; kk1 += SK) { + // 16 次子迭代的 load + compute + int8_t a_frags[TM][TK][kElemsPerFrag]; + int8_t b_frags[TK][TN][kElemsPerFrag]; + volatile int compiler_barrier; // 防止寄存器溢出 + + // 加载 + 计算 + ... + (void)compiler_barrier; + } +} +``` + +**`volatile int compiler_barrier` 的作用：** Metal 编译器可能将所有 load 提升到所有 compute 之前，导致寄存器爆掉溢出到 device memory。volatile 变量强制编译器插入调度边界。 + +### Step 4 实测数据 + +| Shape (M,K,N) | Step 3 (ms) | Step 4 (ms) | 加速比 | +|----------------|------------|------------|--------| +| (128, 4096, 4096) | 0.408 | **0.363** | **1.12x** | +| (256, 4096, 4096) | 0.410 | **0.399** | **1.03x** | +| (128, 3584, 18944) | 0.797 | **0.781** | **1.02x** | +| (256, 3584, 18944) | 1.130 | **1.082** | **1.04x** | + +(128,4096,4096) 获益最大——BK=512 更高效地处理 4096 的 K 维度。 + +--- + +## 6. Step 5：Swizzle 调度 —— L2 Cache 局部性 + +默认 threadgroup 索引将相邻 TG 分配到相邻的 tile 列。当 N 很大时，会造成 L2 cache 颠簸： + +``` +默认映射: + TG(0,0)→tile(0,0) TG(1,0)→tile(0,1) TG(2,0)→tile(0,2) ... + → 同一行的 TG 访问不同的 B 列 → B 被反复从 DRAM 重新加载 +``` + +**Swizzle** 重映射 TG 索引，使相邻 TG 共享 B 的 tile 列： + +```metal +inline void swizzle_decode(uint2 tgid, uint swizzle_log, uint tiles_n, + thread uint &tid_y, thread uint &tid_x) { + uint tile = 1u << swizzle_log; + tid_y = tgid.y * tile + (tgid.x % tile); + tid_x = tgid.x / tile; +} +``` + +### Step 5 实测数据（= 生产 kernel） + +| Shape (M,K,N) | Step 4 (ms) | Step 5 (ms) | 加速比 | +|----------------|------------|------------|--------| +| (128, 4096, 4096) | 0.363 | **0.361** | **1.01x** | +| (256, 4096, 4096) | 0.399 | **0.397** | **1.01x** | +| (128, 3584, 18944) | 0.781 | **0.794** | 0.98x | +| (256, 3584, 18944) | 1.082 | **1.086** | 1.00x | + +在当前 benchmark 规模下 swizzle 效果不显著 (~1%)。在更大的 N 或更多 GPU 核心（如 M5 Ultra 40 核）下增益会更明显。 + +--- + +## 7. 核心技术：NAXFrag 寄存器布局 + +理解 `cooperative_tensor` 的 fragment 布局是写正确 TensorOps kernel 的关键。 + +`matmul2d(16,32,16)` 指令操作 16×16 逻辑 tile，由 32 个线程（一个 simdgroup）分布持有。每个线程持有 **8 个元素**，排列为 2 行 × 4 列： + +```metal +// 32 线程 → 16×16 tile +// 每个线程：2 行（间距 8）× 4 列（连续） +short2 nax_get_coord(ushort lid) { + short qid = short(lid >> 2); // quad group (0-7) + short fm = ((qid & 4) | ((short(lid) >> 1) & 3)); // 行 + short fn = ((qid & 2) | (short(lid) & 1)) * 4; // 列 + return short2{fn, fm}; // (col, row) +} +``` + +**`matmul2d` 的操作数：** +- **ct_a**: 8 个元素（一个 16×16 fragment） +- **ct_b**: 16 个元素（两个 16×16 fragment 拼接，覆盖 16×32） +- **ct_c**: 16 个元素（与 ct_b 同布局） + +```metal +for (short i = 0; i < 8; i++) ct_a[i] = a_frags[mm][kk][i]; +for (short i = 0; i < 8; i++) { + ct_b[i] = b_frags[kk][nn][i]; + ct_b[8 + i] = b_frags[kk][nn+1][i]; +} +gemm_op.run(ct_a, ct_b, ct_c); +``` + +--- + +## 8. 踩坑与教训 + +### 坑 1：统一内存上的 Threadgroup 暂存 + +从独立 GPU 移植 kernel 时最常犯的错误。NVIDIA/AMD 的 shared memory 是物理独立的 SRAM，暂存到那里是必要的。Apple Silicon 的 threadgroup memory 跟 device memory 共享物理内存，暂存反而增加延迟。 + +**结论：** Apple Silicon 上，除非需要线程间数据共享，否则优先直读 device memory。 + +### 坑 2：Pipeline Cache 未感知源码变更 + +MLX 的 `PipelineCache` 用目录路径字符串判断是否需要重编译： + +```cpp +if (kernel_dir_ == kernel_dir) return; // 旧的！ +``` + +修改 `.metal` 源码但不改目录路径，cache 返回旧 pipeline——编译正确但运行结果错误。 + +**绕过方法：** 每个 kernel 变体用独立目录。 + +### 坑 3：编译器寄存器溢出 + +没有 `volatile int compiler_barrier` 时，Metal 编译器可能把所有 fragment load 提到所有 TensorOps 调用之前，导致寄存器压力爆炸、溢出到 device memory。volatile 充当编译器调度屏障。 + +### 坑 4：`mx.fast.metal_kernel` 与 TensorOps 不兼容 + +MLX 的 `mx.fast.metal_kernel` API 会自动生成 wrapper 函数。这个 wrapper 与 TensorOps 的 `cooperative_tensor` 操作不兼容——`cooperative_tensor` 要求 simdgroup 内所有线程**统一参与**。 + +**解决方案：** TensorOps kernel 必须用 C++ primitive API (`mx::Primitive`) + `PipelineCache`。 + +--- + +## 9. 最终结果 + +### 累积优化汇总 + +Apple M5 Pro 实测，Shape **(128, 4096, 4096)**，INT8×INT8→INT32，bit-exact 验证。 + +| Step | 技术 | 耗时 (ms) | vs 上一步 | vs Step 1 | +|------|------|-----------|-----------|-----------| +| 1 | 朴素 TG 暂存 | 1.002 | — | 1.00x | +| 2 | 直读 device memory | 0.438 | **2.29x** | **2.29x** | +| 3 | 多 tile 调度 | 0.408 | 1.07x | 2.46x | +| 4 | 深 K 循环 (BK=512) | 0.363 | 1.12x | 2.76x | +| 5 | Swizzle 调度 | 0.361 | 1.01x | **2.78x** | + +### 跨 Shape 性能 (Step 5 = 生产 kernel) + +| Shape (M,K,N) | Step 1 (ms) | Step 5 (ms) | 总加速比 | +|----------------|------------|------------|----------| +| (128, 4096, 4096) | 1.002 | 0.361 | **2.78x** | +| (256, 4096, 4096) | 1.179 | 0.397 | **2.97x** | +| (128, 3584, 18944) | 1.817 | 0.794 | **2.29x** | +| (256, 3584, 18944) | 3.120 | 1.086 | **2.87x** | + +### 小 Tile 性能 (Step 3+) + +| Shape (M,K,N) | Step 3 (ms) | Step 5 (ms) | +|----------------|------------|------------| +| (1, 4096, 4096) | 0.239 | 0.214 | +| (16, 4096, 4096) | 0.241 | 0.237 | +| (32, 4096, 4096) | 0.240 | 0.234 | +| (64, 4096, 4096) | 0.247 | 0.250 | + +### 核心结论 + +1. **消除 threadgroup 暂存**是最大单项优化 (2.3x)，统一内存架构的第一要务 +2. **多 tile 调度**对小 M 是刚需，对大 M 增益不大 +3. **BK 深度**有效——BK=512 vs BK=128 带来 ~12% 提升 +4. **Swizzle** 在中等规模 N 下效果不大，大 grid 下更明显 +5. **本文所有数据均为真机实测**——所有 kernel 变体在 `dev/step_kernels/` 下可独立复现 + +--- + +## 附录：复现方法 + +所有 step kernel 变体位于 `dev/step_kernels/` 目录： + +``` +dev/step_kernels/ +├── step1/w8a8_matmul.metal # 朴素 TG 暂存 +├── step2/w8a8_matmul.metal # 直读 device memory +├── step3/w8a8_matmul.metal # 多 tile (BM=128 + BM=32) +├── step4/w8a8_matmul.metal # 深 K 循环 (BK=512) +├── step5/w8a8_matmul.metal # Swizzle (= 生产 kernel) +└── bench_final.py # Benchmark 脚本 +``` + +运行: `python dev/step_kernels/bench_final.py` + +## License + +MIT. 完整 SDK: [Cider] diff --git a/vlm_service/__init__.py b/vlm_service/__init__.py new file mode 100644 index 0000000..25b5122 --- /dev/null +++ b/vlm_service/__init__.py @@ -0,0 +1,13 @@ +"""cider.vlm_service — Qwen3-VL inference service with MLX patches.""" + +from .code_def import CustomGenerationResult, ErrorCode +from .custom_qwen3vl import custom_generate, custom_stream_generate +from .core_infer import HMInference + +__all__ = [ + "CustomGenerationResult", + "ErrorCode", + "custom_generate", + "custom_stream_generate", + "HMInference", +] diff --git a/vlm_service/bench_client.py b/vlm_service/bench_client.py new file mode 100644 index 0000000..7e9cd4d --- /dev/null +++ b/vlm_service/bench_client.py @@ -0,0 +1,101 @@ +#!/usr/bin/env python3 +"""Benchmark client: call server multiple times, report timing stats. +Uses the same request format as client.py (images array + text with tags). +""" +import requests +import json +import time +import base64 +import statistics + +BASE_URL = "http://server_ip:8341/v1/chat/completions" +HEADERS = { + "Content-Type": "application/json", + "Authorization": "Bearer test-key", +} + +def load_image_b64(path): + with open(path, "rb") as f: + return base64.b64encode(f.read()).decode() + +def single_request(images_b64, run_id): + payload = { + "model": "qwen3-vl", + "request_id": run_id, + "messages": [ + {"role": "system", "content": "You are a helpful assistant."}, + { + "role": "user", + "content": ( + "You are a GUI agent. You are given a task and your action history, " + "with screenshots. You need to perform the next action to complete the task.\n\n" + "## Output Format\nthink\ndesc\naction\n\n" + "## Action Space\nclick(start_box='<|box_start|>(x1,y1)<|box_end|>')\nscroll(direction='down or up')\nstop(reason='')\nfinish()\n\n" + "## User Instruction\n### task: 查看Wiley可持续发展目标10的图书有哪些。\n" + "### action history: 第1步：Click Research.，对应的截图为\n" + "第2步：Click SDG Hub.，对应的截图为\n\n" + "当前截图为" + ), + }, + ], + "images": images_b64, + "temperature": 0.7, + "max_tokens": 256, + } + + t0 = time.perf_counter() + resp = requests.post(BASE_URL, json=payload, headers=HEADERS, timeout=120) + total = time.perf_counter() - t0 + resp.raise_for_status() + data = resp.json() + + prefill_time = data.get("prefill_time", 0) + decode_tps = data.get("decode_tps", 0) + text = data["choices"][0]["message"]["content"][:80] if data.get("choices") else "" + + return { + "total": total, + "prefill_time": prefill_time, + "decode_tps": decode_tps, + "text": text, + } + +def main(): + img_dir = "~/work/images/group0" + im0 = load_image_b64(f"{img_dir}/0.png") + im1 = load_image_b64(f"{img_dir}/1.png") + im2 = load_image_b64(f"{img_dir}/2_resize1.png") + images = [im0, im1, im2] + + n_warmup = 1 + n_bench = 3 + + print(f"Images: 3 (group0/0.png, 1.png, 2_resize1.png)") + print(f"Warmup: {n_warmup}, Bench: {n_bench}") + print() + + # Warmup + for i in range(n_warmup): + r = single_request(images, f"warmup_{i}") + print(f"Warmup {i}: total={r['total']:.3f}s, prefill={r['prefill_time']:.3f}s, decode={r['decode_tps']:.1f} t/s") + + # Bench + results = [] + for i in range(n_bench): + r = single_request(images, f"bench_{i}") + results.append(r) + print(f"Run {i}: total={r['total']:.3f}s, prefill={r['prefill_time']:.3f}s, decode={r['decode_tps']:.1f} t/s") + + # Stats + prefills = [r["prefill_time"] for r in results] + decodes = [r["decode_tps"] for r in results] + totals = [r["total"] for r in results] + + print() + print(f"Prefill: median={statistics.median(prefills):.3f}s, mean={statistics.mean(prefills):.3f}s") + print(f"Decode: median={statistics.median(decodes):.1f} t/s, mean={statistics.mean(decodes):.1f} t/s") + print(f"Total: median={statistics.median(totals):.3f}s, mean={statistics.mean(totals):.3f}s") + print(f"Response: {results[-1]['text']}...") + +if __name__ == "__main__": + main() diff --git a/vlm_service/client.py b/vlm_service/client.py new file mode 100644 index 0000000..6530793 --- /dev/null +++ b/vlm_service/client.py @@ -0,0 +1,130 @@ +# client_openai.py +import requests +import base64 +import json + + +base_url="http://xx:8341/v1/chat/completions" +def image_to_base64(image_path: str) -> str: + """将图像转换为 base64""" + with open(image_path, "rb") as f: + return base64.b64encode(f.read()).decode('utf-8') + +def example_openai_non_stream(): + """使用 OpenAI SDK - 非流式""" + print("=" * 60) + print("OpenAI SDK 示例 - 非流式") + print("=" * 60) + im0 = image_to_base64("~/work/images/group0/0.png") + im1 = image_to_base64("~/work/images/group0/1.png") + im2 = image_to_base64("~/work/images/group0/2_resize1.png") + headers = { + "Content-Type": "application/json", + # "Authorization": "Bearer sk-your-key" # 如果需要认证 + } + payload = { + "model": "qwen3-vl", + "messages": [ + { + "content": "You are a helpful assistant.", + "role": "system", + }, + { + "content": "You are a GUI agent. You are given a task and your action history, with screenshots. You need to perform the next action to complete the task. When outputting, output the thought process for the next action between the and tags, its description between the and tags, and the action itself between the and tags.\n\n## Output Format\nthink process\nnext action description\nnext action\n\n## Action Space\nclick(start_box='<|box_start|>(x1,y1)<|box_end|>')\ndoubleclick(start_box='<|box_start|>(x1,y1)<|box_end|>')\nselect(start_box='<|box_start|>(x1,y1)<|box_end|>') # To expand a dropdown menu or select an item from it.\ndrag(start_box='<|box_start|>(x1,y1)<|box_end|>', end_box='<|box_start|>(x3,y3)<|box_end|>') # Drag an element from the start coordinate (x1,y1) to the end coordinate (x3,y3).\nhotkey(key='') # Trigger a keyboard shortcut.\nwait() # Sleep for 5s and take a screenshot to check for any changes.\ncall_user() # Request human assistance\ntype(content='', start_box='<|box_start|>(x1,y1)<|box_end|>') # First click on the textbox, and then type the content.\nstop(reason='') # If the item can not found in the image, give the reason\nscroll(direction='down or up or right or left') # If the all screenshot need to scroll on the `direction` side.\nscrollmenu(start_box='<|box_start|>(x1,y1),(x2,y2)<|box_end|>', direction='down or up or right or left') # If the part of screenshot need to scroll on the `direction` side and give the area need to scroll\ntable_get_data() # The table data extraction begins.\ntable_get_data_finish() # The table data extraction is completed.\nfinish() # The task is completed.\n\n## User Instruction\n### task: 查看Wiley可持续发展目标10的图书有哪些。\n### action history: 第1步：Click the Research下拉框.，对应的截图为\n第2步：Click Sustainable Development Goals Hub.，对应的截图为\n\n当前截图为", + "role": "user" + }, + ], + "images": [im0, im1, im2], + "request_id": "tiandu", + "temperature": 0.7, + } + response = requests.post(base_url, json=payload, headers=headers) + response.raise_for_status() + result = response.json() + # 打印结果 + print(f"Request ID: {result['id']}") + print(f"AI: {result['choices'][0]['message']['content']}") + print(f"Prefill Time: {result.get('prefill_time', 0):.3f}s") + print(f"Decode Time: {result.get('decode_time', 0):.3f}s") + print(f"E2E Time: {result.get('e2e_time', 0):.3f}s") + + +def example_openai_stream(): + """使用 OpenAI SDK - 流式""" + print("=" * 60) + print("OpenAI SDK 示例 - 流式") + print("=" * 60) + + im0 = image_to_base64("~/work/images/group0/0.png") + im1 = image_to_base64("~/work/images/group0/1.png") + im2 = image_to_base64("~/work/images/group0/2_resize1.png") + with open("message.txt", "r") as f: + raw = json.load(f) + payload = { + "model": "qwen3-vl", + "request_id": "tiandu", + "temperature": 0.7, + "stream": True + } + messages = [] + for msg in raw: + if msg["role"] == "system": + messages.append({"role": "system", "content": msg["content"][0]["text"]}) + elif msg["role"] == "user": + final_content = "" + for content in msg["content"]: + if content["type"] == "text": + final_content += content["text"] + if "截图" in content["text"]: + final_content += "" + elif content["type"] == "image_url": + pass + messages.append({"role": "user", "content": final_content}) + payload["messages"] = messages + payload["images"] = [image_to_base64("x.png")] + # payload["request_id"] = "tiandu" + + payload = { + "model": "qwen3-vl", + "messages": [ + { + "content": "You are a helpful assistant.", + "role": "system", + }, + { + "content": "You are a GUI agent. You are given a task and your action history, with screenshots. You need to perform the next action to complete the task. When outputting, output the thought process for the next action between the and tags, its description between the and tags, and the action itself between the and tags.\n\n## Output Format\nthink process\nnext action description\nnext action\n\n## Action Space\nclick(start_box='<|box_start|>(x1,y1)<|box_end|>')\ndoubleclick(start_box='<|box_start|>(x1,y1)<|box_end|>')\nselect(start_box='<|box_start|>(x1,y1)<|box_end|>') # To expand a dropdown menu or select an item from it.\ndrag(start_box='<|box_start|>(x1,y1)<|box_end|>', end_box='<|box_start|>(x3,y3)<|box_end|>') # Drag an element from the start coordinate (x1,y1) to the end coordinate (x3,y3).\nhotkey(key='') # Trigger a keyboard shortcut.\nwait() # Sleep for 5s and take a screenshot to check for any changes.\ncall_user() # Request human assistance\ntype(content='', start_box='<|box_start|>(x1,y1)<|box_end|>') # First click on the textbox, and then type the content.\nstop(reason='') # If the item can not found in the image, give the reason\nscroll(direction='down or up or right or left') # If the all screenshot need to scroll on the `direction` side.\nscrollmenu(start_box='<|box_start|>(x1,y1),(x2,y2)<|box_end|>', direction='down or up or right or left') # If the part of screenshot need to scroll on the `direction` side and give the area need to scroll\ntable_get_data() # The table data extraction begins.\ntable_get_data_finish() # The table data extraction is completed.\nfinish() # The task is completed.\n\n## User Instruction\n### task: 查看Wiley可持续发展目标10的图书有哪些。\n### action history: 第1步：Click the Research下拉框.，对应的截图为\n第2步：Click Sustainable Development Goals Hub.，对应的截图为\n\n当前截图为", + "role": "user" + }, + ], + "images": [im0, im1, im2], + "request_id": "tiu", + "temperature": 0.7, + "stream": True + } + response = requests.post(base_url, json=payload, stream=True) + response.raise_for_status() + + print("Streaming response:") + for line in response.iter_lines(): + if line: + line = line.decode('utf-8') + if line.startswith('data: '): + data = line[6:] + if data == '[DONE]': + break + try: + chunk = json.loads(data) + if 'choices' in chunk and chunk['choices']: + delta = chunk['choices'][0].get('delta', {}) + content = delta.get('content', '') + if content: + print(content, end='', flush=True) + + except: + pass + print("\n") + + +if __name__ == "__main__": + example_openai_non_stream() + # example_openai_stream() diff --git a/vlm_service/code_def.py b/vlm_service/code_def.py new file mode 100644 index 0000000..7481882 --- /dev/null +++ b/vlm_service/code_def.py @@ -0,0 +1,23 @@ +from enum import Enum +from typing import Optional, List +from dataclasses import dataclass + +class ErrorCode(str, Enum): + SUCCESS = "success" + INVLID_MESSAGE_FORMAT = "invalid message format. Message should contain images not pure text." + VIT_FAILED = "Failed to extract image features. " + PREFILL_TOO_LONG = "The context length exceeds the maximum limit." + DECODE_TOO_LONG = "Decoding exceeds maximum length." + +@dataclass +class CustomGenerationResult: + text: str = "" + token: Optional[int] = None + logprobs: Optional[List[float]] = None + prompt_tokens: int = 0 + generation_tokens: int = 0 + total_tokens: int = 0 + prompt_tps: float = 0.0 + generation_tps: float = 0.0 + peak_memory: float = 0.0 + code: ErrorCode = ErrorCode.SUCCESS diff --git a/vlm_service/config.py b/vlm_service/config.py new file mode 100644 index 0000000..9cca1da --- /dev/null +++ b/vlm_service/config.py @@ -0,0 +1,123 @@ +import yaml +from pathlib import Path +from typing import Optional, Dict, Any +from dataclasses import dataclass, field +from loguru import logger + + +@dataclass +class ModelConfig: + """模型配置""" + model_name_or_path: str + + +@dataclass +class W8A8Config: + """W8A8 INT8 TensorOps 配置""" + mode: str = "auto" # "auto" | "on" | "off" + + +@dataclass +class SamplingConfig: + """采样配置""" + temperature: float = 1.0 + top_k: Optional[int] = None + top_p: float = 1.0 + repetition_penalty: float = 1.0 + max_new_tokens: int = 1024 + + +@dataclass +class ServerConfig: + """服务配置""" + host: str = "0.0.0.0" + port: int = 8000 + ttl: float = 1800 + max_image_buffer_size: int = 2 + + +@dataclass +class Config: + """总配置""" + model: ModelConfig + w8a8: W8A8Config = field(default_factory=W8A8Config) + sampling: SamplingConfig = field(default_factory=SamplingConfig) + server: ServerConfig = field(default_factory=ServerConfig) + + @classmethod + def from_yaml(cls, config_path: str) -> "Config": + """从 YAML 文件加载配置""" + config_path = Path(config_path) + if not config_path.exists(): + raise FileNotFoundError(f"Config file not found: {config_path}") + + with open(config_path, "r", encoding="utf-8") as f: + data = yaml.safe_load(f) + + # Parse w8a8 config + w8a8_raw = data.get("w8a8", {}) + if isinstance(w8a8_raw, str): + w8a8_cfg = W8A8Config(mode=w8a8_raw) + elif isinstance(w8a8_raw, dict): + w8a8_cfg = W8A8Config(**w8a8_raw) + else: + w8a8_cfg = W8A8Config() + + return cls( + model=ModelConfig(data["model_name_or_path"]), + w8a8=w8a8_cfg, + sampling=SamplingConfig(**data.get("sampling", {})), + server=ServerConfig(**data.get("server", {})), + ) + + @classmethod + def from_dict(cls, data: Dict[str, Any]) -> "Config": + """从字典加载配置""" + return cls( + model=ModelConfig(**data["model"]), + w8a8=W8A8Config(**data.get("w8a8", {})), + sampling=SamplingConfig(**data.get("sampling", {})), + server=ServerConfig(**data.get("server", {})), + ) + + def to_dict(self) -> Dict[str, Any]: + """转换为字典""" + return { + "model": self.model.__dict__, + "w8a8": self.w8a8.__dict__, + "sampling": self.sampling.__dict__, + "server": self.server.__dict__, + } + + def validate(self): + """验证配置""" + if not Path(self.model.model_name_or_path).exists(): + raise FileNotFoundError( + f"model not found: {self.model.model_name_or_path}" + ) + if self.w8a8.mode not in ("auto", "on", "off"): + raise ValueError( + f"w8a8.mode must be 'auto', 'on', or 'off', got '{self.w8a8.mode}'" + ) + logger.info("Config validation passed") + + +# 全局配置实例 +_config: Optional[Config] = None + + +def get_config() -> Config: + """获取全局配置""" + global _config + if _config is None: + raise RuntimeError("Config not initialized. Call load_config() first.") + return _config + + +def load_config(config_path: str = "config.yaml") -> Config: + """加载配置""" + global _config + _config = Config.from_yaml(config_path) + _config.validate() + logger.info(f"Config loaded from {config_path}") + return _config diff --git a/vlm_service/convert_qwen3vl_smooth_to_mlx.py b/vlm_service/convert_qwen3vl_smooth_to_mlx.py new file mode 100644 index 0000000..52338eb --- /dev/null +++ b/vlm_service/convert_qwen3vl_smooth_to_mlx.py @@ -0,0 +1,326 @@ +#!/usr/bin/env python3 +"""Convert LLMCompressor quantized model → MLX fp16 format. + +Supported input formats (auto-detected): + 1. pack-quantized: .weight_packed (int32) + .weight_scale + .weight_shape + 2. raw int8: .weight (int8) + .weight_scale + 3. Already fp16/bf16 (passthrough with key remap only) + +Key remapping (auto-detected): + If keys start with "model.language_model.*" / "model.visual.*" (HF Qwen3VL): + model.language_model.X → language_model.model.X + model.visual.X → vision_tower.X + lm_head.weight → dropped (tied weight) + If keys already match MLX format: no remapping needed. + +Usage: + python convert_smooth_to_mlx.py /path/to/source /path/to/output [--verify] +""" +import argparse, json, os, shutil, glob +import numpy as np +from pathlib import Path + + +# ─── Key Remapping ────────────────────────────────────────────────────────────── + +REMAP_RULES = [ + # (src_prefix, dst_prefix) + ("model.language_model.", "language_model.model."), + ("model.visual.", "vision_tower."), +] + +DROP_KEYS = {"lm_head.weight"} + + +def detect_key_format(keys): + """Detect whether keys need remapping.""" + for k in keys: + if k.startswith("model.language_model.") or k.startswith("model.visual."): + return "hf_qwen3vl" + if k.startswith("language_model.") or k.startswith("vision_tower."): + return "mlx_native" + return "unknown" + + +def remap_key(key, fmt): + """Remap key based on detected format. Returns None if key should be dropped.""" + if key in DROP_KEYS: + return None + + if fmt == "hf_qwen3vl": + for src_prefix, dst_prefix in REMAP_RULES: + if key.startswith(src_prefix): + return dst_prefix + key[len(src_prefix):] + # Unknown prefix in HF format → keep as-is (shouldn't happen) + return key + else: + # Already MLX native or unknown → no remap + return key + + +# ─── Unpacking ────────────────────────────────────────────────────────────────── + +def unpack_int8_from_int32(packed, original_shape): + """Unpack 4 x int8 values packed into int32 (little-endian byte order).""" + import torch + out_features = packed.shape[0] + target_out, target_in = int(original_shape[0]), int(original_shape[1]) + assert out_features == target_out + + raw_bytes = packed.contiguous().numpy().view(np.int8) + raw_bytes = raw_bytes.reshape(out_features, -1)[:, :target_in] + return torch.from_numpy(raw_bytes.copy()).to(torch.int8) + + +# ─── Shard Discovery ─────────────────────────────────────────────────────────── + +def get_shard_files(src_dir): + """Get shard files from index.json. Falls back to model-*-of-* glob.""" + src_dir = Path(src_dir) + idx_path = src_dir / "model.safetensors.index.json" + + if idx_path.exists(): + with open(idx_path) as f: + wm = json.load(f).get("weight_map", {}) + shard_names = sorted(set(wm.values())) + shard_files = [str(src_dir / n) for n in shard_names if (src_dir / n).exists()] + if shard_files: + return shard_files + + # Fallback: only sharded files (avoid stale single model.safetensors) + shards = sorted(glob.glob(str(src_dir / "model-*-of-*.safetensors"))) + if shards: + return shards + + # Last resort: single file + single = src_dir / "model.safetensors" + if single.exists(): + return [str(single)] + + raise FileNotFoundError(f"No safetensors files found in {src_dir}") + + +# ─── Main Conversion ─────────────────────────────────────────────────────────── + +def convert(src_dir, dst_dir): + from safetensors.torch import load_file, save_file + import torch + + src_dir, dst_dir = Path(src_dir), Path(dst_dir) + dst_dir.mkdir(parents=True, exist_ok=True) + + with open(src_dir / "config.json") as f: + config = json.load(f) + + qconfig = config.get("quantization_config", {}) + fmt = qconfig.get("format", "none") + print(f" Source quantization format: {fmt}") + + # Copy non-safetensor files (tokenizer, etc.) + for item in src_dir.iterdir(): + if item.suffix == ".safetensors" or item.name in {"model.safetensors.index.json", "config.json"}: + continue + if item.is_file(): + shutil.copy2(item, dst_dir / item.name) + + # Load tensors + st_files = get_shard_files(src_dir) + all_tensors = {} + for st_file in st_files: + print(f" Loading {os.path.basename(st_file)}...") + for key, tensor in load_file(st_file).items(): + all_tensors[key] = tensor + + print(f" Loaded {len(all_tensors)} tensors from {len(st_files)} shard(s)") + + # Detect key format + key_fmt = detect_key_format(all_tensors.keys()) + print(f" Key format: {key_fmt}") + + # ─── Identify quantized layers ─── + # Case 1: pack-quantized (weight_packed + weight_scale + weight_shape) + packed_bases = set() + for k in all_tensors: + if k.endswith(".weight_packed"): + packed_bases.add(k.removesuffix(".weight_packed")) + + # Case 2: raw int8 (weight is int8 + weight_scale exists) + raw_int8_bases = set() + for k in all_tensors: + if k.endswith(".weight_scale"): + base = k.removesuffix(".weight_scale") + if base not in packed_bases: + w_key = f"{base}.weight" + if w_key in all_tensors and all_tensors[w_key].dtype == torch.int8: + raw_int8_bases.add(base) + + print(f" Pack-quantized layers: {len(packed_bases)}") + print(f" Raw INT8 layers: {len(raw_int8_bases)}") + + # Mark all auxiliary keys + aux_keys = set() + for base in packed_bases: + aux_keys.add(f"{base}.weight_packed") + aux_keys.add(f"{base}.weight_scale") + aux_keys.add(f"{base}.weight_shape") + aux_keys.add(f"{base}.weight") # may not exist + for base in raw_int8_bases: + aux_keys.add(f"{base}.weight") + aux_keys.add(f"{base}.weight_scale") + + # ─── Build output ─── + output_tensors = {} + n_dequant = 0 + n_passthrough = 0 + n_dropped = 0 + + # 1) Dequant pack-quantized + for base in sorted(packed_bases): + packed = all_tensors[f"{base}.weight_packed"] + scale = all_tensors[f"{base}.weight_scale"] + shape_t = all_tensors[f"{base}.weight_shape"] + original_shape = shape_t.tolist() + + w_int8 = unpack_int8_from_int32(packed, original_shape) + w_fp = (w_int8.to(torch.float32) * scale.to(torch.float32)).to(torch.bfloat16) + + out_key = remap_key(f"{base}.weight", key_fmt) + if out_key is not None: + output_tensors[out_key] = w_fp + n_dequant += 1 + else: + n_dropped += 1 + + # 2) Dequant raw int8 + for base in sorted(raw_int8_bases): + w = all_tensors[f"{base}.weight"] + scale = all_tensors[f"{base}.weight_scale"] + w_fp = (w.to(torch.float32) * scale.to(torch.float32)).to(torch.bfloat16) + + out_key = remap_key(f"{base}.weight", key_fmt) + if out_key is not None: + output_tensors[out_key] = w_fp + n_dequant += 1 + else: + n_dropped += 1 + + # 3) Passthrough everything else + for key in sorted(all_tensors.keys()): + if key in aux_keys: + continue + + out_key = remap_key(key, key_fmt) + if out_key is None: + n_dropped += 1 + continue + + tensor = all_tensors[key] + if tensor.dtype == torch.bfloat16: + tensor = tensor.to(torch.bfloat16) + elif tensor.dtype == torch.float32: + tensor = tensor.to(torch.bfloat16) + + output_tensors[out_key] = tensor + n_passthrough += 1 + + if n_dequant % 50 != 0: + print(f" Dequantized {n_dequant} layers total") + + # ─── Save ─── + out_file = "model.safetensors" + save_file(output_tensors, str(dst_dir / out_file)) + print(f" Saved {out_file}: {len(output_tensors)} tensors") + + weight_map = {k: out_file for k in sorted(output_tensors.keys())} + with open(dst_dir / "model.safetensors.index.json", "w") as f: + json.dump({"metadata": {"total_size": 0}, "weight_map": weight_map}, f, indent=2) + + # Write clean config + new_config = dict(config) + new_config.pop("quantization_config", None) + with open(dst_dir / "config.json", "w") as f: + json.dump(new_config, f, indent=2) + + print(f"\n ✓ Summary: {n_dequant} dequantized, {n_passthrough} passthrough, {n_dropped} dropped") + print(f" Output sample:") + for k in sorted(output_tensors.keys())[:5]: + print(f" {k}: {output_tensors[k].dtype} {list(output_tensors[k].shape)}") + + return {"dequantized": n_dequant, "passthrough": n_passthrough, "dropped": n_dropped} + + +# ─── Verification ────────────────────────────────────────────────────────────── + +def verify(src_dir, dst_dir, n_check=5): + """Spot-check dequant accuracy.""" + from safetensors.torch import load_file + import torch + + src_all = {} + for f in get_shard_files(src_dir): + src_all.update(load_file(f)) + dst_all = {} + for f in sorted(glob.glob(str(Path(dst_dir) / "*.safetensors"))): + dst_all.update(load_file(f)) + + key_fmt = detect_key_format(src_all.keys()) + checked = 0 + + # Check pack-quantized + for k in sorted(src_all): + if checked >= n_check: + break + if not k.endswith(".weight_packed"): + continue + base = k.removesuffix(".weight_packed") + packed = src_all[k] + scale = src_all[f"{base}.weight_scale"] + shape_t = src_all[f"{base}.weight_shape"] + w_int8 = unpack_int8_from_int32(packed, shape_t.tolist()) + ref = (w_int8.to(torch.float32) * scale.to(torch.float32)).to(torch.bfloat16) + + out_key = remap_key(f"{base}.weight", key_fmt) + if out_key and out_key in dst_all: + out = dst_all[out_key] + diff = (ref.float() - out.float()).abs().max().item() + status = "✓" if diff == 0 else ("~" if diff < 0.01 else "✗") + print(f" {status} {out_key}: shape={list(out.shape)} max_diff={diff:.2e}") + checked += 1 + + # Check raw int8 + for k in sorted(src_all): + if checked >= n_check: + break + if not k.endswith(".weight_scale"): + continue + base = k.removesuffix(".weight_scale") + if f"{base}.weight_packed" in src_all: + continue + w_key = f"{base}.weight" + if w_key not in src_all or src_all[w_key].dtype != torch.int8: + continue + ref = (src_all[w_key].to(torch.float32) * src_all[k].to(torch.float32)).to(torch.bfloat16) + out_key = remap_key(w_key, key_fmt) + if out_key and out_key in dst_all: + out = dst_all[out_key] + diff = (ref.float() - out.float()).abs().max().item() + status = "✓" if diff == 0 else ("~" if diff < 0.01 else "✗") + print(f" {status} {out_key}: shape={list(out.shape)} max_diff={diff:.2e}") + checked += 1 + + if checked == 0: + print(" (no quantized layers to verify)") + + +if __name__ == "__main__": + p = argparse.ArgumentParser(description="Convert LLMCompressor quantized → MLX fp16") + p.add_argument("src", help="Source directory") + p.add_argument("dst", help="Output directory") + p.add_argument("--verify", action="store_true", help="Verify dequant accuracy") + args = p.parse_args() + + stats = convert(args.src, args.dst) + if args.verify: + print("\n--- Verify ---") + verify(args.src, args.dst) + print(f"\nDone. Load with: mlx_lm.load('{args.dst}')") diff --git a/vlm_service/core_infer.py b/vlm_service/core_infer.py new file mode 100644 index 0000000..ee9ab42 --- /dev/null +++ b/vlm_service/core_infer.py @@ -0,0 +1,260 @@ +from .code_def import ErrorCode +import time +import logging +from collections import defaultdict +from typing import List, Optional, Dict +import mlx_vlm as pm +from .custom_qwen3vl import * + +logging.basicConfig(level=logging.ERROR) +import warnings +import threading +from typing import Optional + +warnings.simplefilter(action="ignore", category=UserWarning) +warnings.simplefilter(action="ignore", category=FutureWarning) +import torch + +import numpy as np +from dataclasses import dataclass +from mlx_vlm import generate + +TARGET_TYPE = torch.float16 + + +class Timer: + """带自动统计功能的计时器""" + + _records: Dict[str, List[float]] = defaultdict(list) + _enabled = True + + def __init__(self, name: str = "Code block", verbose: bool = False): + self.name = name + self.verbose = verbose + self.elapsed = 0 + + def __enter__(self): + self.start = time.perf_counter() + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + self.elapsed = time.perf_counter() - self.start + if Timer._enabled: + Timer._records[self.name].append(self.elapsed) + if self.verbose: + print(f"[{self.name}] {self.elapsed:.4f}s") + return False + + @classmethod + def report(cls, sort_by: str = "total") -> None: + if not cls._records: + print("No timing records.") + return + print("\n" + "=" * 70) + print(f"{'Name':<30} {'Count':>8} {'Total':>10} {'Mean':>10} {'Min':>10} {'Max':>10}") + print("=" * 70) + stats = [] + for name, times in cls._records.items(): + stats.append({ + "name": name, "count": len(times), "total": sum(times), + "mean": sum(times) / len(times), "min": min(times), "max": max(times), + }) + if sort_by in ["total", "mean", "count"]: + stats.sort(key=lambda x: x[sort_by], reverse=True) + elif sort_by == "name": + stats.sort(key=lambda x: x["name"]) + for s in stats: + print(f"{s['name']:<30} {s['count']:>8} {s['total']:>10.4f}s {s['mean']:>10.4f}s {s['min']:>10.4f}s {s['max']:>10.4f}s") + print("=" * 70) + print(f"Total time: {sum(s['total'] for s in stats):.4f}s") + print() + + @classmethod + def get_stats(cls, name: Optional[str] = None) -> Dict: + if name: + times = cls._records.get(name, []) + if not times: + return {} + return { + "count": len(times), "total": sum(times), + "mean": sum(times) / len(times), + "min": min(times), "max": max(times), "times": times, + } + else: + return {k: cls.get_stats(k) for k in cls._records.keys()} + + @classmethod + def reset(cls, name: Optional[str] = None) -> None: + if name: + cls._records[name] = [] + else: + cls._records.clear() + + @classmethod + def disable(cls) -> None: + cls._enabled = False + + @classmethod + def enable(cls) -> None: + cls._enabled = True + + +class HMInference: + _instance: Optional['HMInference'] = None + _lock = threading.Lock() + _initialized = False + + def __new__(cls, *args, **kwargs): + if cls._instance is None: + with cls._lock: + if cls._instance is None: + cls._instance = super().__new__(cls) + return cls._instance + + def __init__(self, + model_path, + temperature=1.0, + topk=None, + topp=1.0, + repetition_penalty=1.0, + max_new_tokens=1024, + w8a8="auto"): + if HMInference._initialized: + return + with HMInference._lock: + if HMInference._initialized: + return + self.model, self.processor = pm.load(model_path) + + # ── W8A8 INT8 TensorOps via cider ── + # Replaces all Linear layers with CiderLinear: + # prefill mode → W8A8 INT8 TensorOps (~15-19% faster) + # decode mode → original weights (zero overhead) + self._w8a8_enabled = False + if w8a8 != "off": + try: + from cider import convert_model, set_mode, is_available + if w8a8 == "auto" and not is_available(): + logging.info( + "[W8A8] Hardware does not support INT8 TensorOps " + "(requires M5+), using default inference" + ) + else: + import mlx.core as mx + try: + stats = convert_model(self.model.language_model) + except: + stats = convert_model(self.model) + mx.eval(self.model.parameters()) + self._w8a8_set_mode = set_mode + self._w8a8_enabled = True + logging.info(f"[W8A8] cider enabled: {stats}") + except Exception as e: + if w8a8 == "on": + raise + logging.warning( + f"[W8A8] Init failed, using default inference: {e}" + ) + + self.temperature = temperature + self.topk = topk + self.topp = topp + self.repetition_penalty = repetition_penalty + self.max_new_tokens = max_new_tokens + HMInference._initialized = True + + def complete_stream(self, messages, images, buf_vis_feats, + buf_vis_stack_feats, **kwargs): + """流式推理接口""" + temperature = kwargs.pop("temperature", self.temperature) + topk = kwargs.pop("topk", self.topk) + topp = kwargs.pop("topp", self.topp) + repetition_penalty = kwargs.pop("repetition_penalty", self.repetition_penalty) + max_new_tokens = kwargs.pop("max_new_tokens", self.max_new_tokens) + + prompt = self.processor.tokenizer.apply_chat_template( + messages, tokenize=False, add_generation_prompt=True) + org_image_placeholder = "" + new_image_placeholder = "<|vision_start|><|image_pad|><|vision_end|>" + pi = len(images) + while pi > 0: + pi -= 1 + pos = prompt.rfind(org_image_placeholder) + if pos >= 0: + prompt = prompt[:pos] + prompt[pos:].replace( + org_image_placeholder, new_image_placeholder) + else: + break + + if self._w8a8_enabled: + self._w8a8_set_mode("prefill") + + for resp in custom_stream_generate( + self.model, + self.processor, + prompt, + images, + buf_vis_features=buf_vis_feats, + buf_vis_stack_features=buf_vis_stack_feats, + max_tokens=max_new_tokens, + temperature=temperature, + top_p=topp, + top_k=topk, + repetition_penalty=repetition_penalty, + verbose=True, + prefill_step_size=2048, + on_first_token=( + lambda: self._w8a8_set_mode("decode") + ) if self._w8a8_enabled else None, + ): + yield resp.code, resp.text, { + "prefill_time": resp.prompt_tokens / resp.prompt_tps, + "decode_tps": resp.generation_tps + } + + def complete(self, messages, images, buf_vis_feats, buf_vis_stack_feats, + **kwargs): + """非流式推理接口""" + temperature = kwargs.pop("temperature", self.temperature) + topk = kwargs.pop("topk", self.topk) + topp = kwargs.pop("topp", self.topp) + repetition_penalty = kwargs.pop("repetition_penalty", self.repetition_penalty) + max_new_tokens = kwargs.pop("max_new_tokens", self.max_new_tokens) + + prompt = self.processor.tokenizer.apply_chat_template( + messages, tokenize=False, add_generation_prompt=True) + org_image_placeholder = "" + new_image_placeholder = "<|vision_start|><|image_pad|><|vision_end|>" + pi = len(images) + while pi > 0: + pi -= 1 + pos = prompt.rfind(org_image_placeholder) + if pos >= 0: + prompt = prompt[:pos] + prompt[pos:].replace( + org_image_placeholder, new_image_placeholder) + else: + break + + if self._w8a8_enabled: + self._w8a8_set_mode("prefill") + + resp = custom_generate(self.model, + self.processor, + prompt, + images, + buf_vis_features=buf_vis_feats, + buf_vis_stack_features=buf_vis_stack_feats, + max_tokens=max_new_tokens, + temperature=temperature, + top_p=topp, + top_k=topk, + repetition_penalty=repetition_penalty, + prefill_step_size=2048, + on_first_token=( + lambda: self._w8a8_set_mode("decode") + ) if self._w8a8_enabled else None) + + return resp.code, resp.text, { + "prefill_time": resp.prompt_tokens / resp.prompt_tps, + "decode_tps": resp.generation_tps + } diff --git a/vlm_service/custom_qwen3vl.py b/vlm_service/custom_qwen3vl.py new file mode 100644 index 0000000..093bb77 --- /dev/null +++ b/vlm_service/custom_qwen3vl.py @@ -0,0 +1,925 @@ +from typing import Optional +import mlx.core as mx +from mlx_vlm.models.base import InputEmbeddingsFeatures +from mlx_vlm.models.qwen3_vl.language import LanguageModel as _OrigLanguageModel +import mlx_vlm.models.qwen3_vl as qwen_mod +import mlx.nn as nn +from tqdm import tqdm +from typing import Dict, Any, Callable, Tuple, List, Generator, Union +import functools +from mlx_vlm.generate import ( + maybe_quantize_kv_cache, + normalize_resize_shape, + wired_limit, +) +from mlx_lm.sample_utils import make_logits_processors, make_sampler +from mlx_vlm.models import cache +from mlx_vlm.utils import StoppingCriteria, ThinkingBudgetCriteria, prepare_inputs +import time +from .code_def import CustomGenerationResult, ErrorCode +from transformers import PreTrainedTokenizer +import numpy as np + +MAX_PROMPT_TOKENS = 8192 + + +def find_vis_token_positions(input_ids, target_token=151655): + if input_ids.ndim == 2: + input_ids = input_ids[0] + + # 找到所有匹配的位置 + mask = input_ids == target_token # [N] bool array + start_idx = mx.argmax(mask) + n_tokens = mx.sum(mask) + end_idx = start_idx + n_tokens + return start_idx.item(), end_idx.item() + + +class CustomLanguageModel(_OrigLanguageModel): + """Fix get_rope_index for multi-image: upstream uses mx.sum of vision_start + indices which collapses to a wrong scalar when >1 image is present. + Replaced with vectorized shifted-pair matching. + + Also fix __call__ for chunked prefill: upstream else-branch uses simple + arange+delta which is wrong for multi-image. When _position_ids is + pre-computed, always slice from it regardless of cache_offset.""" + + def __call__( + self, + inputs: mx.array, + inputs_embeds=None, + mask=None, + cache=None, + visual_pos_masks=None, + deepstack_visual_embeds=None, + **kwargs, + ): + from mlx_vlm.models.qwen3_vl.language import LanguageModelOutput + + n_to_process = kwargs.get("n_to_process", None) + if n_to_process is not None: + visual_pos_masks = ( + visual_pos_masks[:, n_to_process:] + if visual_pos_masks is not None + else None + ) + + position_ids = kwargs.pop("position_ids", None) + pixel_values = kwargs.pop("pixel_values", None) + image_grid_thw = kwargs.pop("image_grid_thw", None) + video_grid_thw = kwargs.pop("video_grid_thw", None) + + if pixel_values is not None: + self._rope_deltas = None + self._position_ids = None + + cache_offset = 0 + if cache and cache[0] is not None: + offset = cache[0].offset + if isinstance(offset, int): + cache_offset = offset + elif isinstance(offset, mx.array): + cache_offset = (offset if offset.ndim == 0 else offset[0]).item() + else: + raise ValueError(f"Unexpected cache offset type: {type(offset)}") + + rope_mask = mask + if mask is not None and mask.shape[-1] != inputs.shape[-1]: + rope_mask = None + + if position_ids is None and (rope_mask is None or rope_mask.ndim == 2): + if ( + self._position_ids is not None + and cache_offset + inputs.shape[1] <= self._position_ids.shape[2] + ): + # Slice from pre-computed position_ids (fixes chunked prefill) + seq_length = inputs.shape[1] + position_ids = self._position_ids[ + :, :, cache_offset : cache_offset + seq_length + ] + elif ( + (cache is not None and cache[0] is not None and cache_offset == 0) + or self._rope_deltas is None + or cache is None + ): + position_ids, rope_deltas = self.get_rope_index( + inputs, image_grid_thw, video_grid_thw, rope_mask + ) + self._rope_deltas = rope_deltas + self._position_ids = position_ids + else: + batch_size, seq_length = inputs.shape + delta = mx.array( + cache_offset + self._rope_deltas if cache is not None else 0 + ) + position_ids = mx.arange(seq_length).reshape(1, -1) + position_ids = mx.broadcast_to(position_ids, (batch_size, seq_length)) + if cache_offset is not None: + if delta.ndim == 0: + delta = mx.expand_dims(delta, axis=0) + if delta.shape[0] < batch_size: + delta = mx.tile(delta, (batch_size, 1)) + else: + delta = delta[:batch_size] + position_ids = mx.add(position_ids, delta)[None, ...] + position_ids = mx.broadcast_to( + position_ids, (3, batch_size, seq_length) + ) + + out = self.model( + inputs, + cache=cache, + inputs_embeds=inputs_embeds, + position_ids=position_ids, + visual_pos_masks=visual_pos_masks, + deepstack_visual_embeds=deepstack_visual_embeds, + ) + if self.args.tie_word_embeddings: + out = self.model.embed_tokens.as_linear(out) + else: + out = self.lm_head(out) + return LanguageModelOutput(logits=out) + + def get_rope_index( + self, + input_ids: mx.array, + image_grid_thw: Optional[mx.array] = None, + video_grid_thw: Optional[mx.array] = None, + attention_mask: Optional[mx.array] = None, + ): + batch_size, seq_length = input_ids.shape + position_ids = mx.arange(seq_length, dtype=mx.int32) + position_ids = mx.broadcast_to(position_ids[None, :], (batch_size, seq_length)) + spatial_merge_size = self.config.vision_config.spatial_merge_size + image_token_id = self.config.image_token_id + video_token_id = self.config.video_token_id + vision_start_token_id = self.config.vision_start_token_id + mrope_position_deltas = [] + if input_ids is not None and ( + image_grid_thw is not None or video_grid_thw is not None + ): + total_input_ids = input_ids + if attention_mask is None: + attention_mask = mx.ones_like(input_ids) + position_ids = mx.ones( + (3, input_ids.shape[0], input_ids.shape[1]), dtype=input_ids.dtype + ) + image_index, video_index = 0, 0 + for i, input_ids in enumerate(total_input_ids): + input_ids = mx.where( + attention_mask[i] == 1, input_ids, mx.zeros_like(input_ids) + ) + # Vectorized shifted-pair matching (fixes multi-image bug) + vs_mask = input_ids[:-1] == vision_start_token_id + next_toks = input_ids[1:] + image_nums = int(mx.sum(vs_mask & (next_toks == image_token_id)).item()) + video_nums = int(mx.sum(vs_mask & (next_toks == video_token_id)).item()) + input_tokens = input_ids.tolist() + llm_pos_ids_list: list = [] + st = 0 + remain_images, remain_videos = image_nums, video_nums + for _ in range(image_nums + video_nums): + if image_token_id in input_tokens and remain_images > 0: + ed_image = input_tokens.index(image_token_id, st) + else: + ed_image = len(input_tokens) + 1 + if video_token_id in input_tokens and remain_videos > 0: + ed_video = input_tokens.index(video_token_id, st) + else: + ed_video = len(input_tokens) + 1 + if ed_image < ed_video: + t, h, w = ( + image_grid_thw[image_index][0], + image_grid_thw[image_index][1], + image_grid_thw[image_index][2], + ) + image_index += 1 + remain_images -= 1 + ed = ed_image + else: + t, h, w = ( + video_grid_thw[video_index][0], + video_grid_thw[video_index][1], + video_grid_thw[video_index][2], + ) + video_index += 1 + remain_videos -= 1 + ed = ed_video + llm_grid_t, llm_grid_h, llm_grid_w = ( + t.item(), + h.item() // spatial_merge_size, + w.item() // spatial_merge_size, + ) + text_len = ed - st + st_idx = ( + llm_pos_ids_list[-1].max() + 1 + if len(llm_pos_ids_list) > 0 + else 0 + ) + index = mx.arange(text_len).reshape(1, text_len) + index = mx.broadcast_to(index, (3, text_len)) + index = index + st_idx + llm_pos_ids_list.append(index) + t_index = mx.arange(llm_grid_t).reshape(llm_grid_t, 1) + t_index = mx.broadcast_to( + t_index, (llm_grid_t, llm_grid_h * llm_grid_w) + ) + t_index = t_index.flatten() + h_index = mx.arange(llm_grid_h).reshape(1, llm_grid_h, 1) + h_index = mx.broadcast_to( + h_index, (llm_grid_t, llm_grid_h, llm_grid_w) + ) + h_index = h_index.flatten() + w_index = mx.arange(llm_grid_w).reshape(1, 1, llm_grid_w) + w_index = mx.broadcast_to( + w_index, (llm_grid_t, llm_grid_h, llm_grid_w) + ) + w_index = w_index.flatten() + llm_pos_ids_list.append( + mx.stack([t_index, h_index, w_index]) + text_len + st_idx + ) + st = ed + llm_grid_t * llm_grid_h * llm_grid_w + if st < len(input_tokens): + st_idx = ( + llm_pos_ids_list[-1].max() + 1 + if len(llm_pos_ids_list) > 0 + else 0 + ) + text_len = len(input_tokens) - st + t_index = mx.arange(text_len).reshape(1, text_len) + t_index = mx.broadcast_to(t_index, (3, text_len)) + llm_pos_ids_list.append(t_index + st_idx) + llm_positions = mx.concatenate(llm_pos_ids_list, axis=1).reshape(3, -1) + mask = mx.array(attention_mask[i] == 1) + expanded_mask = mx.expand_dims(mask, axis=0) + expanded_mask = mx.broadcast_to(expanded_mask, (3, 1, mask.shape[0])) + expanded_positions = mx.expand_dims(llm_positions, axis=1) + new_positions = mx.where( + expanded_mask, expanded_positions, position_ids[:, i : i + 1, :] + ) + updated_position_ids = mx.concatenate( + [ + position_ids[:, :i, :], + new_positions, + position_ids[:, i + 1 :, :], + ], + axis=1, + ) + position_ids = updated_position_ids + mrope_position_deltas.append( + llm_positions.max() + 1 - len(total_input_ids[i]) + ) + mrope_position_deltas = mx.array(mrope_position_deltas)[0] + return position_ids, mrope_position_deltas + else: + if attention_mask is not None: + position_ids = mx.cumsum(attention_mask.astype(mx.int64), axis=-1) - 1 + position_ids = mx.where( + attention_mask == 0, mx.ones_like(position_ids), position_ids + ) + position_ids = mx.expand_dims(position_ids[0], axis=0) + position_ids = mx.tile(position_ids, (3, 1, 1)) + max_position_ids = position_ids.max(0, keepdims=False)[0].max( + -1, keepdims=True + )[0] + mrope_position_deltas = max_position_ids + 1 - attention_mask.shape[-1] + else: + position_ids = mx.arange(input_ids.shape[1]).reshape(1, -1) + position_ids = mx.broadcast_to( + position_ids, (3, input_ids.shape[0], input_ids.shape[1]) + ) + mrope_position_deltas = mx.zeros( + [input_ids.shape[0], 1], + dtype=input_ids.dtype, + ) + return position_ids, mrope_position_deltas + + +class CustomQwen3VLModel(qwen_mod.Model): + + def __init__(self, config: qwen_mod.ModelConfig): + nn.Module.__init__(config) + self.enable_pruning = False + self.config = config + self.vision_tower = qwen_mod.VisionModel(config.vision_config) + # Replace upstream LanguageModel with fixed version + self.language_model = CustomLanguageModel(config.text_config, config) + + def get_input_embeddings( + self, + input_ids: Optional[mx.array] = None, + pixel_values: Optional[mx.array] = None, + **kwargs, + ): + image_grid_thw = kwargs.get("image_grid_thw", None) + video_grid_thw = kwargs.get("video_grid_thw", None) + mask = kwargs.get("mask", None) + grid_thw = image_grid_thw if image_grid_thw is not None else video_grid_thw + + if pixel_values is None: + # Reset position state for text-only generation + self.language_model._position_ids = None + self.language_model._rope_deltas = None + return InputEmbeddingsFeatures( + inputs_embeds=self.language_model.model.embed_tokens(input_ids) + ) + + # 获取缓存的特征 + buf_vis_features = kwargs.get("buf_vis_features", None) + buf_vis_stack_features = kwargs.get("buf_vis_stack_features", None) + + dtype = self.vision_tower.patch_embed.proj.weight.dtype + pixel_values = pixel_values.astype(dtype) + + # Get the input embeddings from the language model + inputs_embeds = self.language_model.model.embed_tokens(input_ids) + + # Get the ouptut hidden states from the vision model + hidden_states, deepstack_image_embeds = self.vision_tower( + pixel_values, grid_thw + ) + if buf_vis_features is not None and buf_vis_stack_features is not None: + # 检查 buf_vis_stack_features 是否是包含3个元素的列表 + assert isinstance(buf_vis_features, list) and isinstance( + buf_vis_stack_features, list + ) + if len(buf_vis_stack_features) == 3 and len(buf_vis_features) > 0: + + # 1. concat hidden_states (shape: (1, seq_len, d)) + # buf_vis_features 应该也是 (1, seq_len', d) + axis = 0 + assert hidden_states.ndim == buf_vis_features[0].ndim + if hidden_states.ndim == 3: + axis = 1 + hidden_states = mx.concatenate( + [buf_vis_features[0], hidden_states], axis=axis + ) + + # 就地更新 buf_vis_features + buf_vis_features[0] = hidden_states + + # 2. concat deepstack_image_embeds (列表，每个元素 shape: (seq_len, d)) + # 对应位置的元素进行concat + for i in range(3): + if i < len(deepstack_image_embeds): + # 在 axis=0 (seq_len维度) 上concat + axis = 0 if deepstack_image_embeds[i].ndim == 2 else 1 + concatenated = mx.concatenate( + [buf_vis_stack_features[i], deepstack_image_embeds[i]], + axis=axis, + ) + # 就地覆盖 buf_vis_stack_features + buf_vis_stack_features[i] = concatenated + + # 使用concat后的特征 + deepstack_image_embeds = buf_vis_stack_features + else: + buf_vis_features.append(hidden_states[:]) # 或者 hidden_states.copy() + for embed in deepstack_image_embeds: + buf_vis_stack_features.append(embed[:]) # 或者 embed.copy() + visual_pos_masks = None + deepstack_visual_embeds = None + + # Insert special image tokens in the input_ids + inputs_embeds, image_mask = self.merge_input_ids_with_image_features( + hidden_states, + inputs_embeds, + input_ids, + self.config.image_token_index, + self.config.video_token_index, + ) + image_mask = image_mask[..., 0] + visual_pos_masks = image_mask + mx.eval(deepstack_image_embeds) + deepstack_visual_embeds = deepstack_image_embeds + + # Pre-calculate position_ids for chunked prefill + if image_grid_thw is not None or video_grid_thw is not None: + position_ids, rope_deltas = self.language_model.get_rope_index( + input_ids, image_grid_thw, video_grid_thw, mask + ) + self.language_model._position_ids = position_ids + self.language_model._rope_deltas = rope_deltas + + return InputEmbeddingsFeatures( + inputs_embeds=inputs_embeds, + visual_pos_masks=visual_pos_masks, + deepstack_visual_embeds=deepstack_visual_embeds, + ) + + def __call__( + self, + input_ids: mx.array, + pixel_values: Optional[mx.array] = None, + mask: Optional[mx.array] = None, + cache=None, + **kwargs, + ): + input_embeddings_features = self.get_input_embeddings( + input_ids, pixel_values, **kwargs + ) + kwargs.update( + { + "pixel_values": pixel_values, + **input_embeddings_features.to_dict(), + } + ) + + logits = self.language_model(input_ids, mask=mask, cache=cache, **kwargs) + return logits + + +qwen_mod.Model = CustomQwen3VLModel + + +def custom_generate_step( + input_ids: mx.array, + model: nn.Module, + pixel_values, + mask, + *, + max_tokens: int = 256, + temperature: float = 0.0, + repetition_penalty: Optional[float] = None, + repetition_context_size: Optional[int] = 20, + top_p: float = 1.0, + logit_bias: Optional[Dict[int, float]] = None, + prompt_cache: Optional[List[Any]] = None, + max_kv_size: Optional[int] = None, + kv_bits: Optional[int] = None, + kv_group_size: int = 64, + quantized_kv_start: int = 0, + sampler: Optional[Callable[[mx.array], mx.array]] = None, + logits_processors: Optional[List[Callable[[mx.array, mx.array], mx.array]]] = None, + prefill_step_size: Optional[int] = 2048, + **kwargs, +) -> Generator[Tuple[mx.array, mx.array], None, None]: + """ + A generator producing token ids based on the given prompt from the model. + + Args: + input_ids (mx.array): The input prompt token ids. + model (nn.Module): The model to use for generation. + pixel_values: The pixel values for vision models (optional). + mask: The attention mask (optional). + max_tokens (int): Maximum number of tokens to generate. Default: ``256``. + temperature (float): The temperature for sampling, if 0 the argmax is used. + Default: ``0``. + repetition_penalty (float, optional): The penalty factor for repeating + tokens. + repetition_context_size (int, optional): The number of tokens to + consider for repetition penalty. Default: ``20``. + top_p (float, optional): Nucleus sampling, higher means model considers + more less likely words. + logit_bias (dictionary, optional): Additive logit bias. + prompt_cache (list, optional): Pre-existing KV cache for the prompt. + max_kv_size (int, optional): Maximum KV cache size. + kv_bits (int, optional): Number of bits for KV cache quantization. + kv_group_size (int): Group size for KV cache quantization. Default: ``64``. + quantized_kv_start (int): Start index for quantized KV cache. Default: ``0``. + sampler (Callable[mx.array, mx.array], optional): A sampler for sampling a + token from a vector of log probabilities. Default: ``None``. + logits_processors (List[Callable[[mx.array, mx.array], mx.array]], optional): + A list of functions that take tokens and logits and return the processed + logits. Default: ``None``. + prefill_step_size (int): Number of tokens to process per prefill step. + Chunked prefill processes prompts in smaller chunks to reduce peak + memory usage. Default: ``2048``. + + Yields: + Generator[Tuple[mx.array, mx.array], None, None]: A generator producing + one token and a vector of log probabilities. + """ + + quantize_cache_fn = functools.partial( + maybe_quantize_kv_cache, + quantized_kv_start=quantized_kv_start, + kv_group_size=kv_group_size, + kv_bits=kv_bits, + ) + + if sampler is None: + sampler = make_sampler(temperature, top_p) + + processors = make_logits_processors( + logit_bias, repetition_penalty, repetition_context_size + ) + if logits_processors is not None: + processors.extend(logits_processors) + + y = input_ids + tokens = mx.array([], dtype=input_ids.dtype) + + thinking_budget_criteria = kwargs.pop("thinking_budget_criteria", None) + + # Create the KV cache for generation + if prompt_cache is None: + prompt_cache = cache.make_prompt_cache( + model.language_model, + max_kv_size=max_kv_size, + ) + + def _step(y, inputs_embeds=None): + nonlocal tokens, kwargs + + with mx.stream(mx.gpu): + if "decoder_input_ids" in kwargs: + outputs = model.language_model( + cache=prompt_cache, + **kwargs, + ) + else: + outputs = model.language_model( + y, + inputs_embeds=inputs_embeds, + cache=prompt_cache, + **kwargs, + ) + + logits = outputs.logits[:, -1, :] + + if len(processors) > 0 and len(y) > 0: + tokens = mx.concat([tokens, y.flatten()]) + + for processor in processors: + logits = processor(tokens, logits) + + quantize_cache_fn(prompt_cache) + + logprobs = logits - mx.logsumexp(logits) + y = sampler(logprobs) + + if outputs.cross_attention_states is not None: + kwargs = {"cross_attention_states": outputs.cross_attention_states} + elif outputs.encoder_outputs is not None: + kwargs = {"encoder_outputs": outputs.encoder_outputs} + else: + kwargs = {} + + return y, logprobs.squeeze(0) + + with mx.stream(mx.gpu): + + # Get input embeddings (handles both multimodal and text-only) + embedding_output = model.get_input_embeddings( + input_ids, pixel_values, mask=mask, **kwargs + ) + inputs_embeds = embedding_output.inputs_embeds + + kwargs.update( + { + k: v + for k, v in embedding_output.to_dict().items() + if k != "inputs_embeds" and v is not None + } + ) + if prefill_step_size is not None and inputs_embeds.shape[1] > prefill_step_size: + # Pre-compute RoPE position_ids with FULL input_ids before chunking. + # Without this, get_rope_index sees only the first chunk and computes + # wrong positions when vision tokens span across chunk boundaries. + image_grid_thw = kwargs.get("image_grid_thw", None) + video_grid_thw = kwargs.get("video_grid_thw", None) + position_ids, rope_deltas = model.language_model.get_rope_index( + input_ids, image_grid_thw, video_grid_thw, None + ) + model.language_model._position_ids = position_ids + model.language_model._rope_deltas = rope_deltas + + # Remove pixel_values/image_grid_thw from kwargs to prevent + # language_model.__call__ from resetting _position_ids/_rope_deltas + # (it clears them when pixel_values is not None). + # Vision encoding is already done in get_input_embeddings above. + kwargs.pop("pixel_values", None) + kwargs.pop("image_grid_thw", None) + kwargs.pop("video_grid_thw", None) + + # Chunked prefill with embeddings + # Must slice visual_pos_masks and deepstack per-chunk to match h's seq_len + full_vpm = kwargs.pop("visual_pos_masks", None) + full_dse = kwargs.pop("deepstack_visual_embeds", None) + total_tokens = inputs_embeds.shape[1] + chunk_offset = 0 + with tqdm(total=total_tokens, desc="Prefill", unit="tok") as pbar: + while inputs_embeds.shape[1] > 1: + n_to_process = min(prefill_step_size, inputs_embeds.shape[1] - 1) + # Slice visual_pos_masks for this chunk + chunk_vpm = None + chunk_dse = None + if full_vpm is not None: + chunk_vpm = full_vpm[ + :, chunk_offset : chunk_offset + n_to_process + ] + if full_dse is not None: + # deepstack embeds need to be sliced to match image tokens in this chunk + # Count image token positions in this chunk from full_vpm + if chunk_vpm is not None: + + chunk_mask = np.array(chunk_vpm[0]) + n_img_in_chunk = int(chunk_mask.sum()) + if n_img_in_chunk > 0: + # Find the global image token index range for this chunk + full_mask = np.array(full_vpm[0]) + img_before = int(full_mask[:chunk_offset].sum()) + chunk_dse = [ + e[img_before : img_before + n_img_in_chunk] + for e in full_dse + ] + else: + chunk_dse = None + model.language_model( + inputs=input_ids[:, :n_to_process], + inputs_embeds=inputs_embeds[:, :n_to_process], + cache=prompt_cache, + visual_pos_masks=chunk_vpm, + deepstack_visual_embeds=chunk_dse, + **kwargs, + ) + quantize_cache_fn(prompt_cache) + mx.eval([c.state for c in prompt_cache]) + chunk_offset += n_to_process + inputs_embeds = inputs_embeds[:, n_to_process:] + input_ids = input_ids[:, n_to_process:] + mx.clear_cache() + pbar.update(n_to_process) + + # Last token goes through _step; need remaining vpm/dse + if full_vpm is not None: + kwargs["visual_pos_masks"] = full_vpm[:, chunk_offset:] + if full_dse is not None: + full_mask = np.array(full_vpm[0]) + img_before = int(full_mask[:chunk_offset].sum()) + n_img_remain = int(full_mask[chunk_offset:].sum()) + if n_img_remain > 0: + kwargs["deepstack_visual_embeds"] = [ + e[img_before : img_before + n_img_remain] for e in full_dse + ] + input_ids = input_ids[:, -1:] + + y, logprobs = _step(input_ids, inputs_embeds=inputs_embeds) + + mx.async_eval(y) + + n = 0 + while True: + if n != max_tokens: + next_y, next_logprobs = _step(y[None]) + mx.async_eval(next_y) + if n == 0: + mx.eval(y) + if n == max_tokens: + break + + yield y.item(), logprobs + if n % 256 == 0: + mx.clear_cache() + + if thinking_budget_criteria is not None: + next_y = thinking_budget_criteria.apply_forced_token(next_y) + y, logprobs = next_y, next_logprobs + n += 1 + + +def custom_stream_generate( + model: nn.Module, + processor: PreTrainedTokenizer, + prompt: str, + image: Union[str, List[str]] = None, + audio: Union[str, List[str]] = None, + **kwargs, +) -> Union[str, Generator[str, None, None]]: + """ + A generator producing text based on the given prompt from the model. + + Args: + model (nn.Module): The model to use for generation. + processor (PreTrainedTokenizer): The tokenizer/processor. + prompt (str): The input prompt text. + image (Union[str, List[str]], optional): Image path(s) or URL(s). + audio (Union[str, List[str]], optional): Audio file path(s). + prefill_step_size (int, optional): Number of tokens to process per prefill + step. When set, enables chunked prefill which processes long prompts in + smaller chunks to reduce peak memory usage. + kwargs: Additional options passed to :func:`generate_step`. + See :func:`generate_step` for more details. + + Yields: + Generator[GenerationResult]: A generator producing GenerationResult objects + containing the generated text, tokens, and statistics. + """ + tokenizer = processor.tokenizer if hasattr(processor, "tokenizer") else processor + + # Set up thinking budget criteria if requested + thinking_budget = kwargs.pop("thinking_budget", None) + thinking_end_token = kwargs.pop("thinking_end_token", "") + thinking_start_token = kwargs.pop("thinking_start_token", "") + enable_thinking = kwargs.pop("enable_thinking", False) + + # Skip special tokens + skip_special_tokens = kwargs.pop("skip_special_tokens", False) + skip_special_token_ids = ( + set(tokenizer.all_special_ids) + if skip_special_tokens and hasattr(tokenizer, "all_special_ids") + else [] + ) + + add_special_tokens = ( + not hasattr(processor, "chat_template") + if model.config.model_type in ["gemma3", "gemma3n", "gemma4"] + else True + ) + + resize_shape = normalize_resize_shape(kwargs.pop("resize_shape", None)) + image_token_index = getattr(model.config, "image_token_index", None) + + if kwargs.get("input_ids", None) is not None: + input_ids = kwargs.pop("input_ids") + pixel_values = kwargs.pop("pixel_values", None) + mask = kwargs.pop("mask", None) + else: + inputs = prepare_inputs( + processor, + images=image, + audio=audio, + prompts=prompt, + image_token_index=image_token_index, + resize_shape=resize_shape, + add_special_tokens=add_special_tokens, + **kwargs, + ) + input_ids = inputs.get("input_ids", None) + pixel_values = inputs.get("pixel_values", None) + mask = inputs.get("attention_mask", None) + data_kwargs = { + k: v + for k, v in inputs.items() + if k not in ["input_ids", "pixel_values", "attention_mask"] + } + kwargs.update(data_kwargs) + + if thinking_budget is not None: + thinking_start_token_id = tokenizer.encode( + thinking_start_token, add_special_tokens=False + )[-1] + enable_thinking = enable_thinking and ( + thinking_start_token_id in input_ids.flatten().tolist() + ) + tokenizer.thinking_budget_criteria = ThinkingBudgetCriteria( + tokenizer=tokenizer, + thinking_budget=thinking_budget, + thinking_end_token=thinking_end_token, + thinking_start_token=thinking_start_token, + enable_thinking=enable_thinking, + ) + kwargs["thinking_budget_criteria"] = tokenizer.thinking_budget_criteria + else: + tokenizer.thinking_budget_criteria = None + + if input_ids.size > MAX_PROMPT_TOKENS: + # yield一个错误结果并直接返回 + yield CustomGenerationResult( + text=f"Error: Input token length ({input_ids.size}) exceeds maximum allowed ({MAX_PROMPT_TOKENS})", + token=None, + logprobs=None, + prompt_tokens=input_ids.size, + generation_tokens=0, + total_tokens=input_ids.size, + prompt_tps=0.0, + generation_tps=0.0, + peak_memory=0.0, + code=ErrorCode.PREFILL_TOO_LONG, # 可以添加一个error标志 + ) + return # 直接返回，不进行生成 + + with wired_limit(model): + detokenizer = processor.detokenizer + detokenizer.reset() + thinking_criteria = getattr(tokenizer, "thinking_budget_criteria", None) + gen = custom_generate_step(input_ids, model, pixel_values, mask, **kwargs) + tic = time.perf_counter() + + for n, (token, logprobs) in enumerate(gen): + if n == 0: + prompt_time = time.perf_counter() - tic + prompt_tps = input_ids.size / prompt_time + tic = time.perf_counter() + + # Check thinking budget and force token if needed + if thinking_criteria is not None: + thinking_criteria(token) + + # Stop generation if the token is in the eos_token_ids + if tokenizer.stopping_criteria(token): + break + + detokenizer.add_token(token, skip_special_token_ids=skip_special_token_ids) + + # Yield the last segment if streaming + yield CustomGenerationResult( + text=detokenizer.last_segment, + token=token, + logprobs=logprobs, + prompt_tokens=input_ids.size, + generation_tokens=n + 1, + total_tokens=input_ids.size + n + 1, + prompt_tps=prompt_tps, + generation_tps=(n + 1) / (time.perf_counter() - tic), + peak_memory=mx.get_peak_memory() / 1e9, + ) + + detokenizer.finalize() + yield CustomGenerationResult( + text=detokenizer.last_segment, + token=token, + logprobs=logprobs, + prompt_tokens=input_ids.size, + generation_tokens=n + 1, + total_tokens=input_ids.size + n + 1, + prompt_tps=prompt_tps, + generation_tps=(n + 1) / (time.perf_counter() - tic), + peak_memory=mx.get_peak_memory() / 1e9, + ) + # Cleanup after generation + mx.clear_cache() + + +def custom_generate( + model: nn.Module, + processor: PreTrainedTokenizer, + prompt: str, + image: Union[str, List[str]] = None, + audio: Union[str, List[str]] = None, + **kwargs, +): + """ + Generate text from the model. + + Args: + model (nn.Module): The language model. + tokenizer (PreTrainedTokenizer): The tokenizer. + prompt (str): The string prompt. + temperature (float): The temperature for sampling (default 0). + max_tokens (int): The maximum number of tokens (default 100). + verbose (bool): If ``True``, print tokens and timing information + (default ``False``). + formatter (Optional[Callable]): A function which takes a token and a + probability and displays it. + repetition_penalty (float, optional): The penalty factor for repeating tokens. + repetition_context_size (int, optional): The number of tokens to consider for repetition penalty. + """ + + text = "" + last_response = None + + eos_tokens = kwargs.get("eos_tokens", None) + stopping_criteria = kwargs.get("stopping_criteria", None) + + # Get the tokenizer + tokenizer = processor.tokenizer if hasattr(processor, "tokenizer") else processor + + # Add custom EOS tokens to the stopping criteria + if eos_tokens is not None: + tokenizer.stopping_criteria.add_eos_token_ids(eos_tokens) + + # Use custom stopping criteria + elif stopping_criteria is not None: + if isinstance(stopping_criteria, StoppingCriteria) or callable( + stopping_criteria + ): + tokenizer.stopping_criteria = stopping_criteria + else: + raise ValueError( + "stopping_criteria must be an instance of StoppingCriteria or a callable" + ) + else: + tokenizer.stopping_criteria.reset(model.config.eos_token_id) + + for response in custom_stream_generate( + model, processor, prompt, image, audio, **kwargs + ): + text += response.text + last_response = response + + if len(text) == 0: + return CustomGenerationResult( + text=text, + token=None, + logprobs=None, + prompt_tokens=0, + generation_tokens=0, + total_tokens=0, + prompt_tps=0.0, + generation_tps=0.0, + peak_memory=mx.get_peak_memory() / 1e9, + ) + + return CustomGenerationResult( + text=text, + token=last_response.token, + logprobs=last_response.logprobs, + prompt_tokens=last_response.prompt_tokens, + generation_tokens=last_response.generation_tokens, + total_tokens=last_response.total_tokens, + prompt_tps=last_response.prompt_tps, + generation_tps=last_response.generation_tps, + peak_memory=last_response.peak_memory, + ) diff --git a/vlm_service/server.py b/vlm_service/server.py new file mode 100644 index 0000000..75ba292 --- /dev/null +++ b/vlm_service/server.py @@ -0,0 +1,732 @@ +import asyncio +import uuid +import json +import base64 +from fastapi import FastAPI, HTTPException +from fastapi.responses import StreamingResponse +from pydantic import BaseModel, Field +from contextlib import asynccontextmanager +from queue import Queue +import uvicorn +from typing import List, Dict, Optional +import torch +import argparse +import time +import threading +from PIL import Image +from .core_infer import HMInference, ErrorCode +import logging +from io import BytesIO +from .config import load_config, get_config +from pathlib import Path + +# ============= Pydantic Models for OpenAI Compatible API ============= +logging.basicConfig( + level=logging.INFO, + format="%(levelname)s %(asctime)s %(filename)s: %(lineno)d] %(message)s", + datefmt="%Y-%m-%d %H:%M:%S", +) +logger = logging.getLogger(__name__) + + +def base64_to_pil(base64_string): + """base64 转 PIL.Image""" + image_data = base64.b64decode(base64_string) + image = Image.open(BytesIO(image_data)) + return image + + +class Message(BaseModel): + role: str + content: str | List[Dict] + + +class ChatCompletionRequest(BaseModel): + model: str = "qwen2.5-vl" + messages: List[Message] + images: Optional[List] = None + temperature: Optional[float] = None + top_p: Optional[float] = None + top_k: Optional[int] = None + repetition_penalty: Optional[float] = None + max_tokens: int = 2048 + stream: bool = False + request_id: Optional[str] = None + + +class Usage(BaseModel): + prompt_tokens: int = 0 + completion_tokens: int = 0 + total_tokens: int = 0 + + +class ChatCompletionResponseChoice(BaseModel): + index: int + message: Message + finish_reason: Optional[str] = None + + +class ChatCompletionResponse(BaseModel): + id: str + object: str = "chat.completion" + created: int + model: str + choices: List[ChatCompletionResponseChoice] + usage: Usage + prefill_time: float = 0.0 + decode_tps: float = 0.0 + + +class ChatCompletionStreamChoice(BaseModel): + index: int + delta: Dict + finish_reason: Optional[str] = None + + +class ChatCompletionStreamResponse(BaseModel): + id: str + object: str = "chat.completion.chunk" + created: int + model: str + choices: List[ChatCompletionStreamChoice] + + +class InferenceRequest: + """推理请求""" + + def __init__( + self, + request_id: str, + messages: List[Message], + images: List[Image.Image], + params: Dict, + ): + self.request_id = request_id + self.messages = messages + self.images = images + self.params = params + self.result_queue = Queue() + self.stream = params.get("stream", False) + + +class RequestQueueManager: + """请求队列管理器（单例模式下的请求排队）""" + + def __init__(self): + self.queue = Queue() + + def add_request(self, request: InferenceRequest): + """添加请求到队列""" + self.queue.put(request) + + def get_next_request(self) -> Optional[InferenceRequest]: + """获取下一个请求""" + if not self.queue.empty(): + return self.queue.get(timeout=1) + return None + + def is_empty(self) -> bool: + return self.queue.empty() + + def size(self) -> int: + """队列大小""" + return self.queue.qsize() + + +# ============= Request Context Manager ============= + + +class RequestContext: + """管理单个请求的上下文（多轮对话的图像特征缓存）""" + + def __init__(self, request_id: str): + self.request_id = request_id + self.image_features_buffer: List[torch.Tensor] = [] + self.image_stack_feature_buffer: List[torch.Tensor] = [] + self.created_at = time.time() + self.last_accessed = self.created_at + + def get_image_features_buffer(self) -> List[torch.Tensor]: + """获取图像特征缓存""" + self.last_accessed = time.time() + return self.image_features_buffer, self.image_stack_feature_buffer + + +class RequestContextManager: + """管理所有请求的上下文""" + + def __init__(self, ttl: int = 3600): + self.contexts: Dict[str, RequestContext] = {} + self.ttl = ttl # 上下文存活时间（秒） + self.lock = threading.Lock() + + def get_or_create_context(self, request_id: str) -> RequestContext: + """获取或创建请求上下文""" + with self.lock: + if request_id not in self.contexts: + self.contexts[request_id] = RequestContext(request_id) + return self.contexts[request_id] + + def cleanup_expired_contexts(self): + """清理过期的上下文""" + with self.lock: + current_time = time.time() + expired_keys = [ + key + for key, ctx in self.contexts.items() + if current_time - ctx.last_accessed > self.ttl + ] + for key in expired_keys: + del self.contexts[key] + if expired_keys: + logger.info(f"Cleaned up {len(expired_keys)} expired contexts") + + +class InferenceService: + """推理服务:管理单例HMInference和请求队列""" + + _instance: Optional["InferenceService"] = None + _lock = threading.Lock() + + def __new__(cls, *args, **kwargs): + if cls._instance is None: + with cls._lock: + if cls._instance is None: + cls._instance = super().__new__(cls) + return cls._instance + + def __init__(self, cfg): + if hasattr(self, "_initialized") and self._initialized: + return + + self._cfg = cfg + + # 请求上下文管理器 + self.context_manager = RequestContextManager(cfg.server.ttl) + + # 请求队列管理器 + self.queue_manager = RequestQueueManager() + + # 模型将在worker thread中加载（MLX stream亲和性） + self.inference_engine = None + self._model_ready = threading.Event() + + # 启动处理线程 + self.worker_thread = threading.Thread( + target=self._process_requests, daemon=True + ) + self.worker_thread.start() + + # 启动清理线程 + self.cleanup_thread = threading.Thread(target=self._cleanup_loop, daemon=True) + self.cleanup_thread.start() + + self._initialized = True + logger.info("InferenceService initialized") + + def _cleanup_loop(self): + """定期清理过期上下文""" + while True: + time.sleep(300) # 每5分钟清理一次 + self.context_manager.cleanup_expired_contexts() + + def _process_requests(self): + """处理请求队列（单线程串行处理，模型在此线程加载）""" + # Load model in worker thread (MLX stream affinity) + cfg = self._cfg + self.inference_engine = HMInference( + cfg.model.model_name_or_path, + cfg.sampling.temperature, + cfg.sampling.top_k, + cfg.sampling.top_p, + cfg.sampling.repetition_penalty, + cfg.sampling.max_new_tokens, + w8a8=cfg.w8a8.mode, + ) + self._model_ready.set() + logger.info("Model loaded in worker thread") + while True: + request = self.queue_manager.get_next_request() + if request is None: + time.sleep(0.01) # 避免CPU空转 + continue + + try: + if request.stream: + self._process_stream_request(request) + else: + self._process_non_stream_request(request) + except Exception as e: + import traceback + + tb = traceback.format_exc() + logger.error( + "Error processing request %s: %s\n%s" % (request.request_id, e, tb) + ) + request.result_queue.put({"status": str(e)}) + + def _process_non_stream_request(self, request: InferenceRequest): + """处理非流式请求""" + # 获取请求上下文 + ctx = self.context_manager.get_or_create_context(request.request_id) + buf_vis_feats, buf_vis_stack_feats = ctx.get_image_features_buffer() + + # 执行推理 + code, generated_text, timing = self.inference_engine.complete( + request.messages, + request.images, + buf_vis_feats, + buf_vis_stack_feats, + **request.params, + ) + + # 返回结果 + result = { + "status": code, + "text": generated_text if code == ErrorCode.SUCCESS else "", + "prefill_time": timing["prefill_time"], + "decode_tps": timing["decode_tps"], + } + if code != ErrorCode.SUCCESS: + result["error"] = code + request.result_queue.put(result) + + def _process_stream_request(self, request: InferenceRequest): + """处理流式请求 - 使用 complete_stream""" + ctx = self.context_manager.get_or_create_context(request.request_id) + buf_vis_feats, buf_vis_stack_feats = ctx.get_image_features_buffer() + + try: + stream_gen = self.inference_engine.complete_stream( + request.messages, + request.images, + buf_vis_feats, + buf_vis_stack_feats, + **request.params, + ) + + for code, text, timing in stream_gen: + if code != ErrorCode.SUCCESS: + request.result_queue.put( + { + "status": code, + "done": True, + "prefill_time": timing["prefill_time"], + "decode_tps": timing["decode_tps"], + "error": code, + } + ) + return + else: + request.result_queue.put( + { + "text": text, + "done": False, + } + ) + request.result_queue.put( + { + "text": "", + "done": True, + "prefill_time": timing["prefill_time"], + "decode_tps": timing["decode_tps"], + } + ) + except Exception as e: + logger.error(f"Stream error: {e}") + request.result_queue.put({"error": str(e), "done": True}) + + async def submit_request(self, request: InferenceRequest) -> Queue: + """提交请求到队列""" + self.queue_manager.add_request(request) + return request.result_queue + + +# 全局变量 +inference_service: Optional[InferenceService] = None +_global_config = None + + +def init_config(config_path: str = "config.yaml"): + """初始化全局配置""" + global _global_config + if _global_config is None: + _global_config = load_config(config_path) + return _global_config + + +@asynccontextmanager +async def lifespan(app: FastAPI): + """应用生命周期管理""" + global inference_service + + # ✅ 安全获取配置（带默认值） + try: + cfg = get_config() + except RuntimeError: + # 如果配置还没加载，使用默认配置路径 + logger.warning("Config not loaded, using default config.yaml") + cfg = init_config("config.yaml") + + inference_service = InferenceService(cfg) + # Wait for model to load in worker thread (async-friendly) + while not inference_service._model_ready.is_set(): + await asyncio.sleep(0.5) + logger.info("Service started") + yield + logger.info("Service shutting down") + + +app = FastAPI(title="Mininglamp OpenAI Compatible API", lifespan=lifespan) + + +def parse_openai_messages( + messages: List[Message], images: Optional[List[Image.Image]] = None +) -> tuple[List[Dict], List[Image.Image]]: + """ + 解析 OpenAI 格式的消息，提取文本和图像 + 支持两种格式： + 1. content 是字符串 + 2. content 是列表，包含 text 和 image_url + """ + parsed_messages = [] + image_list = [] + if images: + image_list = images + for msg in messages: + role = msg.role + content = msg.content + if isinstance(content, str): + # 纯文本消息 + parsed_messages.append({"role": role, "content": content}) + elif isinstance(content, list): + # 多模态消息 + text_parts = [] + for item in content: + if item.get("type") == "text": + text_parts.append(item["text"]) + elif item.get("type") == "image_url": + # 处理图像 + image_url = item["image_url"]["url"] + if image_url.startswith("data:image"): + # base64 编码的图像 - strip data URI prefix + b64_data = ( + image_url.split(",", 1)[1] + if "," in image_url + else image_url + ) + img = base64_to_pil(b64_data) + image_list.append(img) + text_parts.append("") + elif image_url.startswith("http"): + # URL 图像（需要下载） + import requests + + response = requests.get(image_url) + img = Image.open(BytesIO(response.content)) + image_list.append(img) + text_parts.append("") + else: + # 本地文件路径 + img = Image.open(image_url) + image_list.append(img) + text_parts.append("") + + if text_parts: + parsed_messages.append({"role": role, "content": "".join(text_parts)}) + else: + parsed_messages.append({"role": role, "content": str(content)}) + + return parsed_messages, image_list + + +def merge_params_with_config(request: ChatCompletionRequest) -> Dict: + """ + 合并请求参数和配置默认值 + 请求参数优先级高于配置 + """ + cfg = get_config() + + params = { + "temperature": ( + request.temperature + if request.temperature is not None + else cfg.sampling.temperature + ), + "topp": request.top_p if request.top_p is not None else cfg.sampling.top_p, + "topk": request.top_k if request.top_k is not None else cfg.sampling.top_k, + "repetition_penalty": ( + request.repetition_penalty + if request.repetition_penalty is not None + else cfg.sampling.repetition_penalty + ), + "stream": request.stream, + } + + return params + + +@app.post("/v1/chat/completions") +async def chat_completions(request: ChatCompletionRequest): + """OpenAI 兼容的聊天补全接口""" + global inference_service + + if inference_service is None: + raise HTTPException(status_code=503, detail="Service not initialized") + + request_id = request.request_id or str(uuid.uuid4()) + + try: + # ✅ 处理 images 参数（如果有的话） + external_images = [] + if request.images: + for img_data in request.images: + # img_data 只能是字符串（base64 或 URL） + if not isinstance(img_data, str): + raise HTTPException( + status_code=400, + detail=f"Invalid image data type: {type(img_data)}. Must be base64 string or URL.", + ) + + # 处理不同的字符串格式 + if img_data.startswith("data:image"): + # data:image/jpeg;base64,xxx 格式 + img = base64_to_pil(img_data) + elif img_data.startswith("http://") or img_data.startswith("https://"): + # URL 格式 + import requests + + response = requests.get(img_data) + response.raise_for_status() + img = Image.open(BytesIO(response.content)) + else: + # 纯 base64 字符串（没有 data:image 前缀） + try: + img = base64_to_pil(img_data) + except Exception as e: + raise HTTPException( + status_code=400, detail=f"Failed to decode image: {str(e)}" + ) + + external_images.append(img) + + # ✅ 解析消息，传入 external_images + parsed_messages, images = parse_openai_messages( + request.messages, external_images + ) + + if not images: + raise HTTPException( + status_code=400, + detail="No images found in messages. This model requires images.", + ) + + params = merge_params_with_config(request) + logger.info(f"Request {request_id[:8]}: params={params}, images={len(images)}") + + inference_request = InferenceRequest( + request_id=request_id, + messages=parsed_messages, + images=images, + params=params, + ) + + result_queue = await inference_service.submit_request(inference_request) + + if request.stream: + return StreamingResponse( + stream_generator(result_queue, request_id, request.model), + media_type="text/event-stream", + ) + else: + result = await asyncio.get_event_loop().run_in_executor( + None, result_queue.get + ) + + if result.get("status") != ErrorCode.SUCCESS: + error_msg = result.get("error", "Unknown error") + raise HTTPException(status_code=500, detail=error_msg) + + response = ChatCompletionResponse( + id=request_id, + created=int(time.time()), + model=request.model, + choices=[ + ChatCompletionResponseChoice( + index=0, + message=Message(role="assistant", content=result["text"]), + finish_reason="stop", + ) + ], + usage=Usage( + prompt_tokens=0, + completion_tokens=len(result["text"]), + total_tokens=len(result["text"]), + ), + prefill_time=result.get("prefill_time", 0.0), + decode_tps=result.get("decode_tps", 0.0), + ) + + return response + + except HTTPException: + raise + except Exception as e: + logger.error(f"Error in chat_completions: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + except HTTPException: + raise + except Exception as e: + logger.error(f"Error in chat_completions: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +async def stream_generator(result_queue: Queue, request_id: str, model: str): + """流式响应生成器""" + try: + while True: + # 非阻塞获取结果 + result = await asyncio.get_event_loop().run_in_executor( + None, result_queue.get + ) + + if "error" in result: + # 错误情况 + chunk = ChatCompletionStreamResponse( + id=request_id, + created=int(time.time()), + model=model, + choices=[ + ChatCompletionStreamChoice( + index=0, + delta={"content": f"[ERROR]: {result['error']}"}, + finish_reason="error", + ) + ], + ) + yield f"data: {chunk.model_dump_json()}\n\n" + yield "data: [DONE]\n\n" + break + + if result.get("done", False): + # 结束 + chunk = ChatCompletionStreamResponse( + id=request_id, + created=int(time.time()), + model=model, + choices=[ + ChatCompletionStreamChoice( + index=0, delta={}, finish_reason="stop" + ) + ], + ) + yield f"data: {chunk.model_dump_json()}\n\n" + + # 如果有性能数据，额外发送一个包含性能信息的chunk + if "prefill_time" in result: + perf_chunk = { + "id": request_id, + "object": "chat.completion.chunk", + "created": int(time.time()), + "model": model, + "choices": [], + "performance": { + "prefill_time": result.get("prefill_time", 0.0), + "decode_tps": result.get("decode_tps", 0.0), + }, + } + yield f"data: {json.dumps(perf_chunk)}\n\n" + + yield "data: [DONE]\n\n" + break + + # 正常的token + text = result.get("text", "") + if text: + chunk = ChatCompletionStreamResponse( + id=request_id, + created=int(time.time()), + model=model, + choices=[ + ChatCompletionStreamChoice( + index=0, delta={"content": text}, finish_reason=None + ) + ], + ) + yield f"data: {chunk.model_dump_json()}\n\n" + + except Exception as e: + logger.error(f"Error in stream_generator: {e}") + error_chunk = { + "id": request_id, + "object": "chat.completion.chunk", + "created": int(time.time()), + "model": model, + "choices": [ + { + "index": 0, + "delta": {"content": f"[STREAM ERROR]: {str(e)}"}, + "finish_reason": "error", + } + ], + } + yield f"data: {json.dumps(error_chunk)}\n\n" + yield "data: [DONE]\n\n" + + +@app.get("/health") +async def health(): + """健康检查接口""" + return {"status": "ok"} + + +@app.get("/v1/models") +async def list_models(): + """列出可用模型""" + return { + "object": "list", + "data": [ + { + "id": "qwen3-vl", + "object": "model", + "created": int(time.time()), + "owned_by": "mininglamp", + "author": "tiandu.ws", + } + ], + } + + +@app.get("/v1/queue") +async def queue_status(): + """查询队列状态""" + global inference_service + + if inference_service is None: + raise HTTPException(status_code=503, detail="Service not initialized") + + queue_size = inference_service.queue_manager.size() + + return { + "queue_size": queue_size, + "estimated_wait_seconds": queue_size * 3, # 假设每个请求 3 秒 + "status": "idle" if queue_size == 0 else "busy", + } + + +if __name__ == "__main__": + + # 只保留配置文件路径参数 + cfg_path = str((Path(__file__).parent.parent) / Path("config/config.yaml")) + parser = argparse.ArgumentParser() + parser.add_argument( + "--config", type=str, default=cfg_path, help="Path to config file" + ) + args = parser.parse_args() + # 加载配置 + config = load_config(args.config) + + # 启动服务 + logger.info(f"Starting server on {config.server.host}:{config.server.port}") + uvicorn.run(app, host=config.server.host, port=config.server.port, log_level="info")