diff --git a/.gitignore b/.gitignore index 8dc9d7d0b..e4302254c 100644 --- a/.gitignore +++ b/.gitignore @@ -160,3 +160,6 @@ a.out.* AGENTS.local.md .pi/SYSTEM.md + +# Local model/rotation assets (large binaries) +assets/ diff --git a/BUILD-NOTES.md b/BUILD-NOTES.md new file mode 100644 index 000000000..1f5f5d795 --- /dev/null +++ b/BUILD-NOTES.md @@ -0,0 +1,174 @@ +# OSCAR-llamacpp — Build & Run Notes (RTX 5090 / Blackwell) + +This is a working copy of the OSCAR INT2-KV llama.cpp fork +(`FutureMLS-Lab/OSCAR`, branch `zhongzhu/llamacpp`), preserved at +`github.com/giveen/OSCAR-llamacpp`. + +OSCAR compresses only the **KV cache** to INT2 (`q2_0`) using a calibrated, +per-layer orthogonal rotation baked into the GGUF. Model **weights stay in full +precision** (e.g. Q4_K_M). This lets you run very long contexts (tested to the +model's full 256K) for almost no extra VRAM. + +--- + +## 1. Clone + +```bash +git clone --depth 1 --branch zhongzhu/llamacpp \ + https://github.com/FutureMLS-Lab/OSCAR.git OSCAR-llamacpp +cd OSCAR-llamacpp +``` + +## 2. Build for NVIDIA RTX 5090 (sm_120, CUDA 13.3) + +The system's CUDA compiler-id probe is broken (it tries `compute_52`, which +CUDA 13.3's `ptxas` no longer supports → `sm_52 is not defined`). The working +flag set below **lies about the compiler ID** (`-DCMAKE_CUDA_COMPILER_ID=NVIDIA` ++ explicit version) and forces `native` arch. It is derived from the proven +`/mnt/storage/llama-server/rebuild_llama.sh`. + +The critical flag for OSCAR GPU-side KV is **`-DGGML_CUDA_FA_ALL_QUANTS=ON`**: +it compiles Flash-Attention kernels for the `q2_0` KV quant type, which is what +lets the INT2 KV cache run on the GPU instead of forcing `-nkvo` (CPU KV). + +```bash +cd /mnt/storage/Projects/OSCAR-llamacpp +rm -rf build +export PATH="/usr/local/cuda/bin:$PATH" + +cmake -B build -G Ninja \ + -DCMAKE_C_COMPILER=gcc-15 \ + -DCMAKE_CXX_COMPILER=g++-15 \ + -DCMAKE_CXX_COMPILER_LAUNCHER=ccache \ + -DCMAKE_C_COMPILER_LAUNCHER=ccache \ + -DCMAKE_CUDA_COMPILER_LAUNCHER=ccache \ + -DCMAKE_CUDA_COMPILER=/usr/local/cuda/bin/nvcc \ + -DCMAKE_CUDA_FLAGS="-ccbin /usr/bin/g++-15 -isystem /usr/local/cuda/include" \ + -DCMAKE_CUDA_COMPILER_ID=NVIDIA \ + -DCMAKE_CUDA_COMPILER_VERSION=13.3 \ + -DCMAKE_CUDA_STANDARD_COMPUTED_DEFAULT=17 \ + -DCMAKE_CUDA_EXTENSIONS_COMPUTED_DEFAULT=ON \ + -DCUDAToolkit_ROOT=/usr/local/cuda \ + -DGGML_LTO=ON \ + -DGGML_CPU_KLEIDIAI=OFF \ + -DGGML_CUDA=ON \ + -DGGML_NATIVE=ON \ + -DGGML_CUDA_GRAPHS=ON \ + -DGGML_CUDA_FA_ALL_QUANTS=ON \ + -DCMAKE_BUILD_TYPE=Release \ + -DCMAKE_CUDA_ARCHITECTURES=native \ + -DCMAKE_LINK_DEPENDS_USE_LINKER=OFF \ + -DLLAMA_CURL=ON . + +cmake --build build --config Release -j"$(nproc)" +``` + +> NOTE: do **not** mix generators — the first attempt used Unix Makefiles; this +> script uses Ninja. Purge `build/` if you switch. + +Resulting binaries: `build/bin/llama-server`, `build/bin/llama`. + +## 3. Get a model + +Pre-baked OSCAR GGUF (weights Q4_K_M + rotation baked in), ~7.4 GB: + +- `https://huggingface.co/Zhongzhu/OSCAR-LLAMACPP-Gemma-4-12B-it-INT2-KV` + - `q4km-rot-kv/gemma-4-12b-it-rot-kv.gguf` + +Raw rotation matrices (for self-baking other models), head_dim 256, 48 layers: + +- `https://huggingface.co/Zhongzhu/OSCAR-RotationZoo` → `Gemma4-12B/` + - `k_rotation_qqt_r_h_pbr.pt`, `v_rotation_sst_r_h_pbr.pt` + +Self-bake a `-rot-kv` GGUF from a base GGUF + rotation dir: + +```bash +python3 oscar-rotation/export_rot_kv_gguf.py \ + --base /path/to/base-q4km.gguf \ + --rot-dir /path/to/Gemma4-12B \ + --out /path/to/base-q4km-rot-kv.gguf +``` + +## 4. Run + +### The original crash (and why the new build fixes it) + +With the *first* build (single-arch sm_120, no `GGML_CUDA_FA_ALL_QUANTS`), +running `-fa on` with the KV cache on the GPU aborted at context build: + +``` +ggml-backend.cpp:898: pre-allocated tensor (cache_k_l0 (view)) in a buffer +(CUDA0) that cannot run the operation (SET_ROWS) +``` + +`-fa off` does **not** help — the server refuses because +`V cache quantization requires flash_attn`. + +**Fix (this build):** `-DGGML_CUDA_FA_ALL_QUANTS=ON` compiles the +Flash-Attention kernels for the `q2_0` KV quant, so the INT2 KV cache can live +on the GPU. Drop `-nkvo` and run with GPU KV. Only fall back to `-nkvo` (CPU KV) +if the SET_ROWS crash reappears. + +### Recommended launch (GPU KV) + +```bash +cd /mnt/storage/Projects/OSCAR-llamacpp +G=/path/to/gemma-4-12b-it-rot-kv.gguf + +pkill -f llama-server 2>/dev/null; sleep 2 +CUDA_VISIBLE_DEVICES=0 \ +GGML_CUDA_DISABLE_GRAPHS=1 \ +LLAMA_KV_FUSED_FA=1 \ +LLAMA_KV_NO_HADAMARD=1 \ +LLAMA_KV_CLIP_RATIO=0.96 \ +LLAMA_KV_HP_SINK=512 \ +LLAMA_KV_HP_RECENT=2048 \ +./build/bin/llama-server \ + -m "$G" \ + -fa on \ + -ngl 99 \ + -c 262144 \ + --cache-type-k q2_0 --cache-type-v q2_0 \ + --host 127.0.0.1 --port 8080 +``` + +### Fallback launch (CPU KV, if GPU KV crashes) + +```bash +# same as above but append: + -nkvo +``` + +### Notes on the flags + +| Flag | Why | +|------|-----| +| `-fa on` | **Required.** `q2_0` KV quantization needs flash attention. | +| `--cache-type-k/v q2_0` | The OSCAR INT2 KV cache. | +| `-nkvo` | **Fallback only.** KV on CPU sidesteps the Blackwell `SET_ROWS` CUDA-buffer crash. Avoid with the FA_ALL_QUANTS build. | +| `LLAMA_KV_FUSED_FA=1` | Fused mixed-precision Flash-Attention kernel (OSCAR path). | +| `LLAMA_KV_NO_HADAMARD=1` | OSCAR stores plain INT2 KV (no in-quant Hadamard). | +| `LLAMA_KV_CLIP_RATIO=0.96` | Per-layer clip ratio baked during calibration. | +| `LLAMA_KV_HP_SINK=512` / `LLAMA_KV_HP_RECENT=2048` | Hybrid-prefix KV: 512 sink + 2048 recent tokens in higher precision. | +| `-c 262144` | Full native Gemma-4-12B context (model trains at 256K). | + +`GGML_CUDA_DISABLE_GRAPHS=1` is kept to avoid runtime graph-replay issues on +Blackwell (it does not affect the SET_ROWS reserve path). + +## 5. Verified behavior (RTX 5090, 32 GB) + +- OSCAR KV engaged: `KV cache dtype: K=q2_0 V=q2_0 (n_embd_k_gqa=512, kv_size=262144)`. +- Full 256K context loads with **4 slots**, no truncation warning. +- VRAM at 256K context: **~10.4 GB** (weights Q4_K_M + 2-bit KV). +- Decode: **~45 tok/s** (CPU-KV path; GPU still does the matmuls). + - *With the FA_ALL_QUANTS GPU-KV build, expect higher decode once `-nkvo` is dropped — verify after launch.* +- Health: `GET /health` → `{"status":"ok"}`. +- Smoke generation: prompt "The capital of France is" → "**Paris**." (coherent). + +## 6. Known limitation / future work + +With the old build, `-nkvo` put the KV cache on CPU, capping long-context decode +throughput. The new `GGML_CUDA_FA_ALL_QUANTS=ON` build targets GPU KV directly. +If the SET_ROWS crash persists on GPU, the deeper fix is the scheduler's KV-buffer +backend selection in `src/llama-kv-cache.cpp` (CUDA buffer-type placement) so +`SET_ROWS` into the q2_0 cache runs on the CUDA buffer. diff --git a/ggml/src/ggml-cuda/fattn-common.cuh b/ggml/src/ggml-cuda/fattn-common.cuh index debcb6e54..e87776b1e 100644 --- a/ggml/src/ggml-cuda/fattn-common.cuh +++ b/ggml/src/ggml-cuda/fattn-common.cuh @@ -136,6 +136,48 @@ static __device__ __forceinline__ float vec_dot_fattn_vec_KQ_q4_0( return sum; } +template +static __device__ __forceinline__ float vec_dot_fattn_vec_KQ_q2_0( + const char * __restrict__ K_c, const void * __restrict__ Q_v, const int * __restrict__ Q_q8, const void * __restrict__ Q_ds_v) { + + const block_q2_0 * K_q2_0 = (const block_q2_0 *) K_c; + GGML_UNUSED(Q_v); + + float sum = 0.0f; + + // One int32 word of Q_q8 holds 4 int8 query codes; Q_ds is one float2 per word. + // QK2_0 == 32, so a q2_0 block spans 8 words and m/d are constant within a block. + // OSCAR q2_0 Lloyd-Max reconstruction levels: code -> centroid. value = m + d*centroid[code]. + constexpr float kQ2_0_lm_centroids[4] = {-0.9816f, -0.4528f, 0.4528f, 0.9816f}; +#pragma unroll + for (int k_KQ_0 = 0; k_KQ_0 < int(D/sizeof(int)); k_KQ_0 += nthreads) { + const int k_KQ = k_KQ_0 + (nthreads == WARP_SIZE ? threadIdx.x : threadIdx.x % nthreads); + + const int ib = (k_KQ * 4) / QK2_0; + const float d = __half2float(K_q2_0[ib].d); + const float m = __half2float(K_q2_0[ib].m); + const int u = Q_q8[k_KQ_0/nthreads]; + const float2 Q_ds = ((const float2 *) Q_ds_v)[k_KQ_0/nthreads]; + const float q8scale = Q_ds.x; + const float q8off = Q_ds.y / QI8_1; + +#pragma unroll + for (int b = 0; b < 4; ++b) { + const int k = k_KQ * 4 + b; + const int j = k % QK2_0; + const int by = j / 4; + const int sub = j % 4; + const uint8_t code = (K_q2_0[ib].qs[by] >> (2 * sub)) & 0x03; + const float val = m + d * kQ2_0_lm_centroids[code]; + const int qv = (int)(int8_t)(u >> (8 * b)); + const float q8eff = q8scale * qv - q8off; + sum += val * q8eff; + } + } + + return sum; +} + template static __device__ __forceinline__ float vec_dot_fattn_vec_KQ_q4_1( const char * __restrict__ K_c, const void * __restrict__ Q_v, const int * __restrict__ Q_q8, const void * __restrict__ Q_ds_v) { @@ -544,6 +586,42 @@ static __device__ __forceinline__ void dequantize_V_q5_1(const void * __restrict } } +// OSCAR q2_0 (INT2) V-cache dequant for flash-attn. 2-bit codes index Lloyd-Max +// reconstruction levels for N(0,sigma), scaled by per-block sigma (d) and offset by +// the group mean (m). Must EXACTLY match the CPU reference LM_CENTROIDS in +// ggml-quants.c and Metal's dequantize_q2_0. Ported from the Metal backend. + +template +static __device__ __forceinline__ void dequantize_V_q2_0(const void * __restrict__ vx, void * __restrict__ dst, const int64_t i0) { + // OSCAR q2_0 Lloyd-Max reconstruction levels: code -> centroid. value = m + d*centroid[code]. + constexpr float kQ2_0_lm_centroids[4] = {-0.9816f, -0.4528f, 0.4528f, 0.9816f}; + const block_q2_0 * x = (const block_q2_0 *) vx; + + const int64_t ib = i0 / QK2_0; // QK2_0 == 32 + + const float d = __half2float(x[ib].d); + const float m = __half2float(x[ib].m); + + static_assert(ne == 2 || ne == 4, "bad ne"); + +#pragma unroll + for (int l = 0; l < ne; ++l) { + const int64_t j = (i0 + l) % QK2_0; + const int by = j / 4; + const int sub = j % 4; + const uint8_t code = (x[ib].qs[by] >> (2 * sub)) & 0x03; + const float val = m + d * kQ2_0_lm_centroids[code]; + + if constexpr (std::is_same_v) { + ((half *) dst)[l] = __float2half(val); + } else if constexpr (std::is_same_v) { + ((float *) dst)[l] = val; + } else { + static_assert(std::is_same_v, "bad type"); + } + } +} + template static __device__ __forceinline__ void dequantize_V_q8_0(const void * __restrict__ vx, void * __restrict__ dst, const int64_t i0) { const block_q8_0 * x = (const block_q8_0 *) vx; @@ -591,6 +669,8 @@ constexpr __device__ vec_dot_KQ_t get_vec_dot_KQ() { return vec_dot_fattn_vec_KQ_q5_1; } else if constexpr (type_K == GGML_TYPE_Q8_0) { return vec_dot_fattn_vec_KQ_q8_0; + } else if constexpr (type_K == GGML_TYPE_Q2_0) { + return vec_dot_fattn_vec_KQ_q2_0; } else if constexpr (type_K == GGML_TYPE_BF16) { return vec_dot_fattn_vec_KQ_bf16; } else { @@ -613,6 +693,8 @@ constexpr __device__ dequantize_V_t get_dequantize_V() { return dequantize_V_q5_1; } else if constexpr (type_V == GGML_TYPE_Q8_0) { return dequantize_V_q8_0; + } else if constexpr (type_V == GGML_TYPE_Q2_0) { + return dequantize_V_q2_0; } else if constexpr (type_V == GGML_TYPE_BF16) { return dequantize_V_bf16; } else { diff --git a/ggml/src/ggml-cuda/fattn-vec.cuh b/ggml/src/ggml-cuda/fattn-vec.cuh index b0a6cf67f..1dc27e768 100644 --- a/ggml/src/ggml-cuda/fattn-vec.cuh +++ b/ggml/src/ggml-cuda/fattn-vec.cuh @@ -577,6 +577,7 @@ void ggml_cuda_flash_attn_ext_vec_case(ggml_backend_cuda_context & ctx, ggml_ten extern DECL_FATTN_VEC_CASE(D, type_K, GGML_TYPE_Q5_1); \ extern DECL_FATTN_VEC_CASE(D, type_K, GGML_TYPE_Q8_0); \ extern DECL_FATTN_VEC_CASE(D, type_K, GGML_TYPE_BF16); \ + extern DECL_FATTN_VEC_CASE(D, type_K, GGML_TYPE_Q2_0); EXTERN_DECL_FATTN_VEC_CASES( 64, GGML_TYPE_F16) EXTERN_DECL_FATTN_VEC_CASES( 64, GGML_TYPE_Q4_0) @@ -585,6 +586,7 @@ EXTERN_DECL_FATTN_VEC_CASES( 64, GGML_TYPE_Q5_0) EXTERN_DECL_FATTN_VEC_CASES( 64, GGML_TYPE_Q5_1) EXTERN_DECL_FATTN_VEC_CASES( 64, GGML_TYPE_Q8_0) EXTERN_DECL_FATTN_VEC_CASES( 64, GGML_TYPE_BF16) +EXTERN_DECL_FATTN_VEC_CASES( 64, GGML_TYPE_Q2_0) EXTERN_DECL_FATTN_VEC_CASES(128, GGML_TYPE_F16) EXTERN_DECL_FATTN_VEC_CASES(128, GGML_TYPE_Q4_0) @@ -593,6 +595,7 @@ EXTERN_DECL_FATTN_VEC_CASES(128, GGML_TYPE_Q5_0) EXTERN_DECL_FATTN_VEC_CASES(128, GGML_TYPE_Q5_1) EXTERN_DECL_FATTN_VEC_CASES(128, GGML_TYPE_Q8_0) EXTERN_DECL_FATTN_VEC_CASES(128, GGML_TYPE_BF16) +EXTERN_DECL_FATTN_VEC_CASES(128, GGML_TYPE_Q2_0) EXTERN_DECL_FATTN_VEC_CASES(256, GGML_TYPE_F16) EXTERN_DECL_FATTN_VEC_CASES(256, GGML_TYPE_Q4_0) @@ -601,3 +604,12 @@ EXTERN_DECL_FATTN_VEC_CASES(256, GGML_TYPE_Q5_0) EXTERN_DECL_FATTN_VEC_CASES(256, GGML_TYPE_Q5_1) EXTERN_DECL_FATTN_VEC_CASES(256, GGML_TYPE_Q8_0) EXTERN_DECL_FATTN_VEC_CASES(256, GGML_TYPE_BF16) +EXTERN_DECL_FATTN_VEC_CASES(256, GGML_TYPE_Q2_0) + +// OSCAR INT2 KV (q2_0) at D=512 (Gemma4 head_dim): declare only the q2_0 V-combos +// we explicitly instantiate in fattn.cu, to avoid dangling extern decls for the other types. +extern DECL_FATTN_VEC_CASE(512, GGML_TYPE_Q2_0, GGML_TYPE_Q2_0); +extern DECL_FATTN_VEC_CASE(512, GGML_TYPE_F16, GGML_TYPE_Q2_0); +extern DECL_FATTN_VEC_CASE(512, GGML_TYPE_Q2_0, GGML_TYPE_F16); +extern DECL_FATTN_VEC_CASE(512, GGML_TYPE_Q8_0, GGML_TYPE_Q2_0); +extern DECL_FATTN_VEC_CASE(512, GGML_TYPE_Q2_0, GGML_TYPE_Q8_0); diff --git a/ggml/src/ggml-cuda/fattn.cu b/ggml/src/ggml-cuda/fattn.cu index 1c7777e8a..c0e31f0fd 100644 --- a/ggml/src/ggml-cuda/fattn.cu +++ b/ggml/src/ggml-cuda/fattn.cu @@ -311,6 +311,19 @@ static void ggml_cuda_flash_attn_ext_vec(ggml_backend_cuda_context & ctx, ggml_t FATTN_VEC_CASES_ALL_D(GGML_TYPE_Q8_0, GGML_TYPE_Q8_0) FATTN_VEC_CASES_ALL_D(GGML_TYPE_BF16, GGML_TYPE_Q8_0) + // OSCAR INT2 KV (q2_0): full coverage mirroring the q8_0 matrix so enabling q2_0 + // in supports_op never hits a missing pipeline. Includes D=512 (Gemma4 head_dim). + FATTN_VEC_CASES_ALL_D(GGML_TYPE_Q2_0, GGML_TYPE_Q2_0) + FATTN_VEC_CASE(512, GGML_TYPE_Q2_0, GGML_TYPE_Q2_0) + FATTN_VEC_CASE(512, GGML_TYPE_F16, GGML_TYPE_Q2_0) + FATTN_VEC_CASE(512, GGML_TYPE_Q2_0, GGML_TYPE_F16) + FATTN_VEC_CASE(512, GGML_TYPE_Q8_0, GGML_TYPE_Q2_0) + FATTN_VEC_CASE(512, GGML_TYPE_Q2_0, GGML_TYPE_Q8_0) + FATTN_VEC_CASES_ALL_D(GGML_TYPE_F16, GGML_TYPE_Q2_0) + FATTN_VEC_CASES_ALL_D(GGML_TYPE_Q2_0, GGML_TYPE_F16) + FATTN_VEC_CASES_ALL_D(GGML_TYPE_Q8_0, GGML_TYPE_Q2_0) + FATTN_VEC_CASES_ALL_D(GGML_TYPE_Q2_0, GGML_TYPE_Q8_0) + FATTN_VEC_CASES_ALL_D(GGML_TYPE_F16, GGML_TYPE_BF16) FATTN_VEC_CASES_ALL_D(GGML_TYPE_Q4_0, GGML_TYPE_BF16) FATTN_VEC_CASES_ALL_D(GGML_TYPE_Q4_1, GGML_TYPE_BF16) @@ -440,11 +453,20 @@ static best_fattn_kernel ggml_cuda_get_best_fattn_kernel(const int device, const case GGML_TYPE_Q4_0: case GGML_TYPE_Q8_0: case GGML_TYPE_BF16: + case GGML_TYPE_Q2_0: break; default: return BEST_FATTN_KERNEL_NONE; } + // OSCAR INT2 q2_0 KV cache: only the vector flash-attn kernel (which we ported + // with the q2_0 dequant / KQ-dot) supports this type. Force it regardless of + // head dimensions / tensor-core heuristics (which would otherwise pick MMA-F16 and + // crash). Mirrors how the Metal backend routes q2_0 to its vec kernel. + if (K->type == GGML_TYPE_Q2_0 || V->type == GGML_TYPE_Q2_0) { + return BEST_FATTN_KERNEL_VEC; + } + if (mask && mask->ne[2] != 1) { return BEST_FATTN_KERNEL_NONE; } @@ -558,5 +580,36 @@ void ggml_cuda_flash_attn_ext(ggml_backend_cuda_context & ctx, ggml_tensor * dst } bool ggml_cuda_flash_attn_ext_supported(int device, const ggml_tensor * dst) { + // OSCAR two-tier mixed-precision attention (INT2-LP + F16-HP) has a native + // CUDA kernel (flash-attn-mixed.cu) regardless of head dims, so it is always + // supported when mixed mode is set. + if (ggml_get_op_params_i32(dst, 4) == 1) { + return true; + } return ggml_cuda_get_best_fattn_kernel(device, dst) != BEST_FATTN_KERNEL_NONE; } + +// OSCAR INT2 q2_0 KV-cache flash-attn vec-kernel explicit instantiations. These are +// referenced (declared extern) from fattn-vec.cuh and are not auto-generated by the +// fattn-vec-instance list, so we define them here. Mirrors the q8_0 matrix of cases. +template void ggml_cuda_flash_attn_ext_vec_case< 64, GGML_TYPE_Q2_0, GGML_TYPE_Q2_0>(ggml_backend_cuda_context &, ggml_tensor *); +template void ggml_cuda_flash_attn_ext_vec_case<128, GGML_TYPE_Q2_0, GGML_TYPE_Q2_0>(ggml_backend_cuda_context &, ggml_tensor *); +template void ggml_cuda_flash_attn_ext_vec_case<256, GGML_TYPE_Q2_0, GGML_TYPE_Q2_0>(ggml_backend_cuda_context &, ggml_tensor *); +template void ggml_cuda_flash_attn_ext_vec_case< 64, GGML_TYPE_F16, GGML_TYPE_Q2_0>(ggml_backend_cuda_context &, ggml_tensor *); +template void ggml_cuda_flash_attn_ext_vec_case<128, GGML_TYPE_F16, GGML_TYPE_Q2_0>(ggml_backend_cuda_context &, ggml_tensor *); +template void ggml_cuda_flash_attn_ext_vec_case<256, GGML_TYPE_F16, GGML_TYPE_Q2_0>(ggml_backend_cuda_context &, ggml_tensor *); +template void ggml_cuda_flash_attn_ext_vec_case< 64, GGML_TYPE_Q2_0, GGML_TYPE_F16>(ggml_backend_cuda_context &, ggml_tensor *); +template void ggml_cuda_flash_attn_ext_vec_case<128, GGML_TYPE_Q2_0, GGML_TYPE_F16>(ggml_backend_cuda_context &, ggml_tensor *); +template void ggml_cuda_flash_attn_ext_vec_case<256, GGML_TYPE_Q2_0, GGML_TYPE_F16>(ggml_backend_cuda_context &, ggml_tensor *); +template void ggml_cuda_flash_attn_ext_vec_case< 64, GGML_TYPE_Q8_0, GGML_TYPE_Q2_0>(ggml_backend_cuda_context &, ggml_tensor *); +template void ggml_cuda_flash_attn_ext_vec_case<128, GGML_TYPE_Q8_0, GGML_TYPE_Q2_0>(ggml_backend_cuda_context &, ggml_tensor *); +template void ggml_cuda_flash_attn_ext_vec_case<256, GGML_TYPE_Q8_0, GGML_TYPE_Q2_0>(ggml_backend_cuda_context &, ggml_tensor *); +template void ggml_cuda_flash_attn_ext_vec_case< 64, GGML_TYPE_Q2_0, GGML_TYPE_Q8_0>(ggml_backend_cuda_context &, ggml_tensor *); +template void ggml_cuda_flash_attn_ext_vec_case<128, GGML_TYPE_Q2_0, GGML_TYPE_Q8_0>(ggml_backend_cuda_context &, ggml_tensor *); +template void ggml_cuda_flash_attn_ext_vec_case<256, GGML_TYPE_Q2_0, GGML_TYPE_Q8_0>(ggml_backend_cuda_context &, ggml_tensor *); +// D=512 (Gemma4 head_dim) q2_0 VEC instances: +template void ggml_cuda_flash_attn_ext_vec_case<512, GGML_TYPE_Q2_0, GGML_TYPE_Q2_0>(ggml_backend_cuda_context &, ggml_tensor *); +template void ggml_cuda_flash_attn_ext_vec_case<512, GGML_TYPE_F16, GGML_TYPE_Q2_0>(ggml_backend_cuda_context &, ggml_tensor *); +template void ggml_cuda_flash_attn_ext_vec_case<512, GGML_TYPE_Q2_0, GGML_TYPE_F16>(ggml_backend_cuda_context &, ggml_tensor *); +template void ggml_cuda_flash_attn_ext_vec_case<512, GGML_TYPE_Q8_0, GGML_TYPE_Q2_0>(ggml_backend_cuda_context &, ggml_tensor *); +template void ggml_cuda_flash_attn_ext_vec_case<512, GGML_TYPE_Q2_0, GGML_TYPE_Q8_0>(ggml_backend_cuda_context &, ggml_tensor *); diff --git a/ggml/src/ggml-cuda/fattn.cuh b/ggml/src/ggml-cuda/fattn.cuh index 78705d599..376b5ae10 100644 --- a/ggml/src/ggml-cuda/fattn.cuh +++ b/ggml/src/ggml-cuda/fattn.cuh @@ -2,4 +2,7 @@ void ggml_cuda_flash_attn_ext(ggml_backend_cuda_context & ctx, ggml_tensor * dst); +// OSCAR two-tier mixed-precision fused attention (INT2-LP history + F16-HP). +void ggml_cuda_flash_attn_ext_mixed(ggml_backend_cuda_context & ctx, ggml_tensor * dst); + bool ggml_cuda_flash_attn_ext_supported(int device, const ggml_tensor * dst); diff --git a/ggml/src/ggml-cuda/flash-attn-mixed.cu b/ggml/src/ggml-cuda/flash-attn-mixed.cu new file mode 100644 index 000000000..dc9d8a24c --- /dev/null +++ b/ggml/src/ggml-cuda/flash-attn-mixed.cu @@ -0,0 +1,242 @@ +// +// flash-attn-mixed.cu +// +// Mixed flash attention kernel that reads from two KV cache tiers: +// LP (low precision, q2_0 quantised) – most K/V entries +// HP (high precision, f16) – recent + sink entries +// +// One CUDA thread block handles one query row (iq3,iq2,iq1). The T threads each +// compute the full KQ dot (strided over DK) so no cross-thread reduction of the +// score is needed; each thread owns a disjoint DV/T slice of the value +// accumulator, giving an independent online-softmax per thread that is +// bit-identical across threads. The two tiers (LP then HP) feed one shared +// online-softmax so the merge is mathematically exact. +// op_params[4]==1 marks mixed mode. +// + +#include "ggml-cuda.h" +#include "common.cuh" +#include "fattn-common.cuh" + +// ============================================================ +// Dequantise a single q2_0 block row into fp32 +// ============================================================ + +static void dbg_dequantize_q2_0(const block_q2_0 * x, float * y, int64_t k) { + constexpr float c[4] = {-0.9816f, -0.4528f, 0.4528f, 0.9816f}; + for (int64_t i = 0; i < k; i += QK2_0) { + const float d = __half2float(x->d); + const float m = __half2float(x->m); + const uint8_t * qs = x->qs; + for (int j = 0; j < QK2_0; ++j) { + const int code = (qs[j/4] >> (2*(j%4))) & 0x03; + y[i + j] = m + d * c[code]; + } + ++x; + } +} + +// ============================================================ +// KQ dot product helpers +// ============================================================ + +__device__ __forceinline__ float fp16_to_f32(ggml_fp16_t h) { + return __half2float(__half(h)); +} + +__device__ __forceinline__ float kq_dot_q2_0(const char * kd, const float * q, int64_t DK) { + constexpr float c[4] = {-0.9816f, -0.4528f, 0.4528f, 0.9816f}; + float sum = 0.0f; + const int nblocks = DK / QK2_0; + for (int b = 0; b < nblocks; ++b) { + const block_q2_0 * kb = ((const block_q2_0 *)kd) + b; + const float mean = __half2float(kb->m); + const float d = __half2float(kb->d); + const uint8_t * qs = kb->qs; + for (int j = 0; j < QK2_0; ++j) { + const int code = (qs[j/4] >> (2*(j%4))) & 0x03; + sum += (mean + d * c[code]) * q[b*QK2_0 + j]; + } + } + return sum; +} + +__device__ __forceinline__ float kq_dot_f16(const half * kd, const float * q, int64_t DK) { + float sum = 0.0f; + for (int i = 0; i < DK; ++i) { + sum += __half2float(kd[i]) * q[i]; + } + return sum; +} + +// ============================================================ +// Fused mixed-attention CUDA kernel +// ============================================================ +template +void __global__ flash_attn_ext_mixed_kernel( + const char * q, int64_t nbq1, int64_t nbq2, int64_t nbq3, + const char * k_lp, int64_t nbk1, int64_t nbk2, int64_t nbk3, + const char * v_lp, int64_t nbv1, int64_t nbv2, int64_t nbv3, + const char * mask_lp_data, int64_t mask_lp_nb1, int64_t mask_lp_nb2, int64_t mask_lp_nb3, int64_t mask_lp_ne2, int64_t mask_lp_ne3, + const char * k_hp, int64_t nbkh1, int64_t nbkh2, int64_t nbkh3, + const char * v_hp, int64_t nbvh1, int64_t nbvh2, int64_t nbvh3, + const char * mask_hp_data, int64_t mask_hp_nb1, int64_t mask_hp_nb2, int64_t mask_hp_nb3, int64_t mask_hp_ne2, int64_t mask_hp_ne3, + float * dst, + int64_t DK, int64_t DV, int64_t n_kv, int64_t n_hp, int64_t N, int32_t n_head, int64_t nseq, + float scale, float max_bias, float logit_softcap, + int rk2, int rk3, int rv2, int rv3, int rk2h, int rk3h, int rv2h, int rv3h) +{ + // blockIdx.x encodes (iq3,iq2,iq1) packed linearly + const int64_t ir = blockIdx.x; + const int iq3 = ir / (n_head * N); + const int iq2 = (ir - iq3 * n_head * N) / N; + const int iq1 = ir - iq3 * n_head * N - iq2 * N; + + const uint32_t h = iq2; + const uint32_t n_head_log2 = 1u << (uint32_t) floorf(log2f((float) n_head)); + const float m0 = powf(2.0f, -(max_bias ) / n_head_log2); + const float m1 = powf(2.0f, -(max_bias / 2.0f) / n_head_log2); + const float slope = (max_bias > 0.0f) ? (h < n_head_log2 ? powf(m0, h + 1) : powf(m1, 2 * (h - n_head_log2) + 1)) : 1.0f; + + const float * pq = (const float *)((const char *) q + (iq1 * nbq1 + iq2 * nbq2 + iq3 * nbq3)); + + const int lane = threadIdx.x; + const int slice = DV / T; // each thread owns slice DV/T values of V + const int j0 = lane * slice; + + float v_acc[128]; // DV/T <= 128 (DV<=512, T>=4) + for (int j = 0; j < slice; ++j) v_acc[j] = 0.0f; + + float M = -INFINITY; + float S = 0.0f; + + const ggml_fp16_t * mp_lp = mask_lp_data ? (const ggml_fp16_t *)(mask_lp_data + iq1*mask_lp_nb1 + (iq2%mask_lp_ne2)*mask_lp_nb2 + (iq3%mask_lp_ne3)*mask_lp_nb3) : NULL; + const ggml_fp16_t * mp_hp = mask_hp_data ? (const ggml_fp16_t *)(mask_hp_data + iq1*mask_hp_nb1 + (iq2%mask_hp_ne2)*mask_hp_nb2 + (iq3%mask_hp_ne3)*mask_hp_nb3) : NULL; + + const int ik2 = iq2 / rk2; const int ik3 = iq3 / rk3; + const int iv2 = iq2 / rv2; const int iv3 = iq3 / rv3; + const int ik2h = iq2 / rk2h; const int ik3h = iq3 / rk3h; + const int iv2h = iq2 / rv2h; const int iv3h = iq3 / rv3h; + + // ---- LP tier (q2_0 K/V) ---- + for (int64_t ic = 0; ic < n_kv; ++ic) { + const float mv = mp_lp ? slope * fp16_to_f32(mp_lp[ic]) : 0.0f; + if (mv == -INFINITY) continue; + + const char * kd = (const char *) k_lp + (ic * nbk1 + ik2 * nbk2 + ik3 * nbk3); + float s = kq_dot_q2_0(kd, pq, DK) * scale; + if (logit_softcap != 0.0f) s = logit_softcap * tanhf(s); + s += mv; + + const float Mold = M; + float ms = 1.0f, vs = 1.0f; + const char * vd = (const char *) v_lp + (ic * nbv1 + iv2 * nbv2 + iv3 * nbv3); + if (s > M) { M = s; ms = expf(Mold - M); for (int j = 0; j < slice; ++j) v_acc[j] *= ms; } + else { vs = expf(s - M); } + // V dequant: q2_0 block + const block_q2_0 * vblocks = (const block_q2_0 *) vd; + for (int j = 0; j < slice; ++j) { + const int dim = j0 + j; + const block_q2_0 * vb = &vblocks[dim / QK2_0]; + const int sub = dim % QK2_0; + const int code = (vb->qs[sub/4] >> (2*(sub%4))) & 0x03; + constexpr float cc[4] = {-0.9816f, -0.4528f, 0.4528f, 0.9816f}; + const float val = __half2float(vb->m) + __half2float(vb->d) * cc[code]; + v_acc[j] += vs * val; + } + S = S * ms + vs; + } + + // ---- HP tier (f16 K/V) ---- + for (int64_t ic = 0; ic < n_hp; ++ic) { + const float mv = mp_hp ? slope * fp16_to_f32(mp_hp[ic]) : 0.0f; + if (mv == -INFINITY) continue; + + const char * kd = (const char *) k_hp + (ic * nbkh1 + ik2h * nbkh2 + ik3h * nbkh3); + float s = kq_dot_f16((const half *)kd, pq, DK) * scale; + if (logit_softcap != 0.0f) s = logit_softcap * tanhf(s); + s += mv; + + const float Mold = M; + float ms = 1.0f, vs = 1.0f; + const char * vd = (const char *) v_hp + (ic * nbvh1 + iv2h * nbvh2 + iv3h * nbvh3); + if (s > M) { M = s; ms = expf(Mold - M); for (int j = 0; j < slice; ++j) v_acc[j] *= ms; } + else { vs = expf(s - M); } + const half * vh = (const half *) vd; + for (int j = 0; j < slice; ++j) v_acc[j] += vs * __half2float(vh[j0 + j]); + S = S * ms + vs; + } + + const float S_inv = (S == 0.0f) ? 0.0f : 1.0f / S; + float * out = (float *)((char *) dst + (iq3 * n_head * N + iq1 * N + iq2) * DV * sizeof(float)); + for (int j = 0; j < slice; ++j) out[j0 + j] = v_acc[j] * S_inv; +} + +void ggml_cuda_flash_attn_ext_mixed(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { + const ggml_tensor * q = dst->src[0]; + const ggml_tensor * k_lp = dst->src[1]; + const ggml_tensor * v_lp = dst->src[2]; + const ggml_tensor * mask_lp = dst->src[3]; + const ggml_tensor * k_hp = dst->src[5]; + const ggml_tensor * v_hp = dst->src[6]; + const ggml_tensor * mask_hp = dst->src[7]; + + GGML_TENSOR_LOCALS(int64_t, neq, q, ne); + GGML_TENSOR_LOCALS(size_t, nbq, q, nb); + GGML_TENSOR_LOCALS(int64_t, nek, k_lp, ne); + GGML_TENSOR_LOCALS(size_t, nbk, k_lp, nb); + GGML_TENSOR_LOCALS(int64_t, nev, v_lp, ne); + GGML_TENSOR_LOCALS(size_t, nbv, v_lp, nb); + GGML_TENSOR_LOCALS(int64_t, nekh, k_hp, ne); + GGML_TENSOR_LOCALS(size_t, nbkh, k_hp, nb); + GGML_TENSOR_LOCALS(int64_t, nevh, v_hp, ne); + GGML_TENSOR_LOCALS(size_t, nbvh, v_hp, nb); + + const int64_t DK = nek0; + const int64_t DV = nev0; + const int64_t N = neq1; + const int64_t n_head = neq2; + const int64_t nseq = neq3; + + float scale = 1.0f, max_bias = 0.0f, logit_softcap = 0.0f; + memcpy(&scale, (float *) dst->op_params + 0, sizeof(float)); + memcpy(&max_bias, (float *) dst->op_params + 1, sizeof(float)); + memcpy(&logit_softcap, (float *) dst->op_params + 2, sizeof(float)); + if (logit_softcap != 0.0f) scale /= logit_softcap; + + const int64_t n_kv = nek1; + const int64_t n_hp = nekh1; + + const int rk2 = neq2 / nek2, rk3 = neq3 / nek3; + const int rv2 = neq2 / nev2, rv3 = neq3 / nev3; + const int rk2h = neq2 / nekh2, rk3h = neq3 / nekh3; + const int rv2h = neq2 / nevh2, rv3h = neq3 / nevh3; + + // pick thread count: DV must be divisible by T + int T = (DV >= 256) ? 256 : (DV >= 128) ? 128 : (DV >= 64) ? 64 : 32; + while (T > 1 && DV % T != 0) T /= 2; + GGML_ASSERT(DV % T == 0 && "mixed FA: DV must divide thread count"); + + const int64_t nrows = nseq * n_head * N; + const dim3 grid(nrows); + const dim3 block(T); + + cudaStream_t stream = ctx.stream(); + + flash_attn_ext_mixed_kernel<256><<>>( + (const char *) q->data, + nbq1, nbq2, nbq3, + (const char *) k_lp->data, nbk1, nbk2, nbk3, + (const char *) v_lp->data, nbv1, nbv2, nbv3, + mask_lp ? (const char *) mask_lp->data : nullptr, mask_lp ? mask_lp->nb[1] : 0, mask_lp ? mask_lp->nb[2] : 0, mask_lp ? mask_lp->nb[3] : 0, mask_lp ? mask_lp->ne[2] : 0, mask_lp ? mask_lp->ne[3] : 0, + (const char *) k_hp->data, nbkh1, nbkh2, nbkh3, + (const char *) v_hp->data, nbvh1, nbvh2, nbvh3, + mask_hp ? (const char *) mask_hp->data : nullptr, mask_hp ? mask_hp->nb[1] : 0, mask_hp ? mask_hp->nb[2] : 0, mask_hp ? mask_hp->nb[3] : 0, mask_hp ? mask_hp->ne[2] : 0, mask_hp ? mask_hp->ne[3] : 0, + (float *) dst->data, + DK, DV, n_kv, n_hp, N, n_head, nseq, + scale, max_bias, logit_softcap, + rk2, rk3, rv2, rv3, rk2h, rk3h, rv2h, rv3h); + + GGML_ASSERT(cudaGetLastError() == cudaSuccess); + cudaDeviceSynchronize(); +} diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 18aaa0983..1ece17aab 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -3077,7 +3077,12 @@ static bool ggml_cuda_compute_forward(ggml_backend_cuda_context & ctx, struct gg ggml_cuda_op_argsort(ctx, dst); break; case GGML_OP_FLASH_ATTN_EXT: - ggml_cuda_flash_attn_ext(ctx, dst); + // OSCAR two-tier mixed-precision fused attention (q2_0-LP + f16-HP). + if (ggml_get_op_params_i32(dst, 4) == 1) { + ggml_cuda_flash_attn_ext_mixed(ctx, dst); + } else { + ggml_cuda_flash_attn_ext(ctx, dst); + } break; case GGML_OP_CROSS_ENTROPY_LOSS: ggml_cuda_cross_entropy_loss(ctx, dst); @@ -5207,7 +5212,8 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g { return (op->type == GGML_TYPE_F32 || op->type == GGML_TYPE_F16 || op->type == GGML_TYPE_BF16 || op->type == GGML_TYPE_Q4_0 || op->type == GGML_TYPE_Q4_1 || op->type == GGML_TYPE_Q5_0 || - op->type == GGML_TYPE_Q5_1 || op->type == GGML_TYPE_Q8_0 || op->type == GGML_TYPE_IQ4_NL) && + op->type == GGML_TYPE_Q5_1 || op->type == GGML_TYPE_Q8_0 || op->type == GGML_TYPE_IQ4_NL || + op->type == GGML_TYPE_Q2_0) && op->src[0]->type == GGML_TYPE_F32 && (op->src[1]->type == GGML_TYPE_I64 || op->src[1]->type == GGML_TYPE_I32); } break; diff --git a/ggml/src/ggml-cuda/set-rows.cu b/ggml/src/ggml-cuda/set-rows.cu index e14f96b82..2354d3d13 100644 --- a/ggml/src/ggml-cuda/set-rows.cu +++ b/ggml/src/ggml-cuda/set-rows.cu @@ -1,5 +1,6 @@ #include "set-rows.cuh" #include "cpy-utils.cuh" +#include typedef void (*set_rows_kernel_t)(const char * src, char * dst); @@ -320,16 +321,176 @@ static void set_rows_cuda(ggml_backend_cuda_context & ctx, const ggml_tensor * s } -void ggml_cuda_op_set_rows(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { - const ggml_tensor * src0 = dst->src[0]; - const ggml_tensor * src1 = dst->src[1]; - - GGML_ASSERT(src0->type == GGML_TYPE_F32); - GGML_ASSERT(src1->type == GGML_TYPE_I64 || src1->type == GGML_TYPE_I32); - - if (src1->type == GGML_TYPE_I64) { - set_rows_cuda(ctx, src0, src1, dst); - } else { - set_rows_cuda(ctx, src0, src1, dst); +// OSCAR q2_0 (INT2) SET_ROWS write kernel. Quantizes an f32 K/V row into q2_0 using +// the calibrated OSCAR encode: per-128-group mean, optional outlier clip (clip_ratio +// percentile), per-32-block sigma, and 2-bit Lloyd-Max packing (vs < -0.6745s -> 0, +// < 0 -> 1, < 0.6745s -> 2, else 3). The group mean is replicated into every block's +// m field so the per-block dequant (m + d*centroid) is exact. Requires ne00 % 128 == 0 +// and the LLAMA_KV_NO_HADAMARD path (no in-quant OWHT). Faithful port of Metal's +// kernel_set_rows_q2_0. +template +static __global__ void set_rows_cuda_q2_0( + const char * src0, const char * src1, char * dst, + const int32_t ne01, const int32_t ne11, const int32_t ne12, + const uint64_t nb01, const uint64_t nb02, const uint64_t nb03, + const uint64_t nb10, const uint64_t nb11, const uint64_t nb12, + const uint64_t nb1, const uint64_t nb2, const uint64_t nb3, + const int32_t nk0, const float clip_ratio) { + + const int32_t i03 = blockIdx.z; + const int32_t i02 = blockIdx.y; + const int32_t i01 = blockIdx.x; + if (i01 >= ne01) return; + + const int32_t i12 = i03 % ne12; + const int32_t i11 = i02 % ne11; + const int32_t i10 = i01; + + const idx_t i1 = *((const idx_t *) (src1 + i10*nb10 + i11*nb11 + i12*nb12)); + + block_q2_0 * dst_row = (block_q2_0 *) (dst + (uint64_t)i1*nb1 + i02*nb2 + i03*nb3); + const float * src_row = (const float *) (src0 + i01*nb01 + i02*nb02 + i03*nb03); + + const int ngrp = nk0 / 4; // 128-wide groups per row (4 blocks of 32) + + __shared__ float sh[128]; + __shared__ float sh_sigma[4]; + __shared__ float sh_mean; + __shared__ float sh_thr; + + const unsigned t = threadIdx.x; // 0..31 + + for (int g = 0; g < ngrp; ++g) { + const float * gsrc = src_row + g * 128; + + for (int k = 0; k < 4; ++k) { + sh[t * 4 + k] = gsrc[t * 4 + k]; + } + __syncthreads(); + + // group mean + if (t == 0) { + float s = 0.0f; + for (int j = 0; j < 128; ++j) s += sh[j]; + sh_mean = s / 128.0f; + } + __syncthreads(); + + const float mean = sh_mean; + for (int k = 0; k < 4; ++k) sh[t * 4 + k] -= mean; + __syncthreads(); + + // OSCAR outlier clip: threshold = clip_ratio percentile over the 128 group, + // found by exact rank counting (matches CPU qsort + index selection). + if (clip_ratio > 0.0f && clip_ratio < 1.0f) { + if (t == 0) sh_thr = 0.0f; + __syncthreads(); + + int idx = (int)(clip_ratio * 128.0f); + if (idx >= 128) idx = 127; + for (int k = 0; k < 4; ++k) { + const float a = fabsf(sh[t * 4 + k]); + int lo = 0, le = 0; + for (int j = 0; j < 128; ++j) { + const float aj = fabsf(sh[j]); + lo += (aj < a) ? 1 : 0; + le += (aj <= a) ? 1 : 0; + } + if (lo <= idx && idx < le) sh_thr = a; + } + __syncthreads(); + + const float thr = sh_thr; + for (int k = 0; k < 4; ++k) { + float v = sh[t * 4 + k]; + if (v > thr) v = thr; + if (v < -thr) v = -thr; + sh[t * 4 + k] = v; + } + __syncthreads(); + } + + // per-32-block sigma (RMS) + if (t < 4) { + float ss = 0.0f; + for (int j = 0; j < 32; ++j) ss += sh[t * 32 + j] * sh[t * 32 + j]; + sh_sigma[t] = sqrtf(ss / 32.0f); + } + __syncthreads(); + + // quantize: thread t writes one packed byte (block b = t/8, byte t%8) + const int b = t / 8; + const float sigma = sh_sigma[b]; + const float inv_sigma = (sigma > 1e-8f) ? (1.0f / sigma) : 0.0f; + + uint8_t packed = 0; + for (int k = 0; k < 4; ++k) { + const float vs = sh[t * 4 + k] * inv_sigma; + uint8_t code; + if (vs < -0.6745f) code = 0; + else if (vs < 0.0f) code = 1; + else if (vs < 0.6745f) code = 2; + else code = 3; + packed |= code << (2 * k); + } + + block_q2_0 & blk = dst_row[g * 4 + b]; + blk.qs[t % 8] = packed; + if (t % 8 == 0) { + blk.d = __float2half(sigma); + blk.m = __float2half(mean); + } + __syncthreads(); } } + + void ggml_cuda_op_set_rows(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { + const ggml_tensor * src0 = dst->src[0]; + const ggml_tensor * src1 = dst->src[1]; + + GGML_ASSERT(src0->type == GGML_TYPE_F32); + GGML_ASSERT(src1->type == GGML_TYPE_I64 || src1->type == GGML_TYPE_I32); + + if (dst->type == GGML_TYPE_Q2_0) { + // OSCAR INT2 KV write: dedicated kernel (per-128-group mean + clip + per-block + // Lloyd-Max). clip_ratio comes from LLAMA_KV_CLIP_RATIO (0 disables). + float clip_ratio = 0.0f; + if (const char * e = getenv("LLAMA_KV_CLIP_RATIO")) { + clip_ratio = (float) atof(e); + } + + GGML_TENSOR_BINARY_OP_LOCALS(src0, src1, dst); + + const int32_t nk0 = (int32_t)(ne00 / ggml_blck_size(GGML_TYPE_Q2_0)); + cudaStream_t stream = ctx.stream(); + + const dim3 grid_size(ne01, ne02, ne03); + const dim3 block_size(32, 1, 1); + + const char * src0_d = (const char *) src0->data; + const char * src1_d = (const char *) src1->data; + char * dst_d = (char *) dst->data; + + if (src1->type == GGML_TYPE_I64) { + set_rows_cuda_q2_0<<>>( + src0_d, src1_d, dst_d, + ne01, ne11, ne12, + nb01, nb02, nb03, nb10, nb11, nb12, nb1, nb2, nb3, + nk0, clip_ratio); + } else { + set_rows_cuda_q2_0<<>>( + src0_d, src1_d, dst_d, + ne01, ne11, ne12, + nb01, nb02, nb03, nb10, nb11, nb12, nb1, nb2, nb3, + nk0, clip_ratio); + } + GGML_ASSERT(cudaGetLastError() == cudaSuccess); + return; + } + + if (src1->type == GGML_TYPE_I64) { + set_rows_cuda(ctx, src0, src1, dst); + } else { + set_rows_cuda(ctx, src0, src1, dst); + } + } diff --git a/ggml/src/ggml-cuda/vecdotq.cuh b/ggml/src/ggml-cuda/vecdotq.cuh index d1741cc8d..a154778a3 100644 --- a/ggml/src/ggml-cuda/vecdotq.cuh +++ b/ggml/src/ggml-cuda/vecdotq.cuh @@ -735,6 +735,40 @@ static __device__ __forceinline__ float vec_dot_q4_0_q8_1( return vec_dot_q4_0_q8_1_impl(v, u, bq4_0->d, bq8_1->ds); } +// OSCAR q2_0 Lloyd-Max centroids (defined inline in fattn-common.cuh; redeclared here +// so this TU can use it without dragging fattn-common.cuh in). + +static __device__ __forceinline__ float vec_dot_q2_0_q8_1( + const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs) { + + const block_q2_0 * bq2_0 = (const block_q2_0 *) vbq + kbx; + + const float d = __half2float(bq2_0->d); + const float m = __half2float(bq2_0->m); + const float d8 = __low2float(bq8_1->ds); + + // OSCAR q2_0 Lloyd-Max reconstruction levels: code -> centroid. value = m + d*centroid[code]. + constexpr float kQ2_0_lm_centroids[4] = {-0.9816f, -0.4528f, 0.4528f, 0.9816f}; + + // reconstruct the 32 q8_1 values (int8) and the 32 q2_0 codes (2-bit) + float S = 0.0f; // Σ q8_k + float C = 0.0f; // Σ centroid[code_k] * q8_k +#pragma unroll + for (int j = 0; j < 8; ++j) { + const uint32_t q8_pack = get_int_b4(bq8_1->qs, j); // 4 int8 values + const uint8_t q2_pack = bq2_0->qs[j]; // 4 2-bit codes + for (int b = 0; b < 4; ++b) { + const int qv = (int)(int8_t)(q8_pack >> (8 * b)); + const uint8_t code = (q2_pack >> (2 * b)) & 0x03; + const float cen = kQ2_0_lm_centroids[code]; + S += (float)qv; + C += cen * (float)qv; + } + } + (void)iqs; // whole-block decode (QK2_0 == 32 == QK8_1) + return d8 * (m * S + d * C); +} + static __device__ __forceinline__ float vec_dot_q4_1_q8_1( const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs) {