Upload patch_llamacpp_hpcq_seg.py

patch_llamacpp_hpcq_seg.py

@@ -0,0 +1,977 @@
+#!/usr/bin/env python3
+"""Patch llama.cpp to natively run HPCQ.SEG.V1 segmented Q2_K / Q3_K models
+with ZERO weight expansion in RAM/VRAM.
+
+How it works
+------------
+* Tensors remain in their native ``GGML_TYPE_Q2_K`` (2.625 bpw) and
+  ``GGML_TYPE_Q3_K`` (3.4375 bpw) formats in RAM/VRAM (and stay ``mmap``-backed
+  on disk).
+* Only the compact ``HQS1`` nibble stream (128 bytes per 256-weight superblock)
+  is bound to the tensor's payload address range in ``libggml-base``.
+* No ``uint16_t`` order tables or ``uint8_t`` unpacked label arrays are
+  allocated on the heap. Instead, ``uint16_t ord[256]`` (512 bytes) is
+  reconstructed on the stack (CPU) or in ``__shared__`` memory (CUDA) per
+  256-weight superblock:
+  - CPU ``dequantize_row_q2_K`` / ``dequantize_row_q3_K``: scatters each block's
+    256 dequantized values via the 512-byte stack ``ord[256]``.
+  - CPU ``vec_dot`` (``ggml-cpu.c``): wraps ``ggml_vec_dot_q2_K_q8_K`` and
+    ``ggml_vec_dot_q3_K_q8_K``. For unsegmented tensors it tail-calls the stock
+    SIMD kernel directly; for segmented tensors it gathers ``block_q8_K y_perm``
+    (292 bytes on the stack) per superblock and calls the native AVX2/AVX512/
+    NEON SIMD ``vec_dot`` on ``(&x[i], &y_perm)``.
+  - CUDA ``dequantize_block_q2_K`` / ``dequantize_block_q3_K`` (``convert.cu``)
+    + ``ggml_cuda_mul_mat`` (``ggml-cuda.cu``): scatters each 256-weight
+    superblock in ``__shared__`` memory and routes segmented ``src0`` to the
+    pooled-scratch ``ggml_cuda_mul_mat_cublas`` path.
+
+Usage
+-----
+    python3 tools/patch_llamacpp_hpcq_seg.py --selftest
+    python3 tools/patch_llamacpp_hpcq_seg.py --llama /home/node/Desktop/Dev/llama.cpp --check
+    python3 tools/patch_llamacpp_hpcq_seg.py --llama /home/node/Desktop/Dev/llama.cpp --apply
+    python3 tools/patch_llamacpp_hpcq_seg.py --llama /home/node/Desktop/Dev/llama.cpp --revert
+"""
+
+from __future__ import annotations
+
+import argparse
+import shutil
+import subprocess
+import sys
+import tempfile
+from pathlib import Path
+
+BEGIN = "HPCQ_SEG_PATCH_BEGIN"
+END = "HPCQ_SEG_PATCH_END"
+HEADER_REL = "ggml/include/hpcq-seg.h"
+
+# --------------------------------------------------------------------------
+# 1. C/C++/CUDA header: ggml/include/hpcq-seg.h
+# --------------------------------------------------------------------------
+
+HEADER = r"""// hpcq-seg.h — Zero-expansion HPCQ.SEG.V1 band-major support for Q2_K / Q3_K.
+// Generated by HPC-Quantize tools/patch_llamacpp_hpcq_seg.py.
+#pragma once
+
+#include <stdint.h>
+#include <stddef.h>
+#include <string.h>
+#include <stdbool.h>
+#include "ggml.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+#define HPCQ_SEG_PREFIX "HPCQ.SEG.V1."
+#define HPCQ_SEG_MAX_ENTRIES 2048
+
+#ifdef __CUDACC__
+#define HPCQ_HD __host__ __device__
+#else
+#define HPCQ_HD
+#endif
+
+struct hpcq_seg_entry {
+    uintptr_t base;        // payload start (host ptr on CPU, device ptr on CUDA)
+    uintptr_t end;         // payload end (exclusive)
+    int64_t   block_bytes; // 84 for Q2_K, 110 for Q3_K
+    int64_t   n_blocks;    // number of 256-weight superblocks
+    uint8_t   enc;         // 0 = 4-bit nibbles (128B/blk), 1 = 8-bit (256B/blk)
+    uint8_t * packed;      // compact packed label stream (128B or 256B per block)
+    void    * dev_packed;  // lazily allocated device copy for CUDA (or NULL)
+};
+
+GGML_API void ggml_hpcq_seg_bind_raw(
+        const void * payload, int64_t n_bytes,
+        const uint8_t * hqs1, size_t hqs1_len, int64_t block_bytes);
+
+GGML_API const struct hpcq_seg_entry * ggml_hpcq_seg_find(
+        const void * ptr, int64_t block_bytes, int64_t * out_block_idx);
+
+GGML_API void ggml_hpcq_seg_set_dev_packed(
+        const struct hpcq_seg_entry * entry, void * dev_packed);
+
+// Reconstruct slot -> original_position permutation `order[256]` on the stack
+// or in __shared__ memory from a single 256-weight block's packed labels.
+// Matches `np.argsort(labels, kind="stable")`: slot `s` holds original pos `order[s]`.
+static inline HPCQ_HD void hpcq_seg_make_order_block(
+        const uint8_t * blk_packed, uint8_t enc, uint16_t * order) {
+    if (enc == 0) {
+        uint16_t start[16];
+        for (int b = 0; b < 16; ++b) start[b] = 0;
+        for (int p = 0; p < 256; ++p) {
+            const uint8_t lab = (blk_packed[p >> 1] >> ((p & 1) * 4)) & 0x0F;
+            start[lab]++;
+        }
+        uint16_t acc = 0;
+        for (int b = 0; b < 16; ++b) {
+            const uint16_t c = start[b];
+            start[b] = acc;
+            acc = (uint16_t) (acc + c);
+        }
+        for (uint16_t p = 0; p < 256; ++p) {
+            const uint8_t lab = (blk_packed[p >> 1] >> ((p & 1) * 4)) & 0x0F;
+            order[start[lab]++] = p;
+        }
+    } else {
+        uint16_t start[256];
+        for (int b = 0; b < 256; ++b) start[b] = 0;
+        for (int p = 0; p < 256; ++p) start[blk_packed[p]]++;
+        uint16_t acc = 0;
+        for (int b = 0; b < 256; ++b) {
+            const uint16_t c = start[b];
+            start[b] = acc;
+            acc = (uint16_t) (acc + c);
+        }
+        for (uint16_t p = 0; p < 256; ++p) {
+            order[start[blk_packed[p]]++] = p;
+        }
+    }
+}
+
+#ifdef __cplusplus
+}
+#endif
+"""
+
+# --------------------------------------------------------------------------
+# 2. ggml/src/ggml.c: registry implementation in libggml-base
+# --------------------------------------------------------------------------
+
+GGML_C_INCLUDE = f"""
+// {BEGIN}: hpcq seg header
+#include "hpcq-seg.h"
+// {END}
+"""
+
+GGML_C_API = f"""
+// {BEGIN}: compact nibble registry (no order/F16 expansion)
+static int hpcq_seg_g_count = 0;
+static struct hpcq_seg_entry hpcq_seg_g_entries[HPCQ_SEG_MAX_ENTRIES];
+
+void ggml_hpcq_seg_bind_raw(
+        const void * payload, int64_t n_bytes,
+        const uint8_t * hqs1, size_t hqs1_len, int64_t block_bytes) {{
+    if (payload == NULL || hqs1 == NULL || n_bytes <= 0 || block_bytes <= 0) return;
+    if (hqs1_len < 10 || hqs1[0] != 'H' || hqs1[1] != 'Q' || hqs1[2] != 'S' || hqs1[3] != '1') return;
+    const uint8_t enc = hqs1[5];
+    const uint32_t n_labels = (uint32_t) hqs1[6] | ((uint32_t) hqs1[7] << 8) |
+                              ((uint32_t) hqs1[8] << 16) | ((uint32_t) hqs1[9] << 24);
+    if (n_labels == 0 || (n_labels % 256) != 0) return;
+    const int64_t n_blocks = (int64_t) (n_labels / 256);
+    if (n_blocks * block_bytes != n_bytes) return;
+    const size_t stride = (enc == 0) ? 128u : 256u;
+    const size_t need = (size_t) n_blocks * stride;
+    if (hqs1_len < 10u + need) return;
+
+    const uintptr_t base = (uintptr_t) payload;
+    const uintptr_t end  = base + (uintptr_t) n_bytes;
+    for (int i = 0; i < hpcq_seg_g_count; ++i) {{
+        if (hpcq_seg_g_entries[i].base == base) {{
+            if (hpcq_seg_g_entries[i].packed) free(hpcq_seg_g_entries[i].packed);
+            uint8_t * copy = (uint8_t *) malloc(need);
+            if (!copy) return;
+            memcpy(copy, hqs1 + 10, need);
+            hpcq_seg_g_entries[i].end         = end;
+            hpcq_seg_g_entries[i].block_bytes = block_bytes;
+            hpcq_seg_g_entries[i].n_blocks    = n_blocks;
+            hpcq_seg_g_entries[i].enc         = enc;
+            hpcq_seg_g_entries[i].packed      = copy;
+            hpcq_seg_g_entries[i].dev_packed  = NULL;
+            return;
+        }}
+    }}
+    if (hpcq_seg_g_count >= HPCQ_SEG_MAX_ENTRIES) return;
+    uint8_t * copy = (uint8_t *) malloc(need);
+    if (!copy) return;
+    memcpy(copy, hqs1 + 10, need);
+    struct hpcq_seg_entry * e = &hpcq_seg_g_entries[hpcq_seg_g_count++];
+    e->base        = base;
+    e->end         = end;
+    e->block_bytes = block_bytes;
+    e->n_blocks    = n_blocks;
+    e->enc         = enc;
+    e->packed      = copy;
+    e->dev_packed  = NULL;
+}}
+
+const struct hpcq_seg_entry * ggml_hpcq_seg_find(
+        const void * ptr, int64_t block_bytes, int64_t * out_block_idx) {{
+    const uintptr_t p = (uintptr_t) ptr;
+    const int n = hpcq_seg_g_count;
+    for (int i = 0; i < n; ++i) {{
+        const struct hpcq_seg_entry * e = &hpcq_seg_g_entries[i];
+        if (p >= e->base && p < e->end && e->block_bytes == block_bytes) {{
+            if (out_block_idx) {{
+                *out_block_idx = (int64_t) ((p - e->base) / (uintptr_t) block_bytes);
+            }}
+            return e;
+        }}
+    }}
+    return NULL;
+}}
+
+void ggml_hpcq_seg_set_dev_packed(const struct hpcq_seg_entry * entry, void * dev_packed) {{
+    if (!entry) return;
+    const int n = hpcq_seg_g_count;
+    for (int i = 0; i < n; ++i) {{
+        if (&hpcq_seg_g_entries[i] == entry) {{
+            hpcq_seg_g_entries[i].dev_packed = dev_packed;
+            return;
+        }}
+    }}
+}}
+// {END}
+"""
+
+# --------------------------------------------------------------------------
+# 3. ggml/src/ggml-quants.c: stack-based scatter in dequantize_row_q2_K / q3_K
+# --------------------------------------------------------------------------
+
+Q2K_DEQ_OLD = """    const int nb = k / QK_K;
+
+    for (int i = 0; i < nb; i++) {
+
+        const float d = GGML_FP16_TO_FP32(x[i].d);
+        const float min = GGML_FP16_TO_FP32(x[i].dmin);
+
+        const uint8_t * q = x[i].qs;
+
+        int is = 0;
+        float dl, ml;
+        for (int n = 0; n < QK_K; n += 128) {
+            int shift = 0;
+            for (int j = 0; j < 4; ++j) {
+
+                uint8_t sc = x[i].scales[is++];
+                dl = d * (sc & 0xF); ml = min * (sc >> 4);
+                for (int l = 0; l < 16; ++l) *y++ = dl * ((int8_t)((q[l] >> shift) & 3)) - ml;
+
+                sc = x[i].scales[is++];
+                dl = d * (sc & 0xF); ml = min * (sc >> 4);
+                for (int l = 0; l < 16; ++l) *y++ = dl * ((int8_t)((q[l+16] >> shift) & 3)) - ml;
+
+                shift += 2;
+            }
+            q += 32;
+        }
+    }
+}"""
+
+Q2K_DEQ_NEW = f"""    const int nb = k / QK_K;
+
+    // {BEGIN}: stack-based per-block scatter for segmented Q2_K
+    int64_t hpcq_b0 = 0;
+    const struct hpcq_seg_entry * hpcq_seg = ggml_hpcq_seg_find(x, (int64_t) sizeof(block_q2_K), &hpcq_b0);
+    const size_t hpcq_stride = (hpcq_seg && hpcq_seg->enc != 0) ? 256u : 128u;
+    uint16_t hpcq_ord[QK_K];
+    // {END}
+
+    for (int i = 0; i < nb; i++) {{
+
+        float * y_out = y + i * QK_K;
+        if (hpcq_seg) {{
+            hpcq_seg_make_order_block(hpcq_seg->packed + (size_t) (hpcq_b0 + i) * hpcq_stride, hpcq_seg->enc, hpcq_ord);
+        }}
+
+        const float d = GGML_FP16_TO_FP32(x[i].d);
+        const float min = GGML_FP16_TO_FP32(x[i].dmin);
+
+        const uint8_t * q = x[i].qs;
+
+        int is = 0;
+        int base = 0;
+        float dl, ml;
+        for (int n = 0; n < QK_K; n += 128) {{
+            int shift = 0;
+            for (int j = 0; j < 4; ++j) {{
+
+                uint8_t sc = x[i].scales[is++];
+                dl = d * (sc & 0xF); ml = min * (sc >> 4);
+                for (int l = 0; l < 16; ++l) {{
+                    const float v = dl * ((int8_t)((q[l] >> shift) & 3)) - ml;
+                    if (hpcq_seg) y_out[hpcq_ord[base + l]] = v; else y_out[base + l] = v;
+                }}
+                base += 16;
+
+                sc = x[i].scales[is++];
+                dl = d * (sc & 0xF); ml = min * (sc >> 4);
+                for (int l = 0; l < 16; ++l) {{
+                    const float v = dl * ((int8_t)((q[l+16] >> shift) & 3)) - ml;
+                    if (hpcq_seg) y_out[hpcq_ord[base + l]] = v; else y_out[base + l] = v;
+                }}
+                base += 16;
+
+                shift += 2;
+            }}
+            q += 32;
+        }}
+    }}
+}}"""
+
+Q3K_DEQ_OLD = """    for (int i = 0; i < nb; i++) {
+
+        const float d_all = GGML_FP16_TO_FP32(x[i].d);
+
+        const uint8_t * GGML_RESTRICT q = x[i].qs;
+        const uint8_t * GGML_RESTRICT hm = x[i].hmask;
+        uint8_t m = 1;
+
+        memcpy(aux, x[i].scales, 12);
+        uint32_t tmp = aux[2];
+        aux[2] = ((aux[0] >> 4) & kmask2) | (((tmp >> 4) & kmask1) << 4);
+        aux[3] = ((aux[1] >> 4) & kmask2) | (((tmp >> 6) & kmask1) << 4);
+        aux[0] = (aux[0] & kmask2) | (((tmp >> 0) & kmask1) << 4);
+        aux[1] = (aux[1] & kmask2) | (((tmp >> 2) & kmask1) << 4);
+
+        int is = 0;
+        float dl;
+        for (int n = 0; n < QK_K; n += 128) {
+            int shift = 0;
+            for (int j = 0; j < 4; ++j) {
+
+                dl = d_all * (scales[is++] - 32);
+                for (int l = 0; l < 16; ++l) {
+                    *y++ = dl * ((int8_t)((q[l+ 0] >> shift) & 3) - ((hm[l+ 0] & m) ? 0 : 4));
+                }
+
+                dl = d_all * (scales[is++] - 32);
+                for (int l = 0; l < 16; ++l) {
+                    *y++ = dl * ((int8_t)((q[l+16] >> shift) & 3) - ((hm[l+16] & m) ? 0 : 4));
+                }
+
+                shift += 2;
+                m <<= 1;
+            }
+            q += 32;
+        }
+
+    }
+}"""
+
+Q3K_DEQ_NEW = f"""    // {BEGIN}: stack-based per-block scatter for segmented Q3_K
+    int64_t hpcq_b0 = 0;
+    const struct hpcq_seg_entry * hpcq_seg = ggml_hpcq_seg_find(x, (int64_t) sizeof(block_q3_K), &hpcq_b0);
+    const size_t hpcq_stride = (hpcq_seg && hpcq_seg->enc != 0) ? 256u : 128u;
+    uint16_t hpcq_ord[QK_K];
+    // {END}
+
+    for (int i = 0; i < nb; i++) {{
+
+        float * y_out = y + i * QK_K;
+        if (hpcq_seg) {{
+            hpcq_seg_make_order_block(hpcq_seg->packed + (size_t) (hpcq_b0 + i) * hpcq_stride, hpcq_seg->enc, hpcq_ord);
+        }}
+
+        const float d_all = GGML_FP16_TO_FP32(x[i].d);
+
+        const uint8_t * GGML_RESTRICT q = x[i].qs;
+        const uint8_t * GGML_RESTRICT hm = x[i].hmask;
+        uint8_t m = 1;
+
+        memcpy(aux, x[i].scales, 12);
+        uint32_t tmp = aux[2];
+        aux[2] = ((aux[0] >> 4) & kmask2) | (((tmp >> 4) & kmask1) << 4);
+        aux[3] = ((aux[1] >> 4) & kmask2) | (((tmp >> 6) & kmask1) << 4);
+        aux[0] = (aux[0] & kmask2) | (((tmp >> 0) & kmask1) << 4);
+        aux[1] = (aux[1] & kmask2) | (((tmp >> 2) & kmask1) << 4);
+
+        int is = 0;
+        int base = 0;
+        float dl;
+        for (int n = 0; n < QK_K; n += 128) {{
+            int shift = 0;
+            for (int j = 0; j < 4; ++j) {{
+
+                dl = d_all * (scales[is++] - 32);
+                for (int l = 0; l < 16; ++l) {{
+                    const float v = dl * ((int8_t)((q[l+ 0] >> shift) & 3) - ((hm[l+ 0] & m) ? 0 : 4));
+                    if (hpcq_seg) y_out[hpcq_ord[base + l]] = v; else y_out[base + l] = v;
+                }}
+                base += 16;
+
+                dl = d_all * (scales[is++] - 32);
+                for (int l = 0; l < 16; ++l) {{
+                    const float v = dl * ((int8_t)((q[l+16] >> shift) & 3) - ((hm[l+16] & m) ? 0 : 4));
+                    if (hpcq_seg) y_out[hpcq_ord[base + l]] = v; else y_out[base + l] = v;
+                }}
+                base += 16;
+
+                shift += 2;
+                m <<= 1;
+            }}
+            q += 32;
+        }}
+
+    }}
+}}"""
+
+# --------------------------------------------------------------------------
+# 4. ggml/src/ggml-cpu/ggml-cpu.c: stack-gathered SIMD vec_dot wrappers
+# --------------------------------------------------------------------------
+
+CPU_VECDOT_WRAPPERS = f"""
+// {BEGIN}: stack-gathered SIMD vec_dot wrappers for segmented Q2_K / Q3_K
+#include "hpcq-seg.h"
+
+static void hpcq_vec_dot_q2_K_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) {{
+    int64_t b0 = 0;
+    const struct hpcq_seg_entry * seg = ggml_hpcq_seg_find(vx, (int64_t) sizeof(block_q2_K), &b0);
+    if (seg == NULL) {{
+        ggml_vec_dot_q2_K_q8_K(n, s, bs, vx, bx, vy, by, nrc);
+        return;
+    }}
+    const block_q2_K * GGML_RESTRICT x = (const block_q2_K *) vx;
+    const block_q8_K * GGML_RESTRICT y = (const block_q8_K *) vy;
+    const int nb = n / QK_K;
+    const size_t stride = (seg->enc != 0) ? 256u : 128u;
+    uint16_t ord[QK_K];
+    block_q8_K y_perm;
+    float sumf = 0.0f;
+    for (int i = 0; i < nb; ++i) {{
+        hpcq_seg_make_order_block(seg->packed + (size_t) (b0 + i) * stride, seg->enc, ord);
+        y_perm.d = y[i].d;
+        for (int g = 0; g < 16; ++g) {{
+            int16_t bsum = 0;
+            for (int l = 0; l < 16; ++l) {{
+                const int8_t qv = y[i].qs[ord[16 * g + l]];
+                y_perm.qs[16 * g + l] = qv;
+                bsum = (int16_t) (bsum + qv);
+            }}
+            y_perm.bsums[g] = bsum;
+        }}
+        float blk_s = 0.0f;
+        ggml_vec_dot_q2_K_q8_K(QK_K, &blk_s, bs, &x[i], bx, &y_perm, by, 1);
+        sumf += blk_s;
+    }}
+    *s = sumf;
+}}
+
+static void hpcq_vec_dot_q3_K_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) {{
+    int64_t b0 = 0;
+    const struct hpcq_seg_entry * seg = ggml_hpcq_seg_find(vx, (int64_t) sizeof(block_q3_K), &b0);
+    if (seg == NULL) {{
+        ggml_vec_dot_q3_K_q8_K(n, s, bs, vx, bx, vy, by, nrc);
+        return;
+    }}
+    const block_q3_K * GGML_RESTRICT x = (const block_q3_K *) vx;
+    const block_q8_K * GGML_RESTRICT y = (const block_q8_K *) vy;
+    const int nb = n / QK_K;
+    const size_t stride = (seg->enc != 0) ? 256u : 128u;
+    uint16_t ord[QK_K];
+    block_q8_K y_perm;
+    float sumf = 0.0f;
+    for (int i = 0; i < nb; ++i) {{
+        hpcq_seg_make_order_block(seg->packed + (size_t) (b0 + i) * stride, seg->enc, ord);
+        y_perm.d = y[i].d;
+        for (int g = 0; g < 16; ++g) {{
+            int16_t bsum = 0;
+            for (int l = 0; l < 16; ++l) {{
+                const int8_t qv = y[i].qs[ord[16 * g + l]];
+                y_perm.qs[16 * g + l] = qv;
+                bsum = (int16_t) (bsum + qv);
+            }}
+            y_perm.bsums[g] = bsum;
+        }}
+        float blk_s = 0.0f;
+        ggml_vec_dot_q3_K_q8_K(QK_K, &blk_s, bs, &x[i], bx, &y_perm, by, 1);
+        sumf += blk_s;
+    }}
+    *s = sumf;
+}}
+// {END}
+"""
+
+CPU_TRAITS_OLD = """    [GGML_TYPE_Q2_K] = {
+        .from_float               = quantize_row_q2_K,
+        .vec_dot                  = ggml_vec_dot_q2_K_q8_K,
+        .vec_dot_type             = GGML_TYPE_Q8_K,
+        .nrows                    = 1,
+    },
+    [GGML_TYPE_Q3_K] = {
+        .from_float               = quantize_row_q3_K,
+        .vec_dot                  = ggml_vec_dot_q3_K_q8_K,
+        .vec_dot_type             = GGML_TYPE_Q8_K,
+        .nrows                    = 1,
+    },"""
+
+CPU_TRAITS_NEW = f"""    [GGML_TYPE_Q2_K] = {{
+        .from_float               = quantize_row_q2_K,
+        .vec_dot                  = hpcq_vec_dot_q2_K_q8_K, // {BEGIN}
+        .vec_dot_type             = GGML_TYPE_Q8_K,
+        .nrows                    = 1,
+    }},
+    [GGML_TYPE_Q3_K] = {{
+        .from_float               = quantize_row_q3_K,
+        .vec_dot                  = hpcq_vec_dot_q3_K_q8_K, // {END}
+        .vec_dot_type             = GGML_TYPE_Q8_K,
+        .nrows                    = 1,
+    }},"""
+
+# --------------------------------------------------------------------------
+# 5. ggml/src/ggml-cuda/convert.cu & ggml-cuda.cu: shared-mem scatter on GPU
+# --------------------------------------------------------------------------
+
+CUDA_CONVERT_OLD_KERNELS = """template<typename dst_t>
+static __global__ void dequantize_block_q2_K(const void * __restrict__ vx, dst_t * __restrict__ yy) {
+    const int64_t i = blockIdx.x;
+
+    dequantize_q2_K(vx, i, yy + i*QK_K, threadIdx.x);
+}
+
+template<typename dst_t>
+static __global__ void dequantize_block_q3_K(const void * __restrict__ vx, dst_t * __restrict__ yy) {
+    const int64_t i = blockIdx.x;
+
+    dequantize_q3_K(vx, i, yy + i*QK_K, threadIdx.x);
+}"""
+
+CUDA_CONVERT_NEW_KERNELS = f"""// {BEGIN}: shared-memory per-block scatter for segmented Q2_K / Q3_K on CUDA
+#include "hpcq-seg.h"
+
+static inline const uint8_t * hpcq_cuda_get_dev_packed(
+        const void * vx, int64_t bs, int64_t * out_b0, uint8_t * out_enc) {{
+    const struct hpcq_seg_entry * seg = ggml_hpcq_seg_find(vx, bs, out_b0);
+    if (!seg) return nullptr;
+    *out_enc = seg->enc;
+    if (seg->dev_packed == nullptr) {{
+        const size_t stride = (seg->enc != 0) ? 256u : 128u;
+        const size_t nbytes = (size_t) seg->n_blocks * stride;
+        void * d_ptr = nullptr;
+        if (cudaMalloc(&d_ptr, nbytes) == cudaSuccess && d_ptr != nullptr) {{
+            if (cudaMemcpy(d_ptr, seg->packed, nbytes, cudaMemcpyHostToDevice) == cudaSuccess) {{
+                ggml_hpcq_seg_set_dev_packed(seg, d_ptr);
+            }} else {{
+                cudaFree(d_ptr);
+                d_ptr = nullptr;
+            }}
+        }}
+    }}
+    return (const uint8_t *) seg->dev_packed;
+}}
+
+template<typename dst_t>
+static __global__ void dequantize_block_q2_K(
+        const void * __restrict__ vx, dst_t * __restrict__ yy,
+        const uint8_t * __restrict__ d_packed = nullptr, int64_t b0 = 0, uint8_t enc = 0) {{
+    const int64_t i = blockIdx.x;
+    if (d_packed == nullptr) {{
+        dequantize_q2_K(vx, i, yy + i*QK_K, threadIdx.x);
+        return;
+    }}
+    __shared__ dst_t sh_y[QK_K];
+    __shared__ uint16_t sh_ord[QK_K];
+    dequantize_q2_K(vx, i, sh_y, threadIdx.x);
+    if (threadIdx.x == 0) {{
+        const size_t stride = (enc != 0) ? 256u : 128u;
+        hpcq_seg_make_order_block(d_packed + (size_t) (b0 + i) * stride, enc, sh_ord);
+    }}
+    __syncthreads();
+    dst_t * y_out = yy + i * QK_K;
+    for (int s = threadIdx.x; s < QK_K; s += 64) {{
+        y_out[sh_ord[s]] = sh_y[s];
+    }}
+}}
+
+template<typename dst_t>
+static __global__ void dequantize_block_q3_K(
+        const void * __restrict__ vx, dst_t * __restrict__ yy,
+        const uint8_t * __restrict__ d_packed = nullptr, int64_t b0 = 0, uint8_t enc = 0) {{
+    const int64_t i = blockIdx.x;
+    if (d_packed == nullptr) {{
+        dequantize_q3_K(vx, i, yy + i*QK_K, threadIdx.x);
+        return;
+    }}
+    __shared__ dst_t sh_y[QK_K];
+    __shared__ uint16_t sh_ord[QK_K];
+    dequantize_q3_K(vx, i, sh_y, threadIdx.x);
+    if (threadIdx.x == 0) {{
+        const size_t stride = (enc != 0) ? 256u : 128u;
+        hpcq_seg_make_order_block(d_packed + (size_t) (b0 + i) * stride, enc, sh_ord);
+    }}
+    __syncthreads();
+    dst_t * y_out = yy + i * QK_K;
+    for (int s = threadIdx.x; s < QK_K; s += 64) {{
+        y_out[sh_ord[s]] = sh_y[s];
+    }}
+}}
+// {END}"""
+
+CUDA_CONVERT_OLD_LAUNCH = """template<typename dst_t>
+static void dequantize_row_q2_K_cuda(const void * vx, dst_t * y, const int64_t k, cudaStream_t stream) {
+    const int nb = k / QK_K;
+    dequantize_block_q2_K<<<nb, 64, 0, stream>>>(vx, y);
+}
+
+template<typename dst_t>
+static void dequantize_row_q3_K_cuda(const void * vx, dst_t * y, const int64_t k, cudaStream_t stream) {
+    const int nb = k / QK_K;
+    dequantize_block_q3_K<<<nb, 64, 0, stream>>>(vx, y);
+}"""
+
+CUDA_CONVERT_NEW_LAUNCH = f"""template<typename dst_t>
+static void dequantize_row_q2_K_cuda(const void * vx, dst_t * y, const int64_t k, cudaStream_t stream) {{
+    const int nb = k / QK_K;
+    int64_t b0 = 0; uint8_t enc = 0;
+    const uint8_t * d_packed = hpcq_cuda_get_dev_packed(vx, (int64_t) sizeof(block_q2_K), &b0, &enc);
+    dequantize_block_q2_K<<<nb, 64, 0, stream>>>(vx, y, d_packed, b0, enc);
+}}
+
+template<typename dst_t>
+static void dequantize_row_q3_K_cuda(const void * vx, dst_t * y, const int64_t k, cudaStream_t stream) {{
+    const int nb = k / QK_K;
+    int64_t b0 = 0; uint8_t enc = 0;
+    const uint8_t * d_packed = hpcq_cuda_get_dev_packed(vx, (int64_t) sizeof(block_q3_K), &b0, &enc);
+    dequantize_block_q3_K<<<nb, 64, 0, stream>>>(vx, y, d_packed, b0, enc);
+}}"""
+
+CUDA_MULMAT_OLD = """    if (ggml_cuda_should_use_mmvq(src0->type, cc, ne11)) {
+        ggml_cuda_mul_mat_vec_q(ctx, src0, src1, nullptr, dst);
+        return;
+    }
+    if (ggml_cuda_should_use_mmq(src0->type, cc, ne11, /*n_experts =*/ 0)) {
+        ggml_cuda_mul_mat_q(ctx, src0, src1, nullptr, dst);
+        return;
+    }
+    ggml_cuda_mul_mat_cublas(ctx, src0, src1, dst);"""
+
+CUDA_MULMAT_NEW = f"""    // {BEGIN}: route segmented Q2_K/Q3_K to pooled-scratch cublas dequant path
+    if ((src0->type == GGML_TYPE_Q2_K || src0->type == GGML_TYPE_Q3_K) &&
+        ggml_hpcq_seg_find(src0->data, (int64_t) ggml_type_size(src0->type), nullptr) != nullptr) {{
+        ggml_cuda_mul_mat_cublas(ctx, src0, src1, dst);
+        return;
+    }}
+    // {END}
+    if (ggml_cuda_should_use_mmvq(src0->type, cc, ne11)) {{
+        ggml_cuda_mul_mat_vec_q(ctx, src0, src1, nullptr, dst);
+        return;
+    }}
+    if (ggml_cuda_should_use_mmq(src0->type, cc, ne11, /*n_experts =*/ 0)) {{
+        ggml_cuda_mul_mat_q(ctx, src0, src1, nullptr, dst);
+        return;
+    }}
+    ggml_cuda_mul_mat_cublas(ctx, src0, src1, dst);"""
+
+
+def _read(p: Path) -> str:
+    return p.read_text(encoding="utf-8")
+
+
+def _write(p: Path, s: str) -> None:
+    p.write_text(s, encoding="utf-8")
+
+
+def _backup(p: Path) -> None:
+    bak = p.with_suffix(p.suffix + ".hpcq-bak")
+    if not bak.exists():
+        shutil.copy2(p, bak)
+
+
+def _require_once(text: str, needle: str, what: str) -> None:
+    n = text.count(needle)
+    if n != 1:
+        raise SystemExit(f"anchor for {what}: found {n} occurrence(s), expected 1")
+
+
+def _insert_after(text: str, anchor: str, block: str, what: str) -> str:
+    _require_once(text, anchor, what)
+    i = text.index(anchor) + len(anchor)
+    return text[:i] + block + text[i:]
+
+
+def _replace_block(text: str, old: str, new: str, what: str) -> str:
+    _require_once(text, old, what)
+    return text.replace(old, new, 1)
+
+
+# --------------------------------------------------------------------------
+# Patch steps
+# --------------------------------------------------------------------------
+
+def step_header(root: Path) -> None:
+    _write(root / HEADER_REL, HEADER)
+    old_src_h = root / "ggml" / "src" / "hpcq-seg.h"
+    if old_src_h.exists():
+        old_src_h.unlink()
+
+
+def step_ggml_c(root: Path) -> None:
+    p = root / "ggml" / "src" / "ggml.c"
+    s = _read(p)
+    if BEGIN in s:
+        return
+    _backup(p)
+    anchor = '#include "ggml-impl.h"'
+    s = _insert_after(s, anchor, "\n" + GGML_C_INCLUDE, "ggml.c include")
+    marker = "////////////////////////////////////////////////////////////////////////////////"
+    i = s.rindex(marker)
+    s = s[:i] + GGML_C_API + "\n\n" + s[i:]
+    _write(p, s)
+
+
+def step_ggml_quants(root: Path) -> None:
+    p = root / "ggml" / "src" / "ggml-quants.c"
+    s = _read(p)
+    if BEGIN in s:
+        return
+    _backup(p)
+    anchor = '#include "ggml-quants.h"'
+    s = _insert_after(s, anchor, '\n#include "hpcq-seg.h"\n', "ggml-quants.c include")
+    s = _replace_block(s, Q2K_DEQ_OLD, Q2K_DEQ_NEW, "dequantize_row_q2_K")
+    s = _replace_block(s, Q3K_DEQ_OLD, Q3K_DEQ_NEW, "dequantize_row_q3_K")
+    _write(p, s)
+
+
+def step_ggml_cpu(root: Path) -> None:
+    p = root / "ggml" / "src" / "ggml-cpu" / "ggml-cpu.c"
+    s = _read(p)
+    if BEGIN in s:
+        return
+    _backup(p)
+    anchor = "static const struct ggml_type_traits_cpu type_traits_cpu[GGML_TYPE_COUNT] = {"
+    _require_once(s, anchor, "ggml-cpu.c type_traits_cpu")
+    s = s.replace(anchor, CPU_VECDOT_WRAPPERS + "\n" + anchor, 1)
+    s = _replace_block(s, CPU_TRAITS_OLD, CPU_TRAITS_NEW, "ggml-cpu.c Q2_K/Q3_K traits")
+    _write(p, s)
+
+
+def step_ggml_cuda(root: Path) -> None:
+    conv = root / "ggml" / "src" / "ggml-cuda" / "convert.cu"
+    s = _read(conv)
+    if BEGIN not in s:
+        _backup(conv)
+        s = _replace_block(s, CUDA_CONVERT_OLD_KERNELS, CUDA_CONVERT_NEW_KERNELS, "convert.cu kernels")
+        s = _replace_block(s, CUDA_CONVERT_OLD_LAUNCH, CUDA_CONVERT_NEW_LAUNCH, "convert.cu launch")
+        _write(conv, s)
+
+    cu = root / "ggml" / "src" / "ggml-cuda" / "ggml-cuda.cu"
+    s2 = _read(cu)
+    if BEGIN not in s2:
+        _backup(cu)
+        anchor = '#include "ggml-cuda.h"'
+        s2 = _insert_after(s2, anchor, '\n#include "hpcq-seg.h"', "ggml-cuda.cu include")
+        s2 = _replace_block(s2, CUDA_MULMAT_OLD, CUDA_MULMAT_NEW, "ggml-cuda.cu mul_mat")
+        _write(cu, s2)
+
+
+def step_loader(root: Path) -> None:
+    cpp = root / "src" / "llama-model-loader.cpp"
+    s = _read(cpp)
+    if BEGIN in s:
+        return
+    _backup(cpp)
+
+    anchor = '#include "llama-model-loader.h"'
+    s = _insert_after(s, anchor, '\n#include "hpcq-seg.h"', "loader include")
+
+    # prevent CPU repack buffer selection on segmented tensors so they stay in
+    # native block_q2_K / block_q3_K layout for hpcq_vec_dot_q2_K_q8_K
+    anchor_buft = """        if (!buft) {
+            buft = select_weight_buft(hparams, t_meta, op, info.src1_noncont, buft_list);
+            if (!buft) {
+                throw std::runtime_error(format("failed to find a compatible buffer type for tensor %s", tn.str().c_str()));
+            }
+        }"""
+    new_buft = anchor_buft + f"""
+        // {BEGIN}: keep segmented CPU tensors in standard host/mmap buffer (no 8x8 interleave repack)
+        if (buft && metadata) {{
+            const std::string seg_key = std::string(HPCQ_SEG_PREFIX) + tn.str();
+            if (gguf_find_key(metadata, seg_key.c_str()) >= 0) {{
+                auto * bdev = ggml_backend_buft_get_device(buft);
+                if (bdev && ggml_backend_dev_type(bdev) == GGML_BACKEND_DEVICE_TYPE_CPU &&
+                    buft != ggml_backend_cpu_buffer_type()) {{
+                    buft = ggml_backend_cpu_buffer_type();
+                }}
+            }}
+        }}
+        // {END}"""
+    s = _replace_block(s, anchor_buft, new_buft, "loader select_weight_buft")
+
+    # bind compact HQS1 nibble stream directly to the tensor's payload pointer
+    anchor_tail = "        size_done += n_size;\n    }"
+    new_tail = f"""        // {BEGIN}: bind compact HQS1 nibble stream to payload range (no weight/order expansion)
+        if (metadata && cur->data != NULL && (cur->type == GGML_TYPE_Q2_K || cur->type == GGML_TYPE_Q3_K)) {{
+            const std::string seg_key = std::string(HPCQ_SEG_PREFIX) + ggml_get_name(cur);
+            const int64_t seg_ki = gguf_find_key(metadata, seg_key.c_str());
+            if (seg_ki >= 0 && gguf_get_kv_type(metadata, seg_ki) == GGUF_TYPE_ARRAY &&
+                gguf_get_arr_type(metadata, seg_ki) == GGUF_TYPE_UINT8) {{
+                ggml_hpcq_seg_bind_raw(
+                    cur->data,
+                    (int64_t) ggml_nbytes(cur),
+                    (const uint8_t *) gguf_get_arr_data(metadata, seg_ki),
+                    gguf_get_arr_n(metadata, seg_ki),
+                    (int64_t) ggml_type_size(cur->type));
+            }}
+        }}
+        // {END}
+        size_done += n_size;
+    }}"""
+    s = _replace_block(s, anchor_tail, new_tail, "loader load_all_data tail")
+    _write(cpp, s)
+
+
+STEPS = [
+    ("header", step_header),
+    ("ggml.c", step_ggml_c),
+    ("ggml-quants.c", step_ggml_quants),
+    ("ggml-cpu.c", step_ggml_cpu),
+    ("ggml-cuda", step_ggml_cuda),
+    ("loader", step_loader),
+]
+
+TOUCHED_FILES = (
+    "ggml/src/ggml.c",
+    "ggml/src/ggml-quants.c",
+    "ggml/src/ggml-cpu/ggml-cpu.c",
+    "ggml/src/ggml-cuda/convert.cu",
+    "ggml/src/ggml-cuda/ggml-cuda.cu",
+    "src/llama-model-loader.cpp",
+    "src/llama-model-loader.h",
+)
+
+
+def check(root: Path) -> int:
+    ok = True
+    checks = [
+        ("ggml.c", root / "ggml/src/ggml.c", ['#include "ggml-impl.h"']),
+        ("ggml-quants.c", root / "ggml/src/ggml-quants.c", [Q2K_DEQ_OLD, Q3K_DEQ_OLD]),
+        ("ggml-cpu.c", root / "ggml/src/ggml-cpu/ggml-cpu.c", [CPU_TRAITS_OLD]),
+        ("convert.cu", root / "ggml/src/ggml-cuda/convert.cu", [CUDA_CONVERT_OLD_KERNELS, CUDA_CONVERT_OLD_LAUNCH]),
+        ("ggml-cuda.cu", root / "ggml/src/ggml-cuda/ggml-cuda.cu", [CUDA_MULMAT_OLD]),
+        ("loader.cpp", root / "src/llama-model-loader.cpp", ["        size_done += n_size;\n    }"]),
+    ]
+    for label, p, needles in checks:
+        s = _read(p)
+        if BEGIN in s:
+            print(f"{label:15s}: already patched")
+        else:
+            present = all(n in s for n in needles)
+            print(f"{label:15s}: anchors {'ok' if present else 'MISSING'}")
+            ok &= present
+    return 0 if ok else 1
+
+
+def apply(root: Path) -> int:
+    for name, fn in STEPS:
+        fn(root)
+        print(f"patched {name}")
+    print("done.")
+    return 0
+
+
+def revert(root: Path) -> int:
+    for rel in TOUCHED_FILES:
+        p = root / rel
+        bak = p.with_suffix(p.suffix + ".hpcq-bak")
+        if bak.exists():
+            shutil.copy2(bak, p)
+            print(f"restored {rel}")
+    for hrel in (HEADER_REL, "ggml/src/hpcq-seg.h"):
+        h = root / hrel
+        if h.exists():
+            h.unlink()
+            print(f"removed {hrel}")
+    return 0
+
+
+SELFTEST_CPP = r"""
+#include <cstdio>
+#include <cstdlib>
+#include <cstring>
+#include <vector>
+
+#define GGML_API
+#include "hpcq-seg.h"
+
+int main() {
+    // Verify that hpcq_seg_make_order_block produces the exact stable argsort
+    // (slot -> original_position) for both nibble (enc=0) and byte (enc=1) streams,
+    // and that gathering y_perm by ord[s] matches the scattered dot product.
+    uint8_t labels[256];
+    uint8_t nibbles[128];
+    for (int p = 0; p < 256; ++p) {
+        labels[p] = (uint8_t) ((p * 7 + 3) % 16);
+    }
+    for (int i = 0; i < 128; ++i) {
+        nibbles[i] = (uint8_t) ((labels[2 * i] & 0x0F) | ((labels[2 * i + 1] & 0x0F) << 4));
+    }
+    uint16_t ord0[256], ord1[256];
+    hpcq_seg_make_order_block(nibbles, 0, ord0);
+    hpcq_seg_make_order_block(labels,  1, ord1);
+    for (int s = 0; s < 256; ++s) {
+        if (ord0[s] != ord1[s]) return 1;
+        if (s > 0) {
+            uint8_t l_prev = labels[ord0[s - 1]];
+            uint8_t l_curr = labels[ord0[s]];
+            if (l_prev > l_curr) return 2;
+            if (l_prev == l_curr && ord0[s - 1] >= ord0[s]) return 3;
+        }
+    }
+    // Check dot-product identity: sum_s x_slot[s] * y_orig[ord[s]] == sum_p x_orig[p] * y_orig[p]
+    int x_orig[256], x_slot[256], y_orig[256], y_perm[256];
+    for (int p = 0; p < 256; ++p) {
+        x_orig[p] = (p * 13 - 97);
+        y_orig[p] = (p * 5  - 41);
+    }
+    for (int s = 0; s < 256; ++s) {
+        x_slot[s] = x_orig[ord0[s]];
+        y_perm[s] = y_orig[ord0[s]];
+    }
+    long dot_orig = 0, dot_perm = 0;
+    for (int i = 0; i < 256; ++i) {
+        dot_orig += (long) x_orig[i] * (long) y_orig[i];
+        dot_perm += (long) x_slot[i] * (long) y_perm[i];
+    }
+    if (dot_orig != dot_perm) return 4;
+    std::printf("selftest ok (dot=%ld)\n", dot_perm);
+    return 0;
+}
+"""
+
+
+def selftest() -> int:
+    with tempfile.TemporaryDirectory() as td:
+        td = Path(td)
+        (td / "ggml.h").write_text("#pragma once\n", encoding="utf-8")
+        (td / "hpcq-seg.h").write_text(HEADER, encoding="utf-8")
+        src = td / "selftest.cpp"
+        src.write_text(SELFTEST_CPP, encoding="utf-8")
+        exe = td / "selftest"
+        for cmd in (
+            ["g++", "-O2", "-std=c++11", "-I", str(td), str(src), "-o", str(exe)],
+            [str(exe)],
+        ):
+            r = subprocess.run(cmd, capture_output=True, text=True)
+            if r.returncode != 0:
+                print(r.stderr, file=sys.stderr)
+                return 1
+            if r.stdout:
+                print(r.stdout.strip())
+    return 0
+
+
+def main() -> int:
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--llama", default="/home/node/Desktop/Dev/llama.cpp")
+    g = ap.add_mutually_exclusive_group(required=True)
+    g.add_argument("--check", action="store_true")
+    g.add_argument("--apply", action="store_true")
+    g.add_argument("--revert", action="store_true")
+    g.add_argument("--selftest", action="store_true")
+    a = ap.parse_args()
+    if a.selftest:
+        return selftest()
+    root = Path(a.llama)
+    if not (root / "ggml" / "src" / "ggml-quants.c").exists():
+        raise SystemExit(f"not a llama.cpp tree: {root}")
+    if a.check:
+        return check(root)
+    if a.apply:
+        return apply(root)
+    return revert(root)
+
+
+if __name__ == "__main__":
+    sys.exit(main())