Pr oscar int2 kv

.gitignore

@@ -160,3 +160,6 @@ a.out.*
 
 AGENTS.local.md
 .pi/SYSTEM.md
+
+# Local model/rotation assets (large binaries)
+assets/

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.

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<int D, int nthreads>
+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<int D, int nthreads>
 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 <typename T, int ne>
+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<T, half>) {
+            ((half *) dst)[l] = __float2half(val);
+        } else if constexpr (std::is_same_v<T, float>) {
+            ((float *) dst)[l] = val;
+        } else {
+            static_assert(std::is_same_v<T, void>, "bad type");
+        }
+    }
+}
+
 template <typename T, int ne>
 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<D, nthreads>;
     } else if constexpr (type_K == GGML_TYPE_Q8_0) {
         return vec_dot_fattn_vec_KQ_q8_0<D, nthreads>;
+    } else if constexpr (type_K == GGML_TYPE_Q2_0) {
+        return vec_dot_fattn_vec_KQ_q2_0<D, nthreads>;
     } else if constexpr (type_K == GGML_TYPE_BF16) {
         return vec_dot_fattn_vec_KQ_bf16<D, nthreads>;
     } else {
@@ -613,6 +693,8 @@ constexpr __device__ dequantize_V_t get_dequantize_V() {
         return dequantize_V_q5_1<T, ne>;
     } else if constexpr (type_V == GGML_TYPE_Q8_0) {
         return dequantize_V_q8_0<T, ne>;
+    } else if constexpr (type_V == GGML_TYPE_Q2_0) {
+        return dequantize_V_q2_0<T, ne>;
     } else if constexpr (type_V == GGML_TYPE_BF16) {
         return dequantize_V_bf16<float, ne>;
     } else {

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);