Upload patch_llamacpp_hpcq_seg.py

patch_llamacpp_hpcq_seg.py

@@ -441,12 +441,11 @@ static void hpcq_vec_dot_q2_K_q8_K(int n, float * GGML_RESTRICT s, size_t bs, co
     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);
+        hpcq_seg_make_order_at(seg->packed, seg->enc, b0 + i, ord);
         y_perm.d = y[i].d;
         for (int g = 0; g < 16; ++g) {{
             int16_t bsum = 0;
@@ -474,12 +473,11 @@ static void hpcq_vec_dot_q3_K_q8_K(int n, float * GGML_RESTRICT s, size_t bs, co
     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);
+        hpcq_seg_make_order_at(seg->packed, seg->enc, b0 + i, ord);
         y_perm.d = y[i].d;
         for (int g = 0; g < 16; ++g) {{
             int16_t bsum = 0;
@@ -552,8 +550,7 @@ static inline const uint8_t * hpcq_cuda_get_dev_packed(
     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;
+        const size_t nbytes = hpcq_seg_packed_bytes(seg->n_blocks, seg->enc);
         void * d_ptr = nullptr;
         if (cudaMalloc(&d_ptr, nbytes) == cudaSuccess && d_ptr != nullptr) {{
             if (cudaMemcpy(d_ptr, seg->packed, nbytes, cudaMemcpyHostToDevice) == cudaSuccess) {{
@@ -580,8 +577,7 @@ static __global__ void dequantize_block_q2_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);
+        hpcq_seg_make_order_at(d_packed, enc, b0 + i, sh_ord);
     }}
     __syncthreads();
     dst_t * y_out = yy + i * QK_K;
@@ -603,8 +599,7 @@ static __global__ void dequantize_block_q3_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);
+        hpcq_seg_make_order_at(d_packed, enc, b0 + i, sh_ord);
     }}
     __syncthreads();
     dst_t * y_out = yy + i * QK_K;