Upload patch_llamacpp_hpcq_seg.py

patch_llamacpp_hpcq_seg.py

@@ -627,6 +627,18 @@ static void dequantize_row_q3_K_cuda(const void * vx, dst_t * y, const int64_t k
     dequantize_block_q3_K<<<nb, 64, 0, stream>>>(vx, y, d_packed, b0, enc);
 }}"""
 
+CUDA_FUSE_OLD = """    bool use_mul_mat_vec_q = ggml_is_quantized(src0->type) && !bad_padding_clear && src1->type == GGML_TYPE_F32 &&
+                             dst->type == GGML_TYPE_F32 && src1->ne[1] <= MMVQ_MAX_BATCH_SIZE;"""
+
+CUDA_FUSE_NEW = f"""    // {BEGIN}: do not fuse segmented Q2_K/Q3_K into mmvq
+    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) {{
+        return false;
+    }}
+    // {END}
+    bool use_mul_mat_vec_q = ggml_is_quantized(src0->type) && !bad_padding_clear && src1->type == GGML_TYPE_F32 &&
+                             dst->type == GGML_TYPE_F32 && src1->ne[1] <= MMVQ_MAX_BATCH_SIZE;"""
+
 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;
@@ -654,6 +666,28 @@ CUDA_MULMAT_NEW = f"""    // {BEGIN}: route segmented Q2_K/Q3_K to pooled-scratc
     }}
     ggml_cuda_mul_mat_cublas(ctx, src0, src1, dst);"""
 
+CUDA_MMID_SYNC_OLD = """    if (src1->type != GGML_TYPE_F32 || dst->type != GGML_TYPE_F32) {
+        return true;
+    }"""
+
+CUDA_MMID_SYNC_NEW = f"""    if (src1->type != GGML_TYPE_F32 || dst->type != GGML_TYPE_F32) {{
+        return true;
+    }}
+    // {BEGIN}: segmented MoE experts use the cublas fallback
+    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) {{
+        return true;
+    }}
+    // {END}"""
+
+CUDA_MMID_OLD = """    // [TAG_MUL_MAT_ID_CUDA_GRAPHS]
+    if (src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) {"""
+
+CUDA_MMID_NEW = f"""    // [TAG_MUL_MAT_ID_CUDA_GRAPHS]
+    const bool hpcq_seg_mmid = (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; // {BEGIN}/{END}
+    if (!hpcq_seg_mmid && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) {{"""
+
 
 def _read(p: Path) -> str:
     return p.read_text(encoding="utf-8")
@@ -752,7 +786,10 @@ def step_ggml_cuda(root: Path) -> None:
         _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_FUSE_OLD, CUDA_FUSE_NEW, "ggml-cuda.cu fuse_mmvq")
         s2 = _replace_block(s2, CUDA_MULMAT_OLD, CUDA_MULMAT_NEW, "ggml-cuda.cu mul_mat")
+        s2 = _replace_block(s2, CUDA_MMID_SYNC_OLD, CUDA_MMID_SYNC_NEW, "ggml-cuda.cu mmid_sync")
+        s2 = _replace_block(s2, CUDA_MMID_OLD, CUDA_MMID_NEW, "ggml-cuda.cu mmid")
         _write(cu, s2)