From 355303edab8745e722a2e0e1bd3f6a39399966f2 Mon Sep 17 00:00:00 2001 From: Todor Boinovski Date: Sat, 25 Jul 2026 15:47:29 -0700 Subject: [PATCH 001/190] hexagon: partial im2col support (#26007) * hexagon: add IM2COL op Add Hexagon IM2COL support targeting only patch-embedding convolutions. * hexagon: im2col refactor and cleanup * hex-im2col: instrument and update im2col. * hex-im2col: add local htp_vtcm_layout computation. --- ggml/src/ggml-hexagon/ggml-hexagon.cpp | 34 +++ ggml/src/ggml-hexagon/htp/CMakeLists.txt | 1 + ggml/src/ggml-hexagon/htp/htp-ctx.h | 1 + ggml/src/ggml-hexagon/htp/htp-ops.h | 1 + ggml/src/ggml-hexagon/htp/im2col-ops.c | 306 +++++++++++++++++++++++ ggml/src/ggml-hexagon/htp/main.c | 3 + 6 files changed, 346 insertions(+) create mode 100644 ggml/src/ggml-hexagon/htp/im2col-ops.c diff --git a/ggml/src/ggml-hexagon/ggml-hexagon.cpp b/ggml/src/ggml-hexagon/ggml-hexagon.cpp index 5db8fc84ba6e..bdb8af0820a3 100644 --- a/ggml/src/ggml-hexagon/ggml-hexagon.cpp +++ b/ggml/src/ggml-hexagon/ggml-hexagon.cpp @@ -3281,6 +3281,35 @@ static bool ggml_hexagon_supported_ssm_conv(const struct ggml_hexagon_session * GGML_UNUSED(sess); } +static bool ggml_hexagon_supported_im2col(const struct ggml_hexagon_session * sess, const struct ggml_tensor * op) { + const struct ggml_tensor * src1 = op->src[1]; + const struct ggml_tensor * dst = op; + + const bool is_2D = ((const int32_t *) op->op_params)[6] == 1; + if (!is_2D) { + return false; + } + + // For now support F32->F32 and F32->F16 only. + if (src1->type != GGML_TYPE_F32 || (dst->type != GGML_TYPE_F16 && dst->type != GGML_TYPE_F32)) { + return false; + } + + if (!ggml_is_contiguous(src1) || !ggml_is_contiguous(dst)) { + return false; + } + + // For now keep padded OPs on CPU. Will revisit once we expand coverage past patch-embed OPs. + const int32_t p0 = ((const int32_t *) op->op_params)[2]; + const int32_t p1 = ((const int32_t *) op->op_params)[3]; + if (p0 != 0 || p1 != 0) { + return false; + } + + GGML_UNUSED(sess); + return true; +} + static bool ggml_hexagon_supported_pad(const struct ggml_hexagon_session * sess, const struct ggml_tensor * op) { const struct ggml_tensor * src0 = op->src[0]; const struct ggml_tensor * dst = op; @@ -3430,6 +3459,7 @@ static htp_op_code op_remap_to_htp(const ggml_tensor * t) { case GGML_OP_SOLVE_TRI: return HTP_OP_SOLVE_TRI; case GGML_OP_TRI: return HTP_OP_TRI; case GGML_OP_PAD: return HTP_OP_PAD; + case GGML_OP_IM2COL: return HTP_OP_IM2COL; case GGML_OP_UNARY: switch (ggml_get_unary_op(t)) { @@ -4152,6 +4182,10 @@ static bool ggml_backend_hexagon_device_supports_op(ggml_backend_dev_t dev, cons supp = ggml_hexagon_supported_ssm_conv(sess, op); break; + case GGML_OP_IM2COL: + supp = ggml_hexagon_supported_im2col(sess, op); + break; + case GGML_OP_GATED_DELTA_NET: supp = ggml_hexagon_supported_gated_delta_net(sess, op); break; diff --git a/ggml/src/ggml-hexagon/htp/CMakeLists.txt b/ggml/src/ggml-hexagon/htp/CMakeLists.txt index 4fb526f0c001..b00aa2bc94c3 100644 --- a/ggml/src/ggml-hexagon/htp/CMakeLists.txt +++ b/ggml/src/ggml-hexagon/htp/CMakeLists.txt @@ -42,6 +42,7 @@ add_library(${HTP_LIB} SHARED solve-tri-ops.c pad-ops.c argsort-ops.c + im2col-ops.c ) target_compile_definitions(${HTP_LIB} PRIVATE diff --git a/ggml/src/ggml-hexagon/htp/htp-ctx.h b/ggml/src/ggml-hexagon/htp/htp-ctx.h index 97b8c7f29792..e0f9a0c40d19 100644 --- a/ggml/src/ggml-hexagon/htp/htp-ctx.h +++ b/ggml/src/ggml-hexagon/htp/htp-ctx.h @@ -140,5 +140,6 @@ int op_diag(struct htp_ops_context * octx); int op_solve_tri(struct htp_ops_context * octx); int op_gated_delta_net(struct htp_ops_context * octx); int op_pad(struct htp_ops_context * octx); +int op_im2col(struct htp_ops_context * octx); #endif /* HTP_CTX_H */ diff --git a/ggml/src/ggml-hexagon/htp/htp-ops.h b/ggml/src/ggml-hexagon/htp/htp-ops.h index cad9a4f54007..a138f062aa68 100644 --- a/ggml/src/ggml-hexagon/htp/htp-ops.h +++ b/ggml/src/ggml-hexagon/htp/htp-ops.h @@ -98,6 +98,7 @@ enum htp_op_code { HTP_OP_NORM, HTP_OP_CONCAT, HTP_OP_CLAMP, + HTP_OP_IM2COL, HTP_OP_INVALID }; diff --git a/ggml/src/ggml-hexagon/htp/im2col-ops.c b/ggml/src/ggml-hexagon/htp/im2col-ops.c new file mode 100644 index 000000000000..35fc103df8fe --- /dev/null +++ b/ggml/src/ggml-hexagon/htp/im2col-ops.c @@ -0,0 +1,306 @@ +#pragma clang diagnostic ignored "-Wunused-variable" +#pragma clang diagnostic ignored "-Wunused-function" +#pragma clang diagnostic ignored "-Wunused-but-set-variable" + +#include +#include +#include +#include +#include + +#define GGML_COMMON_DECL_C +#include "ggml-common.h" +#include "htp-ctx.h" +#include "htp-ops.h" +#include "hvx-utils.h" +#include "hex-dma.h" +#include "hex-profile.h" +#include "htp-vtcm.h" + +struct htp_im2col_context { + struct htp_ops_context * octx; + uint32_t npatches_per_thread; // patches = N*OH*OW (pure-DDR kernel) + + uint32_t pe_rows_per_thread; // N*OH rows per worker + uint32_t pe_src_row_bytes; // one output row's source: IC*KH*IW*4, rounded 256 + uint32_t pe_dst_row_bytes; // one output row's dst: OW*patch_stride*2, rounded 256 + + // Patch-embed DMA path VTCM ping-pong. + uint8_t * pe_vtcm_src; // base of the 2x src buffers region + uint8_t * pe_vtcm_dst; // base of the 2x dst buffers region + uint32_t pe_src_size_per_thread; // 2 * pe_src_row_bytes + uint32_t pe_dst_size_per_thread; // 2 * pe_dst_row_bytes +}; + +// Per-op VTCM layout for the patch-embed DMA path +struct htp_im2col_vtcm_layout { + size_t off_src; + size_t off_dst; + size_t src_bytes_per_thread; + size_t dst_bytes_per_thread; + size_t total_bytes; +}; + +static inline void htp_im2col_vtcm_layout_build(struct htp_im2col_vtcm_layout * L, + size_t src_row_bytes, + size_t dst_row_bytes, + uint32_t n_threads) { + L->src_bytes_per_thread = 2 * src_row_bytes; + L->dst_bytes_per_thread = 2 * dst_row_bytes; + + L->off_src = 0; + L->off_dst = L->off_src + L->src_bytes_per_thread * n_threads; + L->total_bytes = L->off_dst + L->dst_bytes_per_thread * n_threads; +} + +#define IM2COL_PATCHEMBED_BODY(FNAME, DST_CTYPE, COPY_FN, SPLAT_FN, DST_ELEM, TAG) \ + static void FNAME(unsigned int nth, unsigned int ith, void * data) { \ + struct htp_im2col_context * ictx = (struct htp_im2col_context *) data; \ + struct htp_ops_context * octx = ictx->octx; \ + struct htp_thread_trace * restrict tr = &octx->ctx->trace[ith]; \ + const struct htp_tensor * restrict src1 = octx->src[1]; \ + const struct htp_tensor * restrict dst = octx->dst; \ + const int32_t s0 = octx->op_params[0]; \ + const int32_t s1 = octx->op_params[1]; \ + const int32_t p0 = octx->op_params[2]; \ + const int32_t p1 = octx->op_params[3]; \ + const int32_t d0 = octx->op_params[4]; \ + const int32_t d1 = octx->op_params[5]; \ + const uint32_t N = src1->ne[3]; \ + const uint32_t IC = src1->ne[2]; \ + const uint32_t IH = src1->ne[1]; \ + const uint32_t IW = src1->ne[0]; \ + const uint32_t KH = octx->src[0]->ne[1]; \ + const uint32_t KW = octx->src[0]->ne[0]; \ + const uint32_t OH = dst->ne[2]; \ + const uint32_t OW = dst->ne[1]; \ + const uint32_t patch_stride = IC * KH * KW; \ + const float * restrict src_data = (const float *) src1->data; \ + DST_CTYPE * restrict dst_data = (DST_CTYPE *) dst->data; \ + const uint32_t npatches = N * OH * OW; \ + const uint32_t patch_start = ictx->npatches_per_thread * ith; \ + const uint32_t patch_end = MIN(patch_start + ictx->npatches_per_thread, npatches); \ + if (patch_start >= patch_end) { \ + return; \ + } \ + htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, patch_start); \ + for (uint32_t p = patch_start; p < patch_end; p++) { \ + const uint32_t iow = p % OW; \ + const uint32_t ioh = (p / OW) % OH; \ + const uint32_t in = p / (OW * OH); \ + DST_CTYPE * restrict dst_patch = dst_data + (uint64_t) p * patch_stride; \ + for (uint32_t iic = 0; iic < IC; iic++) { \ + const float * restrict src_plane = src_data + ((uint64_t) in * IC + iic) * IH * IW; \ + for (uint32_t ikh = 0; ikh < KH; ikh++) { \ + const int32_t iih = (int32_t) ioh * s1 + (int32_t) ikh * d1 - p1; \ + DST_CTYPE * restrict out_run = dst_patch + iic * (KH * KW) + ikh * KW; \ + if (iih < 0 || iih >= (int32_t) IH) { \ + SPLAT_FN(out_run, 0.0f, KW); \ + continue; \ + } \ + const int32_t iiw0 = (int32_t) iow * s0 - p0; \ + const float * restrict src_run = src_plane + (uint64_t) iih * IW + iiw0; \ + if (d0 == 1) { \ + /* contiguous source run: [lo,hi) is in-bounds, tails are zero pad */ \ + const int32_t lo = iiw0 < 0 ? -iiw0 : 0; \ + int32_t hi = (int32_t) IW - iiw0; \ + if (hi > (int32_t) KW) { \ + hi = (int32_t) KW; \ + } \ + if (hi <= lo) { \ + SPLAT_FN(out_run, 0.0f, KW); \ + } else { \ + if (lo > 0) { \ + SPLAT_FN(out_run, 0.0f, (uint32_t) lo); \ + } \ + COPY_FN((uint8_t *) (out_run + lo), (const uint8_t *) (src_run + lo), \ + (uint32_t) (hi - lo)); \ + if (hi < (int32_t) KW) { \ + SPLAT_FN(out_run + hi, 0.0f, (KW - (uint32_t) hi)); \ + } \ + } \ + continue; \ + } \ + for (uint32_t ikw = 0; ikw < KW; ikw++) { \ + const int32_t iiw = (int32_t) iow * s0 + (int32_t) ikw * d0 - p0; \ + out_run[ikw] = (iiw < 0 || iiw >= (int32_t) IW) ? \ + (DST_CTYPE) 0.0f : \ + (DST_CTYPE) src_plane[(uint64_t) iih * IW + iiw]; \ + } \ + } \ + } \ + } \ + htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, patch_start); \ + } + +IM2COL_PATCHEMBED_BODY(im2col_patchembed_thread, __fp16, hvx_copy_f16_f32_uu, hvx_splat_f16_u, sizeof(__fp16), "f32-f16") +IM2COL_PATCHEMBED_BODY(im2col_patchembed_f32_thread, float, hvx_copy_f32_uu, hvx_splat_f32_u, sizeof(float), "f32-f32") + +#define IM2COL_PATCHEMBED_DMA_BODY(FNAME, DST_CTYPE, COPY_FN, SPLAT_FN, DST_ELEM, TAG) \ + static void FNAME(unsigned int nth, unsigned int ith, void * data) { \ + struct htp_im2col_context * ictx = (struct htp_im2col_context *) data; \ + struct htp_ops_context * octx = ictx->octx; \ + struct htp_thread_trace * restrict tr = &octx->ctx->trace[ith]; \ + const struct htp_tensor * restrict src1 = octx->src[1]; \ + const struct htp_tensor * restrict dst = octx->dst; \ + const uint32_t N = src1->ne[3], IC = src1->ne[2], IH = src1->ne[1], IW = src1->ne[0]; \ + const uint32_t KH = octx->src[0]->ne[1], KW = octx->src[0]->ne[0]; \ + const uint32_t OH = dst->ne[2], OW = dst->ne[1]; \ + const uint32_t patch_stride = IC * KH * KW; \ + const float * restrict src_data = (const float *) src1->data; \ + DST_CTYPE * restrict dst_data = (DST_CTYPE *) dst->data; \ + dma_queue * dmaq = octx->ctx->dma[ith]; \ + uint8_t * src_base = ictx->pe_vtcm_src + ith * ictx->pe_src_size_per_thread; \ + uint8_t * dst_base = ictx->pe_vtcm_dst + ith * ictx->pe_dst_size_per_thread; \ + float * srcb = (float *) src_base; \ + DST_CTYPE * dstb = (DST_CTYPE *) dst_base; \ + const uint32_t nrows = N * OH; \ + const uint32_t per_thread = ictx->pe_rows_per_thread; \ + const uint32_t row_start = per_thread * ith; \ + const uint32_t row_end = MIN(row_start + per_thread, nrows); \ + if (row_start >= row_end) \ + return; \ + for (uint32_t r = row_start; r < row_end; r++) { \ + const uint32_t in = r / OH; \ + const uint32_t ioh = r % OH; \ + for (uint32_t ikh = 0; ikh < KH; ikh++) { \ + int32_t iih = (int32_t) ioh * (int32_t) KH + (int32_t) ikh; \ + int ok = (iih >= 0 && iih < (int32_t) IH); \ + for (uint32_t iic = 0; iic < IC; iic++) { \ + float * vdst = srcb + ((uint64_t) (iic * KH + ikh)) * IW; \ + const float * _vsrc = \ + ok ? (src_data + ((uint64_t) (in * IC + iic) * IH + iih) * IW) : (const float *) vdst; \ + dma_queue_push_ddr_to_vtcm( \ + dmaq, dma_make_ptr((uint8_t *) vdst, ok ? (const uint8_t *) _vsrc : (const uint8_t *) vdst), \ + IW * sizeof(float), IW * sizeof(float), ok ? 1 : 0); \ + } \ + } \ + for (uint32_t i = 0; i < IC * KH; i++) \ + dma_queue_pop(dmaq); \ + htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, r); \ + for (uint32_t iow = 0; iow < OW; iow++) { \ + DST_CTYPE * dst_patch = dstb + (uint64_t) iow * patch_stride; \ + for (uint32_t ikh = 0; ikh < KH; ikh++) { \ + int32_t iih = (int32_t) ioh * (int32_t) KH + (int32_t) ikh; \ + for (uint32_t iic = 0; iic < IC; iic++) { \ + DST_CTYPE * out_run = dst_patch + iic * (KH * KW) + ikh * KW; \ + if (iih < 0 || iih >= (int32_t) IH) { \ + SPLAT_FN(out_run, 0.0f, KW); \ + continue; \ + } \ + const float * src_run = srcb + ((uint64_t) (iic * KH + ikh)) * IW + (uint64_t) iow * KW; \ + COPY_FN((uint8_t *) out_run, (const uint8_t *) src_run, KW); \ + } \ + } \ + } \ + htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, r); \ + DST_CTYPE * ddr_row = dst_data + ((uint64_t) (in * OH + ioh) * OW) * patch_stride; \ + dma_queue_push_vtcm_to_ddr(dmaq, dma_make_ptr((uint8_t *) ddr_row, (uint8_t *) dstb), \ + OW * patch_stride * (DST_ELEM), OW * patch_stride * (DST_ELEM), 1); \ + dma_queue_flush(dmaq); \ + } \ + } + +IM2COL_PATCHEMBED_DMA_BODY(im2col_patchembed_dma_thread, __fp16, hvx_copy_f16_f32_uu, hvx_splat_f16_u, sizeof(__fp16), "pe-dma-f16") +IM2COL_PATCHEMBED_DMA_BODY(im2col_patchembed_dma_f32_thread, float, hvx_copy_f32_uu, hvx_splat_f32_u, sizeof(float), "pe-dma-f32") + +static bool im2col_use_patchembed_dma(const struct htp_ops_context * octx) { + const int32_t s0 = octx->op_params[0], s1 = octx->op_params[1]; + const int32_t p0 = octx->op_params[2], p1 = octx->op_params[3]; + const int32_t d0 = octx->op_params[4], d1 = octx->op_params[5]; + const int is_2D = octx->op_params[6] == 1; + if (!is_2D) { + return false; + } + if (octx->dst->type != HTP_TYPE_F16 && octx->dst->type != HTP_TYPE_F32) { + return false; + } + const uint32_t KH = octx->src[0]->ne[1], KW = octx->src[0]->ne[0]; + if (s0 != (int32_t) KW || s1 != (int32_t) KH) { + return false; // non-overlapping + } + if (p0 != 0 || p1 != 0) { + return false; // no padding + } + if (d0 != 1 || d1 != 1) { + return false; // no dilation + } + return true; +} + +// Sizes the per-thread 2x(src,dst) VTCM ping-pong for the patch-embed DMA path. +// Returns false if it doesn't fit the VTCM budget (caller falls back). +static bool im2col_patchembed_dma_fits(struct htp_ops_context * octx, + struct htp_im2col_context * ictx, + uint32_t n_threads) { + const uint32_t IC = octx->src[1]->ne[2], IW = octx->src[1]->ne[0]; + const uint32_t KH = octx->src[0]->ne[1], KW = octx->src[0]->ne[0]; + const uint32_t OW = octx->dst->ne[1]; + const uint32_t patch_stride = IC * KH * KW; + + ictx->pe_src_row_bytes = hex_round_up(IC * KH * IW * sizeof(float), 256); + const uint32_t dst_elem = (octx->dst->type == HTP_TYPE_F16) ? sizeof(__fp16) : sizeof(float); + ictx->pe_dst_row_bytes = hex_round_up(OW * patch_stride * dst_elem, 256); + + // 2 src + 2 dst buffers per thread (ping-pong), src region first then dst. + struct htp_im2col_vtcm_layout L; + htp_im2col_vtcm_layout_build(&L, ictx->pe_src_row_bytes, ictx->pe_dst_row_bytes, n_threads); + if (L.total_bytes > octx->ctx->vtcm_size) { + return false; + } + + uint8_t * const base = octx->ctx->vtcm_base; + ictx->pe_vtcm_src = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src); + ictx->pe_vtcm_dst = VTCM_LAYOUT_PTR(uint8_t, base, L.off_dst); + ictx->pe_src_size_per_thread = (uint32_t) L.src_bytes_per_thread; + ictx->pe_dst_size_per_thread = (uint32_t) L.dst_bytes_per_thread; + return true; +} + +int op_im2col(struct htp_ops_context * octx) { + const struct htp_tensor * src1 = octx->src[1]; + const struct htp_tensor * dst = octx->dst; + + if (src1->type != HTP_TYPE_F32 || (dst->type != HTP_TYPE_F16 && dst->type != HTP_TYPE_F32)) { + FARF(ERROR, "im2col: only (F32 image -> F16/F32 columns) supported"); + return HTP_STATUS_NO_SUPPORT; + } + + const uint32_t N = src1->ne[3]; + const uint32_t OH = dst->ne[2]; + const uint32_t OW = dst->ne[1]; + const uint32_t npatches = N * OH * OW; + const uint32_t n_threads = MIN(octx->n_threads, npatches); + + if ((octx->flags & HTP_OPFLAGS_SKIP_COMPUTE) || n_threads == 0) { + return HTP_STATUS_OK; + } + + struct htp_im2col_context ictx = { 0 }; + ictx.octx = octx; + ictx.npatches_per_thread = (npatches + n_threads - 1) / n_threads; + + // Clean non-overlapping patch-embed -> DMA kernel (if it fits VTCM); + // everything else (padding/dilation/stride edges) -> pure-DDR kernel. + if (im2col_use_patchembed_dma(octx)) { + const uint32_t nrows = N * OH; + const uint32_t pth = MIN(octx->n_threads, nrows); + if (pth > 0 && im2col_patchembed_dma_fits(octx, &ictx, pth)) { + ictx.pe_rows_per_thread = (nrows + pth - 1) / pth; + if (dst->type == HTP_TYPE_F16) { + work_queue_run(octx->ctx->work_queue, im2col_patchembed_dma_thread, &ictx, pth); + } else { + work_queue_run(octx->ctx->work_queue, im2col_patchembed_dma_f32_thread, &ictx, pth); + } + return HTP_STATUS_OK; + } + // else: doesn't fit -> fall through to the pure-DDR kernel below. + } + + if (dst->type == HTP_TYPE_F16) { + work_queue_run(octx->ctx->work_queue, im2col_patchembed_thread, &ictx, n_threads); + } else { + work_queue_run(octx->ctx->work_queue, im2col_patchembed_f32_thread, &ictx, n_threads); + } + return HTP_STATUS_OK; +} diff --git a/ggml/src/ggml-hexagon/htp/main.c b/ggml/src/ggml-hexagon/htp/main.c index 7d65e46436cc..880e20c99597 100644 --- a/ggml/src/ggml-hexagon/htp/main.c +++ b/ggml/src/ggml-hexagon/htp/main.c @@ -781,6 +781,9 @@ static int execute_op(struct htp_ops_context * octx) { case HTP_OP_PAD: return op_pad(octx); + case HTP_OP_IM2COL: + return op_im2col(octx); + case HTP_OP_CONCAT: return op_concat(octx); From 20455a4ad336e958cfe8f82efce2c46cd44c4fa3 Mon Sep 17 00:00:00 2001 From: Xuan-Son Nguyen Date: Sun, 26 Jul 2026 01:08:49 +0200 Subject: [PATCH 002/190] server: support MCP stdio (#26062) * move server_pipe to common * init impl * vendor: update subprocess.h * add server_mcp_stdio * stderr drain * server_mcp_transport * server_mcp_stdio is now framing-only, no json * internal/mcp-stdio: integration + tests + fixes (#26075) * server-mcp: harden transport and wire up the tool integration Builds on the transport/manager architecture (server_mcp_transport + server_pipe) with the hardening and integration the draft did not yet have. Hardening: * Reader and stderr pumps are polled (running-aware) instead of blocking on a read that only ends at EOF. subprocess_terminate() SIGKILLs only the direct child, so a grandchild the MCP server spawned that inherited the pipe would otherwise keep the write end open and hang teardown (both warmup shutdown at startup and process shutdown). The writer is likewise non-blocking + polled. * Windows: resolve the command through PATHEXT so "npx" (npm ships npx.cmd, never npx.exe) spawns, matching POSIX's PATH search; and enumerate the parent environment as UTF-8 (GetEnvironmentStringsW) instead of the active code page. * server_pipe gains an opt-in max_size (default unbounded, so the router's streaming use is unchanged); the MCP reply queue uses it so a server that streams unsolicited notifications between requests cannot grow it without bound. Integration: * --mcp-servers-config / --mcp-servers-json flags; enabling MCP restricts default CORS to localhost, same as --tools. * MCP tools are exposed through /tools (and chat-completions) as _, skipping names that collide with a built-in or another MCP tool. * Manager lifecycle wired into llama_server(): warmup at start, shutdown() from the signal handler before the HTTP server drains, blocking teardown in clean_up(). * SIGPIPE ignored so a child dying mid-write yields EPIPE rather than killing us. Assisted-By: Claude Opus 4.8 * server-mcp: add MCP test suite with grandchild deadlock regression test 21 tests over the /tools endpoint: tool discovery/invocation, timeouts, crash recovery and respawn cooldown, warmup partial failure, malformed and batched notification+response output, tool-definition shape, and prompt shutdown during a slow call. The last test spawns an MCP server that leaves a grandchild inheriting its stdout/stderr and asserts the server both starts and stops promptly. Verified it fails (5s SIGKILL fallback on a deadlocked reader-join) when the pump is made to ignore the running flag, and passes with the polled reader. Assisted-By: Claude Opus 4.8 * clean up * clean up 2 * even stricter life cycle * nits * nits 2 --------- Co-authored-by: Xuan Son Nguyen * fix some edge cases * fix last_error data race * fix response schema + docs * server: fix MCP zombie leak and timeout-induced transport teardown join_pumps() never reaped the child, leaking one zombie per spawn: call subprocess_join() before subprocess_destroy(). A per-call timeout permanently closed from_server and got a healthy transport evicted: add close_on_stop to server_pipe::read() and pass false from send_rpc(), where should_stop is a per-request deadline and a late reply is already skipped on id mismatch. Also drop the unreachable disconnect cancellation in server_mcp_tool::invoke(): support_stream is false, st is always null. (cherry picked from commit e6de1ec043174fd0570b1e60d47f06c7c19d620d) Assisted-by: Claude Opus 4.8 * server: make MCP test fixtures JSON-RPC 2.0 compliant Add the missing notification guard to mcp_malformed_server.py and mcp_burst_server.py (the latter treated id 0 as a notification and replied to unknown ones; its notification table is now unused). Return -32602 instead of -32601 for unknown tools: tools/call is a valid method, the tool name is the invalid parameter. Also fix the test module docstring: tools are named _. (cherry picked from commit 74a08e8c311dabf3b49d06cc6d754b0097ae7a38) Assisted-by: Claude Opus 4.8 --------- Co-authored-by: Piotr Wilkin (ilintar) Co-authored-by: Pascal --- common/arg.cpp | 21 +- common/common.h | 4 + tools/server/CMakeLists.txt | 2 + tools/server/README-dev.md | 4 +- tools/server/server-common.h | 72 +- tools/server/server-mcp.cpp | 836 ++++++++++++++++++ tools/server/server-mcp.h | 176 ++++ tools/server/server-models.cpp | 49 +- tools/server/server-tools.cpp | 73 +- tools/server/server-tools.h | 6 +- tools/server/server.cpp | 35 +- .../server/tests/fixtures/mcp_burst_server.py | 118 +++ .../server/tests/fixtures/mcp_crash_server.py | 114 +++ .../server/tests/fixtures/mcp_echo_server.py | 164 ++++ .../tests/fixtures/mcp_grandchild_server.py | 100 +++ .../tests/fixtures/mcp_malformed_server.py | 113 +++ .../server/tests/fixtures/mcp_slow_server.py | 132 +++ tools/server/tests/unit/test_mcp_servers.py | 718 +++++++++++++++ tools/server/tests/utils.py | 6 + 19 files changed, 2680 insertions(+), 63 deletions(-) create mode 100644 tools/server/server-mcp.cpp create mode 100644 tools/server/server-mcp.h create mode 100644 tools/server/tests/fixtures/mcp_burst_server.py create mode 100644 tools/server/tests/fixtures/mcp_crash_server.py create mode 100755 tools/server/tests/fixtures/mcp_echo_server.py create mode 100644 tools/server/tests/fixtures/mcp_grandchild_server.py create mode 100644 tools/server/tests/fixtures/mcp_malformed_server.py create mode 100644 tools/server/tests/fixtures/mcp_slow_server.py create mode 100644 tools/server/tests/unit/test_mcp_servers.py diff --git a/common/arg.cpp b/common/arg.cpp index a287b907d490..9753441313a7 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -850,8 +850,9 @@ static bool common_params_parse_ex(int argc, char ** argv, common_params_context params.kv_overrides.back().key[0] = 0; } - if (!params.server_tools.empty() && !params.cors_origins_explicit) { - LOG_WRN("server tools are enabled, using localhost as default CORS origin (change via --cors-origins)\n"); + const bool mcp_enabled = !params.mcp_servers_config.empty() || !params.mcp_servers_json.empty(); + if ((!params.server_tools.empty() || mcp_enabled) && !params.cors_origins_explicit) { + LOG_WRN("server tools or MCP servers are enabled, using localhost as default CORS origin (change via --cors-origins)\n"); params.cors_origins = "localhost"; } @@ -3261,6 +3262,22 @@ common_params_context common_params_parser_init(common_params & params, llama_ex params.server_tools = parse_csv_row(value); } ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_TOOLS")); + add_opt(common_arg( + {"--mcp-servers-config"}, "PATH", + "experimental: path to JSON file with MCP server definitions (Cursor-compatible format) - do not enable in untrusted environments (default: none)\n" + "note: for security reasons, this will limit --cors-origins to localhost by default", + [](common_params & params, const std::string & value) { + params.mcp_servers_config = value; + } + ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_MCP_SERVERS_CONFIG")); + add_opt(common_arg( + {"--mcp-servers-json"}, "JSON", + "experimental: inline JSON with MCP server definitions (Cursor-compatible format) - do not enable in untrusted environments (default: none)\n" + "note: for security reasons, this will limit --cors-origins to localhost by default", + [](common_params & params, const std::string & value) { + params.mcp_servers_json = value; + } + ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_MCP_SERVERS_JSON")); add_opt(common_arg( {"-ag", "--agent"}, {"-no-ag", "--no-agent"}, diff --git a/common/common.h b/common/common.h index b5687c10836a..2792521836ae 100644 --- a/common/common.h +++ b/common/common.h @@ -668,6 +668,10 @@ struct common_params { // enable built-in tools std::vector server_tools; + // MCP server configs (Cursor-compatible JSON) + std::string mcp_servers_config; // path to JSON file with MCP server definitions + std::string mcp_servers_json; // inline JSON with MCP server definitions + // router server configs std::string models_dir = ""; // directory containing models for the router server std::string models_preset = ""; // directory containing model presets for the router server diff --git a/tools/server/CMakeLists.txt b/tools/server/CMakeLists.txt index b5c40884fd6e..280bd9e19dca 100644 --- a/tools/server/CMakeLists.txt +++ b/tools/server/CMakeLists.txt @@ -19,6 +19,8 @@ add_library(${TARGET} STATIC server-stream.h server-tools.cpp server-tools.h + server-mcp.cpp + server-mcp.h server-schema.cpp server-schema.h ) diff --git a/tools/server/README-dev.md b/tools/server/README-dev.md index e81336e5e26b..b4ec9f17d3c6 100644 --- a/tools/server/README-dev.md +++ b/tools/server/README-dev.md @@ -189,7 +189,7 @@ This endpoint is intended to be used internally by the Web UI and subject to cha Get a list of tools, each tool has these fields: - `tool` (string): the ID name of the tool, to be used in POST call. Example: `read_file` - `display_name` (string): the name to be displayed on UI. Example: `Read file` -- `type` (string): always be `"builtin"` for now +- `type` (string): `"builtin"` for a built-in tool, or `"mcp"` for a tool exposed by an MCP server - `permissions` (object): a mapping string --> boolean that indicates the permission required by this tool. This is useful for the UI to ask the user before calling the tool. For now, the only permission supported is `"write"` - `definition` (object): the OAI-compat definition of this tool @@ -199,7 +199,7 @@ Invoke a tool call, request body is a JSON object with: - `tool` (string): the name of the tool - `params` (object): a mapping from argument name (string) to argument value -Returns JSON object. There are two response formats: +Returns JSON object. There are two response formats (MCP tools use the same two formats: their result content is concatenated into `plain_text_response`, and RPC or tool errors are surfaced as the `error` string): Format 1: Plain text. The text will be placed into a field called `plain_text_response`, example: diff --git a/tools/server/server-common.h b/tools/server/server-common.h index 583736638032..6ef797ebb473 100644 --- a/tools/server/server-common.h +++ b/tools/server/server-common.h @@ -9,9 +9,15 @@ #define JSON_ASSERT GGML_ASSERT #include +#include +#include +#include +#include +#include +#include +#include #include #include -#include using json = nlohmann::ordered_json; @@ -376,3 +382,67 @@ server_tokens format_prompt_rerank( mtmd_context * mctx, const std::string & query, const std::string & doc); + +// simple implementation of a pipe +// used for streaming data between threads +template +struct server_pipe { + std::mutex mutex; + std::condition_variable cv; + std::queue queue; + std::atomic writer_closed{false}; + std::atomic reader_closed{false}; + + // 0 = unbounded (default) + // > 0, write() drops the oldest item once the queue is full + size_t max_size = 0; + + void close_write() { + writer_closed.store(true, std::memory_order_relaxed); + cv.notify_all(); + } + + void close_read() { + reader_closed.store(true, std::memory_order_relaxed); + cv.notify_all(); + } + + // close_on_stop = true: should_stop means the reader is gone for good, so the writer is told the pipe is broken. + // close_on_stop = false: should_stop is a per-read deadline and further reads still come, so the pipe stays usable. + bool read(T & output, const std::function & should_stop, bool close_on_stop = true) { + std::unique_lock lk(mutex); + constexpr auto poll_interval = std::chrono::milliseconds(500); + while (true) { + if (!queue.empty()) { + output = std::move(queue.front()); + queue.pop(); + return true; + } + if (writer_closed.load()) { + return false; // clean EOF + } + if (should_stop && should_stop()) { // a null should_stop means "never stop" + if (close_on_stop) { + close_read(); // signal broken pipe to writer + } + return false; // cancelled / deadline reached + } + cv.wait_for(lk, poll_interval); + } + } + + bool write(T && data) { + std::lock_guard lk(mutex); + if (reader_closed.load()) { + return false; // broken pipe + } + if (max_size > 0) { + while (queue.size() >= max_size) { + queue.pop(); // drop oldest to stay bounded + } + } + queue.push(std::move(data)); + cv.notify_one(); + return true; + } +}; diff --git a/tools/server/server-mcp.cpp b/tools/server/server-mcp.cpp new file mode 100644 index 000000000000..aa87afeb065a --- /dev/null +++ b/tools/server/server-mcp.cpp @@ -0,0 +1,836 @@ +#include "server-mcp.h" + +#include + +#include +#include +#include +#include +#include +#include +#include + +#if defined(_WIN32) +# include +# include +#else +# include +# include +# include +# include +extern char ** environ; +#endif + +// read NDJSON lines from a child pipe, calling on_line per line until `running` clears, EOF/error, or on_line returns false. +// polled, not blocking: a grandchild can inherit the pipe's write end and hold it open (terminate() kills only the direct child), so a blocking read would hang teardown on an EOF that never comes. +static void mcp_pump_ndjson(FILE * f, std::atomic & running, + const std::function & on_line) { + if (!f) { + return; + } + const int poll_ms = 50; + const size_t max_line = 8 * 1024 * 1024; // drop any single NDJSON line larger than this, so a child that never emits '\n' can't grow buf without bound +#if defined(_WIN32) + HANDLE h = (HANDLE) _get_osfhandle(_fileno(f)); +#else + int fd = fileno(f); + int fl = fcntl(fd, F_GETFL, 0); + if (fl >= 0) { + fcntl(fd, F_SETFL, fl | O_NONBLOCK); + } +#endif + std::string buf; + bool skipping = false; // discarding an over-long line until its terminating newline + char chunk[4096]; + while (running.load()) { + size_t n = 0; +#if defined(_WIN32) + DWORD avail = 0; + if (!PeekNamedPipe(h, NULL, 0, NULL, &avail, NULL)) { + break; // pipe broken / child gone + } + if (avail == 0) { + std::this_thread::sleep_for(std::chrono::milliseconds(poll_ms)); + continue; + } + DWORD to_read = avail < (DWORD) sizeof(chunk) ? avail : (DWORD) sizeof(chunk); + DWORD got = 0; + if (!ReadFile(h, chunk, to_read, &got, NULL) || got == 0) { + break; + } + n = (size_t) got; +#else + struct pollfd pfd; + pfd.fd = fd; + pfd.events = POLLIN; + pfd.revents = 0; + int pr = poll(&pfd, 1, poll_ms); + if (pr < 0) { + if (errno == EINTR) { + continue; + } + break; + } + if (pr == 0) { + continue; // timeout -> re-check running + } + if (pfd.revents & (POLLERR | POLLNVAL)) { + break; + } + ssize_t r = read(fd, chunk, sizeof(chunk)); + if (r < 0) { + if (errno == EINTR || errno == EAGAIN || errno == EWOULDBLOCK) { + continue; + } + break; + } + if (r == 0) { + break; // EOF: child (and any pipe writers) closed the stream + } + n = (size_t) r; +#endif + buf.append(chunk, n); + + // resync after an over-long, unterminated line: discard bytes until the next newline + if (skipping) { + size_t nl = buf.find('\n'); + if (nl == std::string::npos) { + if (buf.size() > max_line) { + buf.clear(); // stay bounded while waiting for a terminator + } + continue; + } + buf.erase(0, nl + 1); + skipping = false; + } + + size_t pos; + while ((pos = buf.find('\n')) != std::string::npos) { + std::string line = buf.substr(0, pos); + buf.erase(0, pos + 1); + if (!line.empty() && line.back() == '\r') { + line.pop_back(); + } + if (line.empty()) { + continue; + } + if (!on_line(std::move(line))) { + return; + } + } + + // a partial line already larger than the cap and still no newline: drop it to avoid unbounded growth + if (buf.size() > max_line) { + SRV_WRN("MCP: dropping oversized line (> %zu bytes) from child pipe\n", max_line); + buf.clear(); + skipping = true; + } + } +} + +// +// server_mcp_server_config +// + +std::vector server_mcp_server_config::parse_from_json(const std::string & json_str) { + return parse_cursor_format(json::parse(json_str)); +} + +std::vector server_mcp_server_config::parse_cursor_format(const json & j) { + std::vector result; + + if (!j.contains("mcpServers") || !j.at("mcpServers").is_object()) { + return result; + } + + for (const auto & [name, cfg] : j.at("mcpServers").items()) { + server_mcp_server_config sc; + sc.name = name; + sc.command = cfg.value("command", std::string()); + sc.cwd = cfg.value("cwd", std::string()); + sc.timeout_ms = cfg.value("timeout_ms", sc.timeout_ms); + + if (cfg.contains("args") && cfg.at("args").is_array()) { + for (const auto & a : cfg.at("args")) { + sc.args.push_back(a.get()); + } + } + if (cfg.contains("env") && cfg.at("env").is_object()) { + for (const auto & [k, v] : cfg.at("env").items()) { + sc.env[k] = v.get(); + } + } + + if (sc.command.empty()) { + SRV_WRN("MCP server '%s' has no command, skipping\n", name.c_str()); + continue; + } + result.push_back(std::move(sc)); + } + + return result; +} + + +// +// server_mcp_transport +// + +static constexpr const char * MCP_PROTOCOL_VERSION = "2024-11-05"; + +static std::string rpc_error_message(const json & resp) { + if (resp.contains("error")) { + const json & e = resp.at("error"); + if (e.is_object()) { + return e.value("message", "unknown error"); + } + if (e.is_string()) { + return e.get(); + } + } + return "unknown error"; +} + +// normalize an MCP tools/call result to the /tools contract (see README-dev.md): +// concat text parts of result.content[], and surface an isError result +static json mcp_result_to_response(const json & result) { + std::string text; + if (result.contains("content") && result.at("content").is_array()) { + for (const auto & part : result.at("content")) { + if (part.is_object() && part.value("type", "") == "text") { + if (!text.empty()) { + text += "\n"; + } + text += part.value("text", ""); + } + } + } + if (result.is_object() && result.value("isError", false)) { + return {{"error", text.empty() ? "MCP tool returned an error" : text}}; + } + return {{"plain_text_response", text}}; +} + +json server_mcp_transport::send_rpc(const json & request, const std::function & should_stop) { + if (!to_server.write(request.dump())) { + return {{"error", {{"code", -32603}, {"message", "transport closed"}}}}; + } + + const bool has_id = request.contains("id"); + const auto deadline = std::chrono::steady_clock::now() + std::chrono::milliseconds(timeout_ms); + auto stop = [&]() { + return (should_stop && should_stop()) || std::chrono::steady_clock::now() >= deadline; + }; + + std::string frame; + while (from_server.read(frame, stop, false)) { + json reply; + try { + reply = json::parse(frame); + } catch (...) { + if (std::chrono::steady_clock::now() >= deadline) { + break; + } + continue; // skip malformed frame + } + // no id: a notification. mismatched id: a stale reply from a timed-out request (ids are monotonic, never a future one) + if (!has_id || (reply.contains("id") && reply.at("id") == request.at("id"))) { + return reply; + } + if (std::chrono::steady_clock::now() >= deadline) { + break; // a flood of notifications must not outrun the deadline + } + } + + if (should_stop && should_stop()) { + return {{"error", {{"code", -32603}, {"message", "cancelled"}}}}; + } + if (std::chrono::steady_clock::now() >= deadline) { + return {{"error", {{"code", -32603}, {"message", "request timed out"}}}}; + } + return {{"error", {{"code", -32603}, {"message", "transport closed"}}}}; +} + +bool server_mcp_transport::ensure_init(const std::function & should_stop) { + if (initialized) { + return true; + } + + json init_req = { + {"jsonrpc", "2.0"}, + {"id", next_id++}, + {"method", "initialize"}, + {"params", { + {"protocolVersion", MCP_PROTOCOL_VERSION}, + {"capabilities", json::object()}, + {"clientInfo", {{"name", "llama.cpp"}, {"version", "1.0"}}}, + }}, + }; + json resp = send_rpc(init_req, should_stop); + if (!resp.contains("result")) { + last_error = "initialize failed: " + rpc_error_message(resp); + return false; + } + + // notifications/initialized: no id, no reply expected + json notif = {{"jsonrpc", "2.0"}, {"method", "notifications/initialized"}}; + to_server.write(notif.dump()); + + initialized = true; + return true; +} + +std::vector server_mcp_transport::list_tools(const std::function & should_stop) { + std::lock_guard lock(rpc_mutex); + if (!ensure_init(should_stop)) { + return {}; + } + if (!tools.empty()) { + return tools; + } + + json req = {{"jsonrpc", "2.0"}, {"id", next_id++}, {"method", "tools/list"}}; + json resp = send_rpc(req, should_stop); + if (!resp.contains("result")) { + last_error = "tools/list failed: " + rpc_error_message(resp); + return {}; + } + + const json & result = resp.at("result"); + if (result.contains("tools") && result.at("tools").is_array()) { + for (const auto & t : result.at("tools")) { + server_mcp_tool_def def; + def.server_name = name; + def.name = t.value("name", ""); + def.description = t.value("description", ""); + if (t.contains("inputSchema")) { + def.input_schema = t.at("inputSchema"); + } + tools.push_back(std::move(def)); + } + } + return tools; +} + +json server_mcp_transport::call_tool(const std::string & tool_name, + const json & arguments, + const std::function & should_stop) { + std::lock_guard lock(rpc_mutex); + if (!ensure_init(should_stop)) { + return {{"error", last_error}}; + } + + json req = { + {"jsonrpc", "2.0"}, + {"id", next_id++}, + {"method", "tools/call"}, + {"params", {{"name", tool_name}, {"arguments", arguments}}}, + }; + json resp = send_rpc(req, should_stop); + if (resp.contains("error")) { + return {{"error", rpc_error_message(resp)}}; + } + if (resp.contains("result")) { + return mcp_result_to_response(resp.at("result")); + } + return {{"error", "invalid response from MCP server"}}; +} + +// +// server_mcp_stdio +// + +struct server_mcp_stdio::process_handle { + subprocess_s sp; + FILE * in = nullptr; // child stdin + FILE * out = nullptr; // child stdout + FILE * err = nullptr; // child stderr +}; + +#if defined(_WIN32) +// config strings are UTF-8 (from JSON) and subprocess.h converts them with CP_UTF8, so inputs must be UTF-8, not the active code page +static std::wstring windows_utf8_to_wide(const std::string & s) { + if (s.empty()) { + return std::wstring(); + } + int n = MultiByteToWideChar(CP_UTF8, 0, s.data(), (int) s.size(), NULL, 0); + if (n <= 0) { + return std::wstring(); + } + std::wstring w((size_t) n, L'\0'); + MultiByteToWideChar(CP_UTF8, 0, s.data(), (int) s.size(), &w[0], n); + return w; +} + +static std::string windows_wide_to_utf8(const wchar_t * s, int len /* -1 for NUL-terminated */) { + int n = WideCharToMultiByte(CP_UTF8, 0, s, len, NULL, 0, NULL, NULL); + if (n <= 0) { + return std::string(); + } + std::string out((size_t) n, '\0'); + WideCharToMultiByte(CP_UTF8, 0, s, len, &out[0], n, NULL, NULL); + if (len == -1 && !out.empty() && out.back() == '\0') { + out.pop_back(); // drop the terminator WideCharToMultiByte counts for -1 + } + return out; +} +#endif + +static std::string mcp_resolve_command(const std::string & command) { +#if defined(_WIN32) + // For Windows: make sure we handle ".exe" correctly, as well as UTF-8 + std::wstring wcmd = windows_utf8_to_wide(command); + wchar_t buf[MAX_PATH * 4]; + const DWORD cap = (DWORD) (sizeof(buf) / sizeof(buf[0])); + + auto search = [&](const wchar_t * ext) -> std::string { + DWORD n = SearchPathW(NULL, wcmd.c_str(), ext, cap, buf, NULL); + return (n > 0 && n < cap) ? windows_wide_to_utf8(buf, (int) n) : std::string(); + }; + + std::string found = search(NULL); // exact path / already-extensioned / .exe on PATH + if (!found.empty()) { + return found; + } + + std::wstring pathext; + DWORD need = GetEnvironmentVariableW(L"PATHEXT", NULL, 0); + if (need > 0) { + pathext.resize(need); + DWORD got = GetEnvironmentVariableW(L"PATHEXT", &pathext[0], need); + pathext.resize(got); + } + if (pathext.empty()) { + pathext = L".COM;.EXE;.BAT;.CMD"; + } + for (size_t start = 0; start <= pathext.size();) { + size_t sep = pathext.find(L';', start); + std::wstring ext = pathext.substr(start, sep == std::wstring::npos ? std::wstring::npos : sep - start); + if (!ext.empty()) { + found = search(ext.c_str()); + if (!found.empty()) { + return found; + } + } + if (sep == std::wstring::npos) { + break; + } + start = sep + 1; + } + return command; // give up and let subprocess.h report the spawn error +#else + return command; +#endif // _WIN32 +} + +static std::vector mcp_parent_env() { + std::vector env; +#if defined(_WIN32) + LPWCH block = GetEnvironmentStringsW(); + if (block) { + for (LPWCH e = block; *e; e += wcslen(e) + 1) { + env.emplace_back(windows_wide_to_utf8(e, -1)); + } + FreeEnvironmentStringsW(block); + } +#else + if (environ) { + for (char ** e = environ; *e; ++e) { + env.emplace_back(*e); + } + } +#endif + return env; +} + +// parent env with the config overrides applied, in "KEY=VALUE" form +static std::vector mcp_build_env(const std::map & overrides) { + std::vector env; + for (auto & e : mcp_parent_env()) { + size_t eq = e.find('='); + std::string key = eq == std::string::npos ? e : e.substr(0, eq); + if (overrides.find(key) == overrides.end()) { + env.push_back(e); + } + } + for (auto & [k, v] : overrides) { + env.push_back(k + "=" + v); + } + return env; +} + +server_mcp_stdio::server_mcp_stdio(const server_mcp_server_config & config) : config(config) { + name = config.name; + timeout_ms = config.timeout_ms; + // bound the reply queue: send_rpc only drains during a call, so unsolicited notifications would otherwise grow it without limit + from_server.max_size = 65536; +} + +server_mcp_stdio::~server_mcp_stdio() { + join_pumps(); +} + +bool server_mcp_stdio::start() { + std::vector argv_s; + argv_s.push_back(mcp_resolve_command(config.command)); + argv_s.insert(argv_s.end(), config.args.begin(), config.args.end()); + + int options = subprocess_option_no_window | subprocess_option_search_user_path; + std::vector envp_s; + if (config.env.empty()) { + options |= subprocess_option_inherit_environment; + } else { + envp_s = mcp_build_env(config.env); + } + + auto to_ptrs = [](std::vector & v) { + std::vector p; + p.reserve(v.size() + 1); + for (auto & s : v) { + p.push_back(s.c_str()); + } + p.push_back(nullptr); + return p; + }; + auto argv = to_ptrs(argv_s); + auto envp = to_ptrs(envp_s); + + auto handle = std::make_unique(); + int rc = subprocess_create_ex(argv.data(), options, + config.env.empty() ? nullptr : envp.data(), + config.cwd.empty() ? nullptr : config.cwd.c_str(), + &handle->sp); + if (rc != 0) { + SRV_WRN("MCP '%s': failed to spawn '%s'\n", config.name.c_str(), config.command.c_str()); + return false; + } + handle->in = subprocess_stdin(&handle->sp); + handle->out = subprocess_stdout(&handle->sp); + handle->err = subprocess_stderr(&handle->sp); + + proc = std::move(handle); + running.store(true); + reader = std::thread([this] { reader_loop(); }); + writer = std::thread([this] { writer_loop(); }); + errlog = std::thread([this] { errlog_loop(); }); + return true; +} + +void server_mcp_stdio::close() { + join_pumps(); +} + +bool server_mcp_stdio::is_alive() const { + return running.load(); +} + +std::string server_mcp_stdio::diagnostics() { + std::string out; + { + std::lock_guard lock(rpc_mutex); // last_error is written by send_rpc's callers + out = last_error; + } + std::lock_guard lk(err_mu); + if (!err_tail.empty()) { + if (!out.empty()) { + out += "; "; + } + out += "last stderr: " + err_tail; + } + return out; +} + +void server_mcp_stdio::reader_loop() { + mcp_pump_ndjson(proc->out, running, [this](std::string && line) { + return from_server.write(std::move(line)); // false => consumer gone, stop + }); + running.store(false); + to_server.close_write(); // stop the writer + from_server.close_write(); // EOF to any waiting caller +} + +// write all of `data` to child stdin, non-blocking and polled so teardown never hangs (a grandchild can hold the read end of a full pipe open). returns false on error/close/shutdown. +static bool mcp_write_all(FILE * f, const std::string & data, std::atomic & running) { + if (!f) { + return false; + } + size_t total = 0; +#if defined(_WIN32) + HANDLE h = (HANDLE) _get_osfhandle(_fileno(f)); + DWORD nowait = PIPE_NOWAIT; + SetNamedPipeHandleState(h, &nowait, NULL, NULL); + while (total < data.size() && running.load()) { + DWORD written = 0; + BOOL ok = WriteFile(h, data.data() + total, (DWORD) (data.size() - total), &written, NULL); + if (ok && written > 0) { + total += written; + continue; + } + if (!ok) { + DWORD err = GetLastError(); + if (err != ERROR_NO_DATA && err != ERROR_PIPE_BUSY) { + return false; + } + } + // backpressure (pipe full) is rare for small JSON-RPC frames; sleep rather than spin. + // no writable-wait exists for a PIPE_NOWAIT anonymous pipe, so this polls like the POSIX poll() path. + std::this_thread::sleep_for(std::chrono::milliseconds(10)); + } +#else + int fd = fileno(f); + int fl = fcntl(fd, F_GETFL, 0); + if (fl >= 0) { + fcntl(fd, F_SETFL, fl | O_NONBLOCK); + } + while (total < data.size() && running.load()) { + ssize_t n = write(fd, data.data() + total, data.size() - total); + if (n > 0) { + total += (size_t) n; + continue; + } + if (n == 0) { + return false; + } + if (errno == EINTR) { + continue; + } + if (errno != EAGAIN && errno != EWOULDBLOCK) { + return false; + } + struct pollfd pfd; + pfd.fd = fd; + pfd.events = POLLOUT; + pfd.revents = 0; + int pr = poll(&pfd, 1, 50); + if (pr < 0) { + if (errno == EINTR) { + continue; + } + return false; + } + if (pfd.revents & (POLLERR | POLLNVAL | POLLHUP)) { + return false; + } + } +#endif + return total == data.size(); +} + +void server_mcp_stdio::writer_loop() { + auto should_stop = [this] { return !running.load(); }; + std::string msg; + while (to_server.read(msg, should_stop)) { + msg.push_back('\n'); + if (!mcp_write_all(proc->in, msg, running)) { + break; // child gone or shutting down + } + } + running.store(false); + to_server.close_read(); // fail fast on any further send_rpc write + from_server.close_write(); // wake any caller waiting for a reply +} + +void server_mcp_stdio::errlog_loop() { + static constexpr size_t ERR_TAIL_MAX = 4096; + // drain stderr (an undrained pipe blocks the child): + // log it, and keep a bounded tail for reporting when the server dies + mcp_pump_ndjson(proc->err, running, [this](std::string && line) { + SRV_DBG("MCP '%s' stderr: %s\n", name.c_str(), line.c_str()); + std::lock_guard lk(err_mu); + err_tail += line; + err_tail += '\n'; + if (err_tail.size() > ERR_TAIL_MAX) { + err_tail.erase(0, err_tail.size() - ERR_TAIL_MAX); + } + return true; + }); +} + +void server_mcp_stdio::join_pumps() { + if (!proc) { + return; + } + running.store(false); + to_server.close_write(); // wake the writer if it waits for a message + from_server.close_write(); // wake any caller waiting for a reply + + subprocess_terminate(&proc->sp); // child death unblocks the blocked fread/fwrite + + if (writer.joinable()) writer.join(); + if (reader.joinable()) reader.join(); + if (errlog.joinable()) errlog.join(); + + subprocess_join(&proc->sp, nullptr); // reap the child: destroy() never waits, so the pid would stay a zombie for the process lifetime + subprocess_destroy(&proc->sp); // safe now: no thread touches the FILE* anymore + proc.reset(); +} + + +// +// server_mcp +// + +static constexpr int MCP_COOLDOWN_SECONDS = 5; +static constexpr int MCP_WARMUP_TIMEOUT_SECONDS = 10; // cap per-server tool discovery at startup + +server_mcp::~server_mcp() { + shutdown(); + + std::vector> to_close; + { + std::lock_guard lock(mutex); + for (auto & [name, t] : transports) { + to_close.push_back(std::move(t)); + } + transports.clear(); + } + for (auto & t : to_close) { + t->close(); + } +} + +std::shared_ptr server_mcp::create_transport(const server_mcp_server_config & cfg) { + return std::make_shared(cfg); +} + +void server_mcp::shutdown() { + stopping.store(true); +} + +const server_mcp_server_config * server_mcp::find_config(const std::string & name) const { + for (const auto & c : configs) { + if (c.name == name) { + return &c; + } + } + return nullptr; +} + +void server_mcp::start(const common_params & params) { + auto append = [this](const std::string & json_str) { + try { + auto parsed = server_mcp_server_config::parse_from_json(json_str); + if (parsed.empty()) { + SRV_WRN("%s", "MCP config: no servers found in JSON\n"); + } + for (auto & p : parsed) { + // names must be unique across both config sources: get_or_create / find_config key on the name + if (find_config(p.name)) { + SRV_WRN("MCP config: duplicate server name '%s', skipping\n", p.name.c_str()); + continue; + } + configs.push_back(std::move(p)); + } + } catch (const std::exception & e) { + throw std::runtime_error(std::string("failed to parse MCP config JSON: ") + e.what()); + } + }; + if (!params.mcp_servers_config.empty()) { + std::ifstream f = fs_open_ifstream(params.mcp_servers_config, std::ios::in); + if (!f) { + throw std::runtime_error("failed to open MCP config file: " + params.mcp_servers_config); + } + std::stringstream ss; + ss << f.rdbuf(); + append(ss.str()); + } + if (!params.mcp_servers_json.empty()) { + append(params.mcp_servers_json); + } + + if (configs.empty()) { + return; + } + + std::vector discovered; + for (const auto & cfg : configs) { + auto t = create_transport(cfg); + if (!t->start()) { + SRV_WRN("MCP warmup: failed to spawn '%s': %s\n", cfg.name.c_str(), t->diagnostics().c_str()); + continue; + } + // bound warmup per server so an unresponsive one can't stall startup for the full per-call timeout + const auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(MCP_WARMUP_TIMEOUT_SECONDS); + auto should_stop = [this, deadline]() { + return stopping.load() || std::chrono::steady_clock::now() >= deadline; + }; + auto tools = t->list_tools(should_stop); + SRV_INF("MCP warmup: '%s' discovered %zu tools\n", cfg.name.c_str(), tools.size()); + discovered.insert(discovered.end(), tools.begin(), tools.end()); + t->close(); + } + + std::lock_guard lock(mutex); + registry.swap(discovered); +} + +std::vector server_mcp::list_tools() const { + std::lock_guard lock(mutex); + return registry; +} + +json server_mcp::call_tool(const std::string & server_name, + const std::string & tool_name, + const json & arguments, + const std::function & should_stop) { + auto transport = get_or_create(server_name); + if (!transport) { + return {{"error", "MCP server unavailable: " + server_name}}; + } + + auto stop = [this, &should_stop]() { + return stopping.load() || (should_stop && should_stop()); + }; + return transport->call_tool(tool_name, arguments, stop); +} + +std::shared_ptr server_mcp::get_or_create(const std::string & name) { + std::vector> to_close; // closed after unlock + std::shared_ptr result; + + { + std::lock_guard lock(mutex); + if (stopping.load()) { + return nullptr; + } + + auto now = std::chrono::steady_clock::now(); + auto dead_it = dead_servers.find(name); + if (dead_it != dead_servers.end()) { + if (now < dead_it->second) { + return nullptr; + } + dead_servers.erase(dead_it); + } + + auto it = transports.find(name); + if (it != transports.end()) { + if (it->second->is_alive()) { + return it->second; + } + SRV_WRN("MCP '%s' is no longer alive: %s\n", name.c_str(), it->second->diagnostics().c_str()); + to_close.push_back(std::move(it->second)); + transports.erase(it); + } + + const server_mcp_server_config * cfg = find_config(name); + if (cfg) { + auto fresh = create_transport(*cfg); + if (fresh->start() && fresh->is_alive()) { + transports[name] = fresh; + result = fresh; + } else { + SRV_WRN("MCP '%s': failed to start: %s\n", name.c_str(), fresh->diagnostics().c_str()); + to_close.push_back(std::move(fresh)); + dead_servers[name] = now + std::chrono::seconds(MCP_COOLDOWN_SECONDS); + } + } + } + + for (auto & t : to_close) { + t->close(); // blocking call, no leaks + } + + return result; +} + diff --git a/tools/server/server-mcp.h b/tools/server/server-mcp.h new file mode 100644 index 000000000000..c7f33a3f797d --- /dev/null +++ b/tools/server/server-mcp.h @@ -0,0 +1,176 @@ +#pragma once + +#include "server-common.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +// +// Configuration (Cursor-compatible "mcpServers" JSON) +// + +struct server_mcp_server_config { + std::string name; // config key, e.g. "filesystem" + std::string command; + std::vector args; + std::map env; // merged over the parent env + std::string cwd; + int timeout_ms = 30000; // per-tool-call timeout + + // throw on parse errors; missing "mcpServers" yields an empty list; entries without a "command" are skipped + static std::vector parse_from_json(const std::string & json_str); + static std::vector parse_cursor_format(const json & j); +}; + +// a tool advertised by an MCP server +struct server_mcp_tool_def { + std::string server_name; + std::string name; // bare tool name, no "_" prefix + std::string description; + json input_schema; // JSON Schema for the arguments, or null +}; + +// +// server_mcp_transport: one MCP server session. +// +// caller --send_rpc--> to_server --[writer]--> framing --> server +// caller <--send_rpc-- from_server <--[reader]-- framing <-- server +// +// each queue item is one complete serialized JSON message. +// subclass owns byte I/O and framing; base owns JSON and the JSON-RPC session (handshake, id correlation). +// + +struct server_mcp_transport { + std::string name; + int timeout_ms = 30000; + + server_pipe to_server; // serialized messages we send to the server + server_pipe from_server; // serialized messages read from the server + + virtual ~server_mcp_transport() = default; + + virtual bool start() = 0; + virtual void close() = 0; // blocking and idempotent + virtual bool is_alive() const = 0; // never blocks behind an in-flight send_rpc() + + // human-readable diagnostics for logging when the transport fails/dies + // (example: last RPC error, plus any transport-specific detail) + // may run on a different thread than send_rpc(), so last_error is read under rpc_mutex + virtual std::string diagnostics() { + std::lock_guard lock(rpc_mutex); + return last_error; + } + + std::vector list_tools(const std::function & should_stop); + + json call_tool(const std::string & tool_name, + const json & arguments, + const std::function & should_stop); + +protected: + // per-transport: send_rpc() holds it across the reply wait, so sharing it would stall every server behind one slow call. guards all members below. + std::mutex rpc_mutex; + uint64_t next_id = 1; // reset to 1 per (re)spawn + bool initialized = false; + std::string last_error; + std::vector tools; + + // both assume rpc_mutex is already held by the public caller + bool ensure_init(const std::function & should_stop); // initialize handshake, once + json send_rpc(const json & request, const std::function & should_stop); // returns the reply or an {"error": ...} +}; + +// +// server_mcp_stdio: child process, NDJSON JSON-RPC over stdio (stderr drained to the debug log) +// + +struct server_mcp_stdio : server_mcp_transport { + explicit server_mcp_stdio(const server_mcp_server_config & config); + ~server_mcp_stdio() override; + + bool start() override; + void close() override; + bool is_alive() const override; + std::string diagnostics() override; + +private: + server_mcp_server_config config; + + // defined in the .cpp so stays out of this header + struct process_handle; + std::unique_ptr proc; + + std::thread reader; // child stdout -> NDJSON de-framing -> from_server + std::thread writer; // to_server -> NDJSON framing -> child stdin + std::thread errlog; // child stderr -> debug log (must be drained or the child blocks) + + // cleared by close() or by the reader on stdout EOF; read without rpc_mutex + std::atomic running{false}; + + // bounded tail of the child's stderr, for diagnostics when it dies + std::mutex err_mu; + std::string err_tail; + + void reader_loop(); + void writer_loop(); + void errlog_loop(); + void join_pumps(); +}; + +// +// server_mcp +// declare before the HTTP context so it outlives every /tools handler. +// + +class server_mcp { +public: + server_mcp() = default; + ~server_mcp(); + + // parse the MCP config from params (file and/or inline JSON), + // then spawn each server once, list its tools, and shut it down + // throws on config parse errors; spawn failures are logged. + void start(const common_params & params); + + // true until start() has parsed at least one server from the config + bool empty() const { return configs.empty(); } + + std::vector list_tools() const; + + // lazily (re)spawns the transport. returns the MCP result or an {"error": ...}. should_stop is OR-ed with the manager's cancel flag. + json call_tool(const std::string & server_name, + const std::string & tool_name, + const json & arguments, + const std::function & should_stop = nullptr); + + // flip the cancel flag so in-flight calls return; blocking teardown is in the destructor. call before the HTTP server drains. + // note: multiple calls are idempotent + void shutdown(); + +private: + std::vector configs; + + mutable std::mutex mutex; // guards transports, dead_servers, registry + + // shared_ptr: call_tool() hands a transport to the caller and drops the lock for the blocking RPC, so a concurrent evict/respawn must not destroy it mid-call + std::map> transports; + std::map dead_servers; // spawn-failure cooldown + std::vector registry; + + std::atomic stopping{false}; + + const server_mcp_server_config * find_config(const std::string & name) const; + + // the only place that names a concrete transport + std::shared_ptr create_transport(const server_mcp_server_config & cfg); + + // nullptr during cooldown or shutdown + std::shared_ptr get_or_create(const std::string & name); +}; diff --git a/tools/server/server-models.cpp b/tools/server/server-models.cpp index 9eac58e9df27..ba63788146fe 100644 --- a/tools/server/server-models.cpp +++ b/tools/server/server-models.cpp @@ -1944,53 +1944,6 @@ void server_models_routes::init_routes() { // server_http_proxy // -// simple implementation of a pipe -// used for streaming data between threads -template -struct pipe_t { - std::mutex mutex; - std::condition_variable cv; - std::queue queue; - std::atomic writer_closed{false}; - std::atomic reader_closed{false}; - void close_write() { - writer_closed.store(true, std::memory_order_relaxed); - cv.notify_all(); - } - void close_read() { - reader_closed.store(true, std::memory_order_relaxed); - cv.notify_all(); - } - bool read(T & output, const std::function & should_stop) { - std::unique_lock lk(mutex); - constexpr auto poll_interval = std::chrono::milliseconds(500); - while (true) { - if (!queue.empty()) { - output = std::move(queue.front()); - queue.pop(); - return true; - } - if (writer_closed.load()) { - return false; // clean EOF - } - if (should_stop()) { - close_read(); // signal broken pipe to writer - return false; // cancelled / reader no longer alive - } - cv.wait_for(lk, poll_interval); - } - } - bool write(T && data) { - std::lock_guard lk(mutex); - if (reader_closed.load()) { - return false; // broken pipe - } - queue.push(std::move(data)); - cv.notify_one(); - return true; - } -}; - static std::string to_lower_copy(const std::string & value) { std::string lowered(value.size(), '\0'); std::transform(value.begin(), value.end(), lowered.begin(), [](unsigned char c) { return std::tolower(c); }); @@ -2100,7 +2053,7 @@ server_http_proxy::server_http_proxy( ) { // shared between reader and writer threads auto cli = std::make_shared(host, port); - auto pipe = std::make_shared>(); + auto pipe = std::make_shared>(); if (scheme == "https") { #ifdef CPPHTTPLIB_OPENSSL_SUPPORT diff --git a/tools/server/server-tools.cpp b/tools/server/server-tools.cpp index a82a3d602590..2af44e49b9d4 100644 --- a/tools/server/server-tools.cpp +++ b/tools/server/server-tools.cpp @@ -10,10 +10,10 @@ #include #include #include -#include #include #include #include +#include namespace fs = std::filesystem; @@ -25,7 +25,7 @@ json server_tool::to_json() const { return { {"display_name", display_name}, {"tool", name}, - {"type", "builtin"}, + {"type", type()}, {"permissions", json{ {"write", permission_write} }}, @@ -1129,6 +1129,49 @@ struct server_tools_res : server_http_res { } }; +// +// server_mcp_tool: exposes one tool from a running MCP server as a server_tool. +// +struct server_mcp_tool : server_tool { + std::string server_name; + std::string tool_name; + server_mcp_tool_def def; + server_mcp & mcp_mgr; + + server_mcp_tool(server_mcp_tool_def d, server_mcp & mgr) + : server_name(d.server_name) + , tool_name(d.name) + , def(std::move(d)) + , mcp_mgr(mgr) + { + name = server_name + "_" + tool_name; + display_name = name; + permission_write = false; + support_stream = false; + } + + std::string type() const override { return "mcp"; } + + json get_definition() const override { + json schema = def.input_schema; + if (schema.is_null() || !schema.is_object()) { + schema = json::object(); + } + return { + {"type", "function"}, + {"function", { + {"name", name}, + {"description", def.description}, + {"parameters", schema}, + }}, + }; + } + + json invoke(json params, server_tool::stream *) const override { + return mcp_mgr.call_tool(server_name, tool_name, params); + } +}; + static server_tool & find_tool(std::vector> & tools, const std::string & name, bool require_stream) { for (auto & t : tools) { if (t->name == name) { @@ -1157,7 +1200,8 @@ static std::vector> build_tools() { return tools; } -void server_tools::setup(const std::vector & enabled_tools) { +void server_tools::setup(const std::vector & enabled_tools, + server_mcp & mcp_mgr) { if (!enabled_tools.empty()) { std::unordered_set enabled_set(enabled_tools.begin(), enabled_tools.end()); auto all_tools = build_tools(); @@ -1188,6 +1232,29 @@ void server_tools::setup(const std::vector & enabled_tools) { } } + // append MCP tools, skipping any that collide with a built-in or another MCP tool of the same "_" name + if (!mcp_mgr.empty()) { + std::unordered_set seen_names; + for (auto & t : tools) { + seen_names.insert(t->name); + } + size_t n_added = 0; + for (const auto & def : mcp_mgr.list_tools()) { + std::string mcp_name = def.server_name + "_" + def.name; + if (seen_names.count(mcp_name)) { + SRV_WRN("MCP tool \"%s\" from server \"%s\" collides with an existing tool, skipping\n", + mcp_name.c_str(), def.server_name.c_str()); + continue; + } + seen_names.insert(mcp_name); + tools.push_back(std::make_unique(def, mcp_mgr)); + n_added++; + } + if (n_added > 0) { + SRV_INF("Added %zu MCP tools\n", n_added); + } + } + handle_get = [this](const server_http_req &) -> server_http_res_ptr { auto res = std::make_unique(); try { diff --git a/tools/server/server-tools.h b/tools/server/server-tools.h index 6f6528f484f8..601399ee9392 100644 --- a/tools/server/server-tools.h +++ b/tools/server/server-tools.h @@ -3,9 +3,11 @@ #include "server-common.h" #include "server-http.h" #include "server-queue.h" +#include "server-mcp.h" #include #include +#include struct server_tool { std::string name; @@ -15,6 +17,7 @@ struct server_tool { virtual ~server_tool() = default; virtual json get_definition() const = 0; + virtual std::string type() const { return "builtin"; } struct stream { server_response & qr; @@ -34,7 +37,8 @@ struct server_tools { server_response queue_res; std::atomic res_id{0}; - void setup(const std::vector & enabled_tools); + void setup(const std::vector & enabled_tools, + server_mcp & mcp_mgr); server_http_context::handler_t handle_get; server_http_context::handler_t handle_post; diff --git a/tools/server/server.cpp b/tools/server/server.cpp index 20effbb14851..b6fef99e8747 100644 --- a/tools/server/server.cpp +++ b/tools/server/server.cpp @@ -88,6 +88,11 @@ static server_http_context::handler_t ex_wrapper(server_http_context::handler_t int llama_server(int argc, char ** argv) { std::setlocale(LC_NUMERIC, "C"); +#ifndef _WIN32 + // Ignore SIGPIPE so the server does not crash if an MCP child exits while we are writing to its stdin + signal(SIGPIPE, SIG_IGN); +#endif + // own arguments required by this example common_params params; @@ -157,6 +162,9 @@ int llama_server(common_params & params, int argc, char ** argv) { params.model_alias.insert(model_name); } + // note: this is guaranteed to out-live ctx_http and tools + server_mcp mcp_mgr; + // struct that contains llama context and inference server_context ctx_server; @@ -326,17 +334,28 @@ int llama_server(common_params & params, int argc, char ** argv) { ctx_http.post("/cors-proxy", ex_wrapper(res_403)); } - // EXPERIMENTAL built-in tools - if (!params.server_tools.empty()) { + try { + mcp_mgr.start(params); + } catch (const std::exception & e) { + SRV_ERR("MCP starting failed: %s\n", e.what()); + return 1; + } + + if (!params.server_tools.empty() || !mcp_mgr.empty()) { try { - tools.setup(params.server_tools); + tools.setup(params.server_tools, mcp_mgr); } catch (const std::exception & e) { SRV_ERR("tools setup failed: %s\n", e.what()); return 1; } ctx_http.get ("/tools", ex_wrapper(tools.handle_get)); ctx_http.post("/tools", ex_wrapper(tools.handle_post)); - warn_names.push_back("built-in tools (experimental)"); + if (!params.server_tools.empty()) { + warn_names.push_back("built-in tools (experimental)"); + } + if (!mcp_mgr.empty()) { + warn_names.push_back("MCP servers (experimental)"); + } } else { ctx_http.get ("/tools", ex_wrapper(res_403)); ctx_http.post("/tools", ex_wrapper(res_403)); @@ -378,7 +397,7 @@ int llama_server(common_params & params, int argc, char ** argv) { if (is_router_server) { SRV_INF("%s", "starting server in router mode. models will be automatically loaded on-demand\n"); - clean_up = [&models_routes]() { + clean_up = [&models_routes, &mcp_mgr]() { SRV_INF("%s: cleaning up before exit...\n", __func__); // stop the session GC first, it finalizes live sessions and wakes pending readers server_stream_session_manager_stop(); @@ -386,6 +405,7 @@ int llama_server(common_params & params, int argc, char ** argv) { models_routes->stopping.store(true); // maybe redundant, but just to be safe models_routes->models.unload_all(); } + mcp_mgr.shutdown(); llama_backend_free(); }; @@ -401,17 +421,19 @@ int llama_server(common_params & params, int argc, char ** argv) { // important to disconnect any SSE clients models_routes->stopping.store(true); } + mcp_mgr.shutdown(); ctx_http.stop(); }; } else { // setup clean up function, to be called before exit - clean_up = [&ctx_http, &ctx_server]() { + clean_up = [&ctx_http, &ctx_server, &mcp_mgr]() { SRV_INF("%s: cleaning up before exit...\n", __func__); // stop the session GC first, it finalizes live sessions and wakes pending readers server_stream_session_manager_stop(); ctx_http.stop(); ctx_server.terminate(); + mcp_mgr.shutdown(); llama_backend_free(); }; @@ -444,6 +466,7 @@ int llama_server(common_params & params, int argc, char ** argv) { SRV_INF("%s", "model loaded\n"); shutdown_handler = [&](int) { + mcp_mgr.shutdown(); // this will unblock start_loop() ctx_server.terminate(); }; diff --git a/tools/server/tests/fixtures/mcp_burst_server.py b/tools/server/tests/fixtures/mcp_burst_server.py new file mode 100644 index 000000000000..22892d9a1d72 --- /dev/null +++ b/tools/server/tests/fixtures/mcp_burst_server.py @@ -0,0 +1,118 @@ +#!/usr/bin/env python3 +""" +Minimal MCP server that writes notification + response in a single write() with no flush. +This reproduces the buffering bug where read_message() can strand the response. +""" +import json +import sys +import os + +TOOLS = [ + { + "name": "echo", + "description": "Echo back the input message", + "inputSchema": { + "type": "object", + "properties": { + "message": {"type": "string"} + }, + "required": ["message"] + } + } +] + +def handle_initialize(params, req_id): + return { + "jsonrpc": "2.0", + "id": req_id, + "result": { + "protocolVersion": "2024-11-05", + "capabilities": {"tools": {}}, + "serverInfo": {"name": "burst-test", "version": "1.0"} + } + } + +def handle_tools_list(params, req_id): + return { + "jsonrpc": "2.0", + "id": req_id, + "result": {"tools": TOOLS} + } + +def handle_tools_call(params, req_id): + tool_name = params.get("name") + arguments = params.get("arguments", {}) + + if tool_name == "echo": + message = arguments.get("message", "") + notif = { + "jsonrpc": "2.0", + "method": "notifications/progress", + "params": {"progress": 50, "total": 100} + } + response = { + "jsonrpc": "2.0", + "id": req_id, + "result": { + "content": [{"type": "text", "text": f"echo: {message}"}] + } + } + # Single os.write() call: both lines land in one pipe packet atomically. + # This is the key difference from mcp_malformed_server.py which flushes between writes. + data = (json.dumps(notif) + "\n" + json.dumps(response) + "\n").encode("utf-8") + os.write(sys.stdout.fileno(), data) + return None # already written + else: + response = { + "jsonrpc": "2.0", + "id": req_id, + "error": {"code": -32602, "message": f"Unknown tool: {tool_name}"} + } + return response + +HANDLERS = { + "initialize": handle_initialize, + "tools/list": handle_tools_list, + "tools/call": handle_tools_call, +} + +def main(): + # Use line-buffered text mode for regular responses, but the burst write + # uses os.write() directly to guarantee a single kernel write(). + sys.stdout = os.fdopen(sys.stdout.fileno(), "w", buffering=1) + sys.stderr = os.fdopen(sys.stderr.fileno(), "w", buffering=1) + + for line in sys.stdin: + line = line.strip() + if not line: + continue + try: + request = json.loads(line) + except json.JSONDecodeError: + continue + + method = request.get("method") + req_id = request.get("id") + params = request.get("params", {}) + + # JSON-RPC 2.0: a message without an id is a notification and must not receive a response + if req_id is None: + continue + + handler = HANDLERS.get(method) + if handler: + response = handler(params, req_id) + if response is not None: + sys.stdout.write(json.dumps(response) + "\n") + sys.stdout.flush() + else: + response = { + "jsonrpc": "2.0", + "id": req_id, + "error": {"code": -32601, "message": f"Method not found: {method}"} + } + sys.stdout.write(json.dumps(response) + "\n") + sys.stdout.flush() + +if __name__ == "__main__": + main() diff --git a/tools/server/tests/fixtures/mcp_crash_server.py b/tools/server/tests/fixtures/mcp_crash_server.py new file mode 100644 index 000000000000..8dffdc61c0e9 --- /dev/null +++ b/tools/server/tests/fixtures/mcp_crash_server.py @@ -0,0 +1,114 @@ +#!/usr/bin/env python3 +""" +MCP server that crashes after receiving a specific tool call. +""" +import json +import sys +import os + +def handle_initialize(params, req_id): + return { + "jsonrpc": "2.0", + "id": req_id, + "result": { + "protocolVersion": "2024-11-05", + "capabilities": {"tools": {}}, + "serverInfo": {"name": "crash-test", "version": "1.0"} + } + } + +def handle_tools_list(params, req_id): + return { + "jsonrpc": "2.0", + "id": req_id, + "result": { + "tools": [ + { + "name": "echo", + "description": "Echo back the input message", + "inputSchema": { + "type": "object", + "properties": { + "message": {"type": "string"} + } + } + }, + { + "name": "crash", + "description": "Crash the server", + "inputSchema": { + "type": "object", + "properties": {} + } + } + ] + } + } + +def handle_tools_call(params, req_id): + tool_name = params.get("name") + arguments = params.get("arguments", {}) + + if tool_name == "echo": + message = arguments.get("message", "") + return { + "jsonrpc": "2.0", + "id": req_id, + "result": { + "content": [{"type": "text", "text": f"echo: {message}"}] + } + } + elif tool_name == "crash": + # Send a partial response then exit + sys.stdout.write(json.dumps({"jsonrpc": "2.0", "id": req_id, "result": {"content": [{"type": "text", "text": "crashing..."}]}}) + "\n") + sys.stdout.flush() + os._exit(1) + else: + return { + "jsonrpc": "2.0", + "id": req_id, + "error": {"code": -32602, "message": f"Unknown tool: {tool_name}"} + } + +HANDLERS = { + "initialize": handle_initialize, + "tools/list": handle_tools_list, + "tools/call": handle_tools_call, +} + +def main(): + sys.stdout = os.fdopen(sys.stdout.fileno(), "w", buffering=1) + sys.stderr = os.fdopen(sys.stderr.fileno(), "w", buffering=1) + + for line in sys.stdin: + line = line.strip() + if not line: + continue + try: + request = json.loads(line) + except json.JSONDecodeError: + continue + + method = request.get("method") + req_id = request.get("id") + params = request.get("params", {}) + + # JSON-RPC 2.0: a message without an id is a notification and must not receive a response + if req_id is None: + continue + + handler = HANDLERS.get(method) + if handler: + response = handler(params, req_id) + else: + response = { + "jsonrpc": "2.0", + "id": req_id, + "error": {"code": -32601, "message": f"Method not found: {method}"} + } + + sys.stdout.write(json.dumps(response) + "\n") + sys.stdout.flush() + +if __name__ == "__main__": + main() diff --git a/tools/server/tests/fixtures/mcp_echo_server.py b/tools/server/tests/fixtures/mcp_echo_server.py new file mode 100755 index 000000000000..7acfb358881a --- /dev/null +++ b/tools/server/tests/fixtures/mcp_echo_server.py @@ -0,0 +1,164 @@ +#!/usr/bin/env python3 +""" +Minimal MCP server for testing. +Implements JSON-RPC 2.0 over stdio (line-delimited JSON). +""" +import json +import sys +import os + +# Ensure we use python3 from the current environment +if sys.platform == "win32": + # On Windows, we need to use the same python interpreter + pass + +TOOLS = [ + { + "name": "echo", + "description": "Echo back the input message", + "inputSchema": { + "type": "object", + "properties": { + "message": {"type": "string", "description": "Message to echo"} + }, + "required": ["message"] + } + }, + { + "name": "add", + "description": "Add two numbers", + "inputSchema": { + "type": "object", + "properties": { + "a": {"type": "number"}, + "b": {"type": "number"} + }, + "required": ["a", "b"] + } + }, + { + "name": "fail_once", + "description": "Fails on first call, succeeds on subsequent calls", + "inputSchema": { + "type": "object", + "properties": {} + } + } +] + +_state = {"fail_once_called": False} + +def handle_initialize(params, req_id): + return { + "jsonrpc": "2.0", + "id": req_id, + "result": { + "protocolVersion": "2024-11-05", + "capabilities": {"tools": {}}, + "serverInfo": {"name": "echo-test", "version": "1.0"} + } + } + +def handle_tools_list(params, req_id): + return { + "jsonrpc": "2.0", + "id": req_id, + "result": {"tools": TOOLS} + } + +def handle_tools_call(params, req_id): + tool_name = params.get("name") + arguments = params.get("arguments", {}) + + if tool_name == "echo": + message = arguments.get("message", "") + return { + "jsonrpc": "2.0", + "id": req_id, + "result": { + "content": [{"type": "text", "text": f"echo: {message}"}] + } + } + elif tool_name == "add": + a = arguments.get("a", 0) + b = arguments.get("b", 0) + return { + "jsonrpc": "2.0", + "id": req_id, + "result": { + "content": [{"type": "text", "text": str(a + b)}] + } + } + elif tool_name == "fail_once": + if not _state["fail_once_called"]: + _state["fail_once_called"] = True + return { + "jsonrpc": "2.0", + "id": req_id, + "error": {"code": -32000, "message": "transient error"} + } + return { + "jsonrpc": "2.0", + "id": req_id, + "result": { + "content": [{"type": "text", "text": "ok"}] + } + } + else: + return { + "jsonrpc": "2.0", + "id": req_id, + "error": {"code": -32602, "message": f"Unknown tool: {tool_name}"} + } + +def handle_ping(params, req_id): + return { + "jsonrpc": "2.0", + "id": req_id, + "result": {} + } + +HANDLERS = { + "initialize": handle_initialize, + "tools/list": handle_tools_list, + "tools/call": handle_tools_call, + "ping": handle_ping, +} + +def main(): + # Use unbuffered output + sys.stdout = os.fdopen(sys.stdout.fileno(), "w", buffering=1) + sys.stderr = os.fdopen(sys.stderr.fileno(), "w", buffering=1) + + for line in sys.stdin: + line = line.strip() + if not line: + continue + try: + request = json.loads(line) + except json.JSONDecodeError: + continue + + method = request.get("method") + req_id = request.get("id") + params = request.get("params", {}) + + # JSON-RPC 2.0: a message without an id is a notification and must not receive a response + if req_id is None: + continue + + handler = HANDLERS.get(method) + if handler: + response = handler(params, req_id) + else: + response = { + "jsonrpc": "2.0", + "id": req_id, + "error": {"code": -32601, "message": f"Method not found: {method}"} + } + + sys.stdout.write(json.dumps(response) + "\n") + sys.stdout.flush() + +if __name__ == "__main__": + main() diff --git a/tools/server/tests/fixtures/mcp_grandchild_server.py b/tools/server/tests/fixtures/mcp_grandchild_server.py new file mode 100644 index 000000000000..2604a77ee80e --- /dev/null +++ b/tools/server/tests/fixtures/mcp_grandchild_server.py @@ -0,0 +1,100 @@ +#!/usr/bin/env python3 +""" +MCP server (NDJSON JSON-RPC over stdio) that spawns a long-lived grandchild which inherits +this process's stdin/stdout/stderr and keeps them open. + +This reproduces the reader-teardown deadlock: killing the direct MCP child (SIGKILL, which is +all subprocess_terminate() does) does NOT close the stdout/stderr pipe write ends, because the +grandchild still holds them. A server that reads those pipes with a blocking read would then +wait forever for an EOF that never arrives, hanging teardown (both warmup shutdown at startup +and process shutdown). The polled, running-aware reader must exit regardless. +""" +import json +import os +import subprocess +import sys + +# Spawn a grandchild that inherits our std handles (fds 0/1/2 = the MCP pipes) and lives well +# past any teardown in the tests. We do NOT redirect its stdio, so it keeps the pipe write ends +# open even after this process is killed. +subprocess.Popen([sys.executable, "-c", "import time; time.sleep(30)"]) + +TOOLS = [ + { + "name": "echo", + "description": "Echo back the input message", + "inputSchema": { + "type": "object", + "properties": {"message": {"type": "string", "description": "Message to echo"}}, + "required": ["message"], + }, + } +] + + +def handle_initialize(params, req_id): + return { + "jsonrpc": "2.0", + "id": req_id, + "result": { + "protocolVersion": "2024-11-05", + "capabilities": {"tools": {}}, + "serverInfo": {"name": "grandchild-test", "version": "1.0"}, + }, + } + + +def handle_tools_list(params, req_id): + return {"jsonrpc": "2.0", "id": req_id, "result": {"tools": TOOLS}} + + +def handle_tools_call(params, req_id): + if params.get("name") == "echo": + message = params.get("arguments", {}).get("message", "") + return { + "jsonrpc": "2.0", + "id": req_id, + "result": {"content": [{"type": "text", "text": f"echo: {message}"}]}, + } + return {"jsonrpc": "2.0", "id": req_id, "error": {"code": -32602, "message": "Unknown tool"}} + + +HANDLERS = { + "initialize": handle_initialize, + "tools/list": handle_tools_list, + "tools/call": handle_tools_call, +} + + +def main(): + sys.stdout = os.fdopen(sys.stdout.fileno(), "w", buffering=1) + sys.stderr = os.fdopen(sys.stderr.fileno(), "w", buffering=1) + + for line in sys.stdin: + line = line.strip() + if not line: + continue + try: + request = json.loads(line) + except json.JSONDecodeError: + continue + + method = request.get("method") + req_id = request.get("id") + params = request.get("params", {}) + + if req_id is None: + continue # notification, no response + + handler = HANDLERS.get(method) + if handler: + response = handler(params, req_id) + else: + response = {"jsonrpc": "2.0", "id": req_id, "error": {"code": -32601, "message": f"Method not found: {method}"}} + + sys.stdout.write(json.dumps(response) + "\n") + sys.stdout.flush() + + +if __name__ == "__main__": + main() diff --git a/tools/server/tests/fixtures/mcp_malformed_server.py b/tools/server/tests/fixtures/mcp_malformed_server.py new file mode 100644 index 000000000000..743333c5fdaa --- /dev/null +++ b/tools/server/tests/fixtures/mcp_malformed_server.py @@ -0,0 +1,113 @@ +#!/usr/bin/env python3 +""" +MCP server that sends malformed responses and notifications during requests. +""" +import json +import sys +import os + +def handle_initialize(params, req_id): + return { + "jsonrpc": "2.0", + "id": req_id, + "result": { + "protocolVersion": "2024-11-05", + "capabilities": {"tools": {}}, + "serverInfo": {"name": "malformed-test", "version": "1.0"} + } + } + +def handle_tools_list(params, req_id): + return { + "jsonrpc": "2.0", + "id": req_id, + "result": { + "tools": [ + { + "name": "echo", + "description": "Echo back the input message", + "inputSchema": { + "type": "object", + "properties": { + "message": {"type": "string"} + } + } + } + ] + } + } + +def handle_tools_call(params, req_id): + tool_name = params.get("name") + arguments = params.get("arguments", {}) + + if tool_name == "echo": + message = arguments.get("message", "") + # Send a notification first (no id field) + notif = { + "jsonrpc": "2.0", + "method": "notifications/progress", + "params": {"progress": 50, "total": 100} + } + sys.stdout.write(json.dumps(notif) + "\n") + sys.stdout.flush() + # Then send the actual response + return { + "jsonrpc": "2.0", + "id": req_id, + "result": { + "content": [{"type": "text", "text": f"echo: {message}"}] + } + } + else: + return { + "jsonrpc": "2.0", + "id": req_id, + "error": {"code": -32602, "message": f"Unknown tool: {tool_name}"} + } + +HANDLERS = { + "initialize": handle_initialize, + "tools/list": handle_tools_list, + "tools/call": handle_tools_call, +} + +def main(): + sys.stdout = os.fdopen(sys.stdout.fileno(), "w", buffering=1) + sys.stderr = os.fdopen(sys.stderr.fileno(), "w", buffering=1) + + for line in sys.stdin: + line = line.strip() + if not line: + continue + try: + request = json.loads(line) + except json.JSONDecodeError: + # Send malformed JSON response + sys.stdout.write("THIS IS NOT JSON\n") + sys.stdout.flush() + continue + + method = request.get("method") + req_id = request.get("id") + params = request.get("params", {}) + + # JSON-RPC 2.0: a message without an id is a notification and must not receive a response + if req_id is None: + continue + + handler = HANDLERS.get(method) + if handler: + response = handler(params, req_id) + else: + response = { + "jsonrpc": "2.0", + "id": req_id, + "error": {"code": -32601, "message": f"Method not found: {method}"} + } + + sys.stdout.write(json.dumps(response) + "\n") + sys.stdout.flush() + +if __name__ == "__main__": + main() diff --git a/tools/server/tests/fixtures/mcp_slow_server.py b/tools/server/tests/fixtures/mcp_slow_server.py new file mode 100644 index 000000000000..7f8e67835acc --- /dev/null +++ b/tools/server/tests/fixtures/mcp_slow_server.py @@ -0,0 +1,132 @@ +#!/usr/bin/env python3 +""" +MCP server that sleeps before responding, for timeout testing. +""" +import json +import sys +import os +import time +import argparse + +TOOLS = [ + { + "name": "sleep", + "description": "Sleep for a given number of seconds", + "inputSchema": { + "type": "object", + "properties": { + "seconds": {"type": "number", "description": "Seconds to sleep"} + }, + "required": ["seconds"] + } + } +] + +def handle_initialize(params, req_id): + return { + "jsonrpc": "2.0", + "id": req_id, + "result": { + "protocolVersion": "2024-11-05", + "capabilities": {"tools": {}}, + "serverInfo": {"name": "slow-test", "version": "1.0"} + } + } + +def handle_tools_list(params, req_id): + return { + "jsonrpc": "2.0", + "id": req_id, + "result": {"tools": TOOLS} + } + +def handle_tools_call(params, req_id): + tool_name = params.get("name") + arguments = params.get("arguments", {}) + + if tool_name == "sleep": + seconds = arguments.get("seconds", 1) + time.sleep(seconds) + return { + "jsonrpc": "2.0", + "id": req_id, + "result": { + "content": [{"type": "text", "text": f"slept {seconds}s"}] + } + } + else: + return { + "jsonrpc": "2.0", + "id": req_id, + "error": {"code": -32602, "message": f"Unknown tool: {tool_name}"} + } + +HANDLERS = { + "initialize": handle_initialize, + "tools/list": handle_tools_list, + "tools/call": handle_tools_call, +} + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--delay", type=float, default=5.0, help="Delay in seconds for sleep tool") + args = parser.parse_args() + + # Override the sleep duration + global handle_tools_call + def handle_tools_call(params, req_id): + tool_name = params.get("name") + arguments = params.get("arguments", {}) + + if tool_name == "sleep": + seconds = arguments.get("seconds", args.delay) + time.sleep(seconds) + return { + "jsonrpc": "2.0", + "id": req_id, + "result": { + "content": [{"type": "text", "text": f"slept {seconds}s"}] + } + } + else: + return { + "jsonrpc": "2.0", + "id": req_id, + "error": {"code": -32602, "message": f"Unknown tool: {tool_name}"} + } + + sys.stdout = os.fdopen(sys.stdout.fileno(), "w", buffering=1) + sys.stderr = os.fdopen(sys.stderr.fileno(), "w", buffering=1) + + for line in sys.stdin: + line = line.strip() + if not line: + continue + try: + request = json.loads(line) + except json.JSONDecodeError: + continue + + method = request.get("method") + req_id = request.get("id") + params = request.get("params", {}) + + # JSON-RPC 2.0: a message without an id is a notification and must not receive a response + if req_id is None: + continue + + handler = HANDLERS.get(method) + if handler: + response = handler(params, req_id) + else: + response = { + "jsonrpc": "2.0", + "id": req_id, + "error": {"code": -32601, "message": f"Method not found: {method}"} + } + + sys.stdout.write(json.dumps(response) + "\n") + sys.stdout.flush() + +if __name__ == "__main__": + main() diff --git a/tools/server/tests/unit/test_mcp_servers.py b/tools/server/tests/unit/test_mcp_servers.py new file mode 100644 index 000000000000..9ad2241bd029 --- /dev/null +++ b/tools/server/tests/unit/test_mcp_servers.py @@ -0,0 +1,718 @@ +#!/usr/bin/env python3 +""" +Tests for MCP server integration via the /tools endpoint. + +Invariants verified: +1. MCP tools appear in /tools listing when configured +2. MCP tools use _ naming +3. MCP tools can be invoked and return correct results +4. Misconfigured MCP servers do not crash the server +5. Multiple MCP servers can be configured simultaneously +6. Warmup populates the tool list at startup +""" +import json +import os +import sys +import tempfile +import time + +import pytest + +from utils import * + +# Path to the test MCP server fixture +FIXTURES_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "fixtures") +MCP_ECHO_SERVER = os.path.join(FIXTURES_DIR, "mcp_echo_server.py") + +server: ServerProcess + + +def _mcp_config_json(servers: dict) -> str: + """Create a JSON config string for --mcp-servers-json.""" + return json.dumps({"mcpServers": servers}) + + +def _start_server_with_mcp(mcp_json: str, **kwargs) -> ServerProcess: + """Helper to start a router server with MCP config.""" + srv = ServerPreset.router() + srv.server_tools = "all" + srv.no_ui = True + srv.server_port = 8085 # avoid conflict with load_all() which uses 8080 + srv.mcp_servers_json = mcp_json + for k, v in kwargs.items(): + setattr(srv, k, v) + srv.start() + return srv + + +def test_mcp_tools_listed_in_tools_endpoint(): + """MCP tools should appear in GET /tools with server:tool naming.""" + global server + mcp_json = _mcp_config_json({ + "echo": { + "command": sys.executable, + "args": [MCP_ECHO_SERVER], + } + }) + server = _start_server_with_mcp(mcp_json) + + try: + res = server.make_request("GET", "/tools") + assert res.status_code == 200, res.body + + tools = res.body + assert isinstance(tools, list), f"Expected list, got {type(tools)}" + + # Find MCP tools - name is in "tool" field or definition.function.name + def get_tool_name(t): + return t.get("tool", "") or t.get("definition", {}).get("function", {}).get("name", "") + + mcp_tools = [t for t in tools if get_tool_name(t).startswith("echo_")] + assert len(mcp_tools) >= 2, f"Expected at least 2 echo_ tools, got {len(mcp_tools)}: {mcp_tools}" + + tool_names = {get_tool_name(t) for t in mcp_tools} + assert "echo_echo" in tool_names + assert "echo_add" in tool_names + + # Verify tool structure + echo_tool = next(t for t in mcp_tools if get_tool_name(t) == "echo_echo") + assert "description" in echo_tool or "definition" in echo_tool + finally: + server.stop() + + +def test_mcp_tool_invocation(): + """MCP tools should be callable via POST /tools and return correct results.""" + global server + mcp_json = _mcp_config_json({ + "echo": { + "command": sys.executable, + "args": [MCP_ECHO_SERVER], + } + }) + server = _start_server_with_mcp(mcp_json) + + try: + # Call echo_echo + res = server.make_request("POST", "/tools", data={ + "tool": "echo_echo", + "params": {"message": "hello world"} + }) + assert res.status_code == 200, res.body + body = res.body + assert "error" not in body, body + # The result format depends on the tool implementation + # For MCP tools, it should contain the tool result + assert "plain_text_response" in body or "result" in body or "content" in body, body + + # Call echo_add + res = server.make_request("POST", "/tools", data={ + "tool": "echo_add", + "params": {"a": 3, "b": 5} + }) + assert res.status_code == 200, res.body + body = res.body + assert "error" not in body, body + finally: + server.stop() + + +def test_mcp_bad_command_does_not_crash(): + """A misconfigured MCP server should not crash the llama-server.""" + global server + mcp_json = _mcp_config_json({ + "nonexistent": { + "command": "this_executable_does_not_exist_12345", + "args": [], + } + }) + server = _start_server_with_mcp(mcp_json) + + try: + # Server should still be healthy + res = server.make_request("GET", "/health") + assert res.status_code == 200, res.body + + # Builtin tools should still work + res = server.make_request("GET", "/tools") + assert res.status_code == 200, res.body + tools = res.body + # Should have builtin tools but no MCP tools from the bad server + mcp_tools = [t for t in tools if t.get("name", "").startswith("nonexistent_")] + assert len(mcp_tools) == 0, f"Expected no nonexistent_ tools, got {mcp_tools}" + finally: + server.stop() + + +def test_mcp_multiple_servers(): + """Multiple MCP servers can be configured simultaneously.""" + global server + mcp_json = _mcp_config_json({ + "echo": { + "command": sys.executable, + "args": [MCP_ECHO_SERVER], + }, + "echo2": { + "command": sys.executable, + "args": [MCP_ECHO_SERVER], + } + }) + server = _start_server_with_mcp(mcp_json) + + try: + res = server.make_request("GET", "/tools") + assert res.status_code == 200, res.body + + tools = res.body + + def get_tool_name(t): + return t.get("tool", "") or t.get("definition", {}).get("function", {}).get("name", "") + + echo_tools = [t for t in tools if get_tool_name(t).startswith("echo_")] + echo2_tools = [t for t in tools if get_tool_name(t).startswith("echo2_")] + + assert len(echo_tools) >= 2, f"Expected echo_ tools, got {echo_tools}" + assert len(echo2_tools) >= 2, f"Expected echo2_ tools, got {echo2_tools}" + finally: + server.stop() + + +def test_mcp_tools_not_listed_when_not_configured(): + """Without MCP config, no MCP tools should appear.""" + global server + server = ServerPreset.router() + server.server_tools = "all" + server.no_ui = True + server.server_port = 8085 + server.start() + + try: + res = server.make_request("GET", "/tools") + assert res.status_code == 200, res.body + + tools = res.body + + def get_tool_name(t): + return t.get("tool", "") or t.get("definition", {}).get("function", {}).get("name", "") + + # Should only have builtin tools, no server: prefixed tools + mcp_tools = [t for t in tools if ":" in get_tool_name(t)] + assert len(mcp_tools) == 0, f"Expected no MCP tools, got {mcp_tools}" + finally: + server.stop() + + +def test_mcp_fail_once_tool_eventual_success(): + """Test that a tool that fails once eventually succeeds (tests instance respawn).""" + global server + mcp_json = _mcp_config_json({ + "echo": { + "command": sys.executable, + "args": [MCP_ECHO_SERVER], + } + }) + server = _start_server_with_mcp(mcp_json) + + try: + # First call should succeed (warmup already spawned and shut down the instance, + # but the first actual tool call will spawn a fresh instance) + res = server.make_request("POST", "/tools", data={ + "tool": "echo_fail_once", + "params": {} + }) + # It might fail on first call if the warmup instance was shut down + # and a new instance is spawned. The fail_once state is per-process, + # so a fresh process will fail once then succeed. + # Actually, warmup spawns, lists, then shuts down. So the first tool call + # spawns a new process which will fail once. + assert res.status_code in (200, 500), res.body + finally: + server.stop() + + +def test_mcp_tools_via_json_config_file(): + """Test that --mcp-servers-config (file) works as well as --mcp-servers-json.""" + global server + config = { + "mcpServers": { + "echo": { + "command": sys.executable, + "args": [MCP_ECHO_SERVER], + } + } + } + + with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as f: + json.dump(config, f) + config_path = f.name + + try: + server = ServerPreset.router() + server.server_tools = "all" + server.no_ui = True + server.server_port = 8085 + server.mcp_servers_config = config_path + server.start() + + res = server.make_request("GET", "/tools") + assert res.status_code == 200, res.body + + tools = res.body + + def get_tool_name(t): + return t.get("tool", "") or t.get("definition", {}).get("function", {}).get("name", "") + + mcp_tools = [t for t in tools if get_tool_name(t).startswith("echo_")] + assert len(mcp_tools) >= 2, f"Expected echo_ tools, got {mcp_tools}" + finally: + os.unlink(config_path) + server.stop() + + +def test_mcp_tools_slot_independent(): + """MCP tools should work without any slot concept; /tools is slot-independent.""" + global server + mcp_json = _mcp_config_json({ + "echo": { + "command": sys.executable, + "args": [MCP_ECHO_SERVER], + } + }) + server = _start_server_with_mcp(mcp_json) + + try: + # Call /tools without any slot binding - should succeed + res = server.make_request("POST", "/tools", data={ + "tool": "echo_echo", + "params": {"message": "hello"} + }) + assert res.status_code == 200, res.body + body = res.body + assert "error" not in body, body + finally: + server.stop() + + +def test_mcp_concurrent_tool_calls(): + """Concurrent POST /tools to same MCP server should all succeed.""" + global server + mcp_json = _mcp_config_json({ + "echo": { + "command": sys.executable, + "args": [MCP_ECHO_SERVER], + } + }) + server = _start_server_with_mcp(mcp_json) + + try: + def call_tool(): + return server.make_request("POST", "/tools", data={ + "tool": "echo_echo", + "params": {"message": "hi"} + }) + + with ThreadPoolExecutor(max_workers=10) as executor: + futures = [executor.submit(call_tool) for _ in range(10)] + results = [f.result() for f in futures] + + for res in results: + assert res.status_code == 200, res.body + assert "error" not in res.body, res.body + finally: + server.stop() + + +def test_mcp_tool_timeout(): + """Tool call should timeout if MCP server is too slow.""" + global server + MCP_SLOW_SERVER = os.path.join(FIXTURES_DIR, "mcp_slow_server.py") + mcp_json = _mcp_config_json({ + "slow": { + "command": sys.executable, + "args": [MCP_SLOW_SERVER, "--delay", "5"], + "timeout_ms": 500 + } + }) + server = _start_server_with_mcp(mcp_json) + + try: + res = server.make_request("POST", "/tools", data={ + "tool": "slow_sleep", + "params": {"seconds": 5} + }) + assert res.status_code == 200, res.body + body = res.body + assert "error" in body, body + finally: + server.stop() + + +def test_mcp_warmup_partial_failure(): + """Good server's tools should appear even if bad server fails warmup.""" + global server + mcp_json = _mcp_config_json({ + "good": { + "command": sys.executable, + "args": [MCP_ECHO_SERVER], + }, + "bad": { + "command": "nonexistent", + "args": [] + } + }) + server = _start_server_with_mcp(mcp_json) + + try: + res = server.make_request("GET", "/tools") + assert res.status_code == 200, res.body + tools = res.body + + def get_tool_name(t): + return t.get("tool", "") or t.get("definition", {}).get("function", {}).get("name", "") + + # good server tools should be present + assert any("good_" in get_tool_name(t) for t in tools), f"Expected good: tools in {tools}" + finally: + server.stop() + + +def test_mcp_notification_during_request(): + """Notification during request should not be returned as response.""" + global server + MCP_MALFORMED_SERVER = os.path.join(FIXTURES_DIR, "mcp_malformed_server.py") + mcp_json = _mcp_config_json({ + "notifying": { + "command": sys.executable, + "args": [MCP_MALFORMED_SERVER], + } + }) + server = _start_server_with_mcp(mcp_json) + + try: + res = server.make_request("POST", "/tools", data={ + "tool": "notifying_echo", + "params": {"message": "hi"} + }) + assert res.status_code == 200, res.body + body = res.body + assert "error" not in body, body + finally: + server.stop() + + +def test_mcp_instance_respawn_after_crash(): + """Tool call after process crash should respawn and succeed.""" + global server + MCP_CRASH_SERVER = os.path.join(FIXTURES_DIR, "mcp_crash_server.py") + mcp_json = _mcp_config_json({ + "crash": { + "command": sys.executable, + "args": [MCP_CRASH_SERVER], + } + }) + server = _start_server_with_mcp(mcp_json) + + try: + # First call succeeds + res1 = server.make_request("POST", "/tools", data={ + "tool": "crash_echo", + "params": {"message": "hi"} + }) + assert res1.status_code == 200, res1.body + assert "error" not in res1.body, res1.body + + # Second call should also succeed (respawned instance) + res2 = server.make_request("POST", "/tools", data={ + "tool": "crash_echo", + "params": {"message": "hi2"} + }) + assert res2.status_code == 200, res2.body + assert "error" not in res2.body, res2.body + finally: + server.stop() + + + + +def test_mcp_fail_once_eventual_success_verified(): + """Verify that fail_once tool eventually succeeds after respawn.""" + global server + mcp_json = _mcp_config_json({ + "echo": { + "command": sys.executable, + "args": [MCP_ECHO_SERVER], + } + }) + server = _start_server_with_mcp(mcp_json) + + try: + # First call may fail (fresh process) + res1 = server.make_request("POST", "/tools", data={ + "tool": "echo_fail_once", + "params": {} + }) + # Second call should succeed + res2 = server.make_request("POST", "/tools", data={ + "tool": "echo_fail_once", + "params": {} + }) + assert res2.status_code == 200, res2.body + assert "error" not in res2.body, res2.body + finally: + server.stop() + + +def test_mcp_config_file_errors(): + """Invalid JSON config and missing file should cause server to fail to start.""" + # Invalid JSON - server should fail to start + server = ServerPreset.router() + server.server_tools = "all" + server.no_ui = True + server.server_port = 8085 + server.mcp_servers_json = "not valid json" + try: + server.start() + assert False, "Server should not have started with invalid MCP JSON config" + except RuntimeError: + pass # Expected: server process dies due to bad config + + # Missing file - server should fail to start + server = ServerPreset.router() + server.server_tools = "all" + server.no_ui = True + server.server_port = 8085 + server.mcp_servers_config = "/nonexistent/path.json" + try: + server.start() + assert False, "Server should not have started with missing config file" + except RuntimeError: + pass # Expected: server process dies due to missing config + + +def test_mcp_empty_tool_list(): + """MCP server reporting zero tools should result in empty tool list.""" + global server + # Create a minimal server that returns empty tools list + empty_server = os.path.join(FIXTURES_DIR, "_empty_mcp_server.py") + with open(empty_server, "w") as f: + f.write('''#!/usr/bin/env python3 +import json, sys, os +def main(): + sys.stdout = os.fdopen(sys.stdout.fileno(), "w", buffering=1) + for line in sys.stdin: + line = line.strip() + if not line: continue + try: request = json.loads(line) + except: continue + method = request.get("method") + req_id = request.get("id") + if method == "initialize": + resp = {"jsonrpc": "2.0", "id": req_id, "result": {"protocolVersion": "2024-11-05", "capabilities": {"tools": {}}, "serverInfo": {"name": "empty", "version": "1.0"}}} + elif method == "tools/list": + resp = {"jsonrpc": "2.0", "id": req_id, "result": {"tools": []}} + else: + resp = {"jsonrpc": "2.0", "id": req_id, "error": {"code": -32601, "message": "Method not found"}} + sys.stdout.write(json.dumps(resp) + "\\n") + sys.stdout.flush() +if __name__ == "__main__": + main() +''') + try: + mcp_json = _mcp_config_json({ + "empty": { + "command": sys.executable, + "args": [empty_server], + } + }) + server = _start_server_with_mcp(mcp_json) + res = server.make_request("GET", "/tools") + assert res.status_code == 200, res.body + tools = res.body + def get_tool_name(t): + return t.get("tool", "") or t.get("definition", {}).get("function", {}).get("name", "") + mcp_tools = [t for t in tools if get_tool_name(t).startswith("empty:")] + assert len(mcp_tools) == 0, f"Expected no empty: tools, got {mcp_tools}" + finally: + os.unlink(empty_server) + server.stop() + + +def test_mcp_rapid_succession_calls(): + """Many rapid calls should increment next_id correctly and correlate responses.""" + global server + mcp_json = _mcp_config_json({ + "echo": { + "command": sys.executable, + "args": [MCP_ECHO_SERVER], + } + }) + server = _start_server_with_mcp(mcp_json) + + try: + for i in range(20): + res = server.make_request("POST", "/tools", data={ + "tool": "echo_echo", + "params": {"message": f"msg{i}"} + }) + assert res.status_code == 200, res.body + assert "error" not in res.body, res.body + finally: + server.stop() + + +def test_mcp_notification_burst(): + """Notification + response in a single write() with no flush should not strand the response.""" + global server + MCP_BURST_SERVER = os.path.join(FIXTURES_DIR, "mcp_burst_server.py") + mcp_json = _mcp_config_json({ + "burst": { + "command": sys.executable, + "args": [MCP_BURST_SERVER], + } + }) + server = _start_server_with_mcp(mcp_json) + + try: + res = server.make_request("POST", "/tools", data={ + "tool": "burst_echo", + "params": {"message": "burst test"} + }) + assert res.status_code == 200, res.body + body = res.body + assert "error" not in body, body + finally: + server.stop() + + +def test_mcp_tool_definition_shape_via_chat_completions(): + """MCP tool definitions returned by GET /tools should have the correct shape for chat/completions.""" + global server + mcp_json = _mcp_config_json({ + "echo": { + "command": sys.executable, + "args": [MCP_ECHO_SERVER], + } + }) + server = _start_server_with_mcp(mcp_json) + + try: + # Get MCP tool definitions + res = server.make_request("GET", "/tools") + assert res.status_code == 200, res.body + tools = res.body + + def get_tool_name(t): + return t.get("tool", "") or t.get("definition", {}).get("function", {}).get("name", "") + + echo_tools = [t for t in tools if get_tool_name(t).startswith("echo_")] + assert len(echo_tools) >= 2, f"Expected echo_ tools, got {echo_tools}" + + echo_tool = next(t for t in echo_tools if get_tool_name(t) == "echo_echo") + definition = echo_tool.get("definition", echo_tool) + + # Verify the definition has the standard function-calling shape + assert definition.get("type") == "function", f"Expected type=function, got {definition.get('type')}" + func = definition.get("function", {}) + assert "name" in func, "Missing function.name" + assert "description" in func, "Missing function.description" + assert "parameters" in func, f"Missing function.parameters, got keys: {list(func.keys())}" + params = func["parameters"] + assert params.get("type") == "object", f"Expected parameters.type=object, got {params.get('type')}" + assert "properties" in params, "Missing parameters.properties" + finally: + server.stop() + + +def test_mcp_slow_tool_call_slot_release(): + """A slow tool call should not stall server shutdown for the full I/O timeout.""" + global server + MCP_SLOW_SERVER = os.path.join(FIXTURES_DIR, "mcp_slow_server.py") + mcp_json = _mcp_config_json({ + "slow": { + "command": sys.executable, + "args": [MCP_SLOW_SERVER, "--delay", "10"], + "timeout_ms": 30000 + } + }) + server = _start_server_with_mcp(mcp_json) + + try: + # Start a slow tool call in a background thread + def slow_call(): + return server.make_request("POST", "/tools", data={ + "tool": "slow_sleep", + "params": {"seconds": 10} + }) + + with ThreadPoolExecutor(max_workers=1) as executor: + future = executor.submit(slow_call) + + # Wait a moment for the call to start + time.sleep(2) + + # Stop the server while the tool call is in progress. + # With global MCP instances, close_all() is called explicitly at shutdown + # (not from slot release), so shutdown should complete promptly. + start_time = time.time() + server.stop() + elapsed = time.time() - start_time + + # The server should stop quickly, not wait for the full 30s I/O timeout. + # With the terminating flag, send_rpc() bails out within one select() + # slice (~50ms). This threshold MUST stay below the 5s force-kill + # fallback in ServerProcess.stop(): without the flag, shutdown stalls + # on the instance mutex and only completes when stop() sends SIGKILL + # at ~5s -- which any threshold above 5 would still accept. + assert elapsed < 3, f"Server stop took {elapsed:.1f}s, expected < 3s" + + # Wait for the future to complete (it will get an error response or timeout) + try: + res = future.result(timeout=5) + # If we got a response, it should be an error since the server stopped + if hasattr(res, 'status_code'): + assert res.status_code in (200, 500, 502, 503, 504), f"Unexpected status: {res.status_code}" + except Exception: + # Thread may have raised due to connection error - that's acceptable + pass + finally: + server.stop() + + +def test_mcp_grandchild_holding_pipes_does_not_deadlock(): + """An MCP server that leaves a grandchild inheriting its stdout/stderr must not deadlock + teardown. + + subprocess_terminate() only SIGKILLs the direct MCP child, so the inherited pipe write ends + stay open and a blocking read on them would never see EOF. That hung both warmup shutdown + (the server would never reach "ready") and process shutdown. The polled, running-aware reader + must exit regardless, so the server both starts and stops promptly here. + """ + global server + MCP_GRANDCHILD_SERVER = os.path.join(FIXTURES_DIR, "mcp_grandchild_server.py") + mcp_json = _mcp_config_json({ + "gc": { + "command": sys.executable, + "args": [MCP_GRANDCHILD_SERVER], + } + }) + + # If warmup teardown deadlocked, the server would never become ready and start() would time out. + server = _start_server_with_mcp(mcp_json) + + try: + # invoking the tool spawns a live transport whose reader thread holds the inherited pipe + res = server.make_request("POST", "/tools", data={ + "tool": "gc_echo", + "params": {"message": "hello"} + }) + assert res.status_code == 200, res.body + assert "error" not in res.body, res.body + + # shutdown must be prompt: a deadlocked reader-join would stall until the 5s SIGKILL + # fallback in ServerProcess.stop(), so the threshold has to stay below that + start = time.time() + server.stop() + elapsed = time.time() - start + assert elapsed < 3, f"server shutdown took {elapsed:.1f}s (expected < 3s) — teardown likely deadlocked" + finally: + server.stop() diff --git a/tools/server/tests/utils.py b/tools/server/tests/utils.py index 5d5c873ac4cc..ae56bc70a15a 100644 --- a/tools/server/tests/utils.py +++ b/tools/server/tests/utils.py @@ -115,6 +115,8 @@ class ServerProcess: backend_sampling: bool = False gcp_compat: bool = False server_tools: str | None = None + mcp_servers_config: str | None = None + mcp_servers_json: str | None = None cors_origins: str | None = None # session variables @@ -265,6 +267,10 @@ def start(self, timeout_seconds: int = DEFAULT_HTTP_TIMEOUT) -> None: server_args.append("--ui-mcp-proxy") if self.server_tools: server_args.extend(["--tools", self.server_tools]) + if self.mcp_servers_config: + server_args.extend(["--mcp-servers-config", self.mcp_servers_config]) + if self.mcp_servers_json: + server_args.extend(["--mcp-servers-json", self.mcp_servers_json]) if self.backend_sampling: server_args.append("--backend_sampling") if self.gcp_compat: From 8bb909374d04d40621340aee5ba2245860027fdc Mon Sep 17 00:00:00 2001 From: Nicky Mouha Date: Sat, 25 Jul 2026 19:10:32 -0400 Subject: [PATCH 003/190] common : use-after-free when loading LoRA adapter fails (#25611) --- common/common.cpp | 1 - 1 file changed, 1 deletion(-) diff --git a/common/common.cpp b/common/common.cpp index a68766cbbbc8..82dd780fd8b3 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -1249,7 +1249,6 @@ common_init_result::common_init_result(common_params & params, bool model_only) lora.reset(llama_adapter_lora_init(model, la.path.c_str())); if (lora == nullptr) { COM_ERR("failed to load lora adapter '%s'\n", la.path.c_str()); - pimpl->model.reset(model); return; } From 7cdd557f76800b5a84ddee2bff6f20178a3e31fe Mon Sep 17 00:00:00 2001 From: Reese Levine Date: Sat, 25 Jul 2026 17:37:18 -0700 Subject: [PATCH 004/190] ggml-webgpu: Fix WASM compilation with OpenMP (#25943) * Fix emscripten compilation with openmp * Separate wasm job to its own workflow * Add flags necessary for newer emsdk * Just disable openmp * Update triggers --- .github/workflows/build-wasm.yml | 90 ++++++++++++++++++++++++++++++ .github/workflows/build-webgpu.yml | 47 +--------------- 2 files changed, 93 insertions(+), 44 deletions(-) create mode 100644 .github/workflows/build-wasm.yml diff --git a/.github/workflows/build-wasm.yml b/.github/workflows/build-wasm.yml new file mode 100644 index 000000000000..aa7ae887dcd5 --- /dev/null +++ b/.github/workflows/build-wasm.yml @@ -0,0 +1,90 @@ +name: CI (wasm) + +on: + workflow_dispatch: # allows manual triggering + push: + branches: + - master + paths: [ + '.github/workflows/build-wasm.yml', + '**/CMakeLists.txt', + '**/.cmake', + '**/*.h', + '**/*.hpp', + '**/*.c', + '**/*.cpp', + '**/*.wgsl', + '**/*.tmpl', + 'ggml/src/ggml-webgpu/wgsl-shaders/embed_wgsl.py' + ] + + pull_request: + types: [opened, synchronize, reopened] + paths: [ + '.github/workflows/build-wasm.yml', + '**/CMakeLists.txt', + '**/.cmake', + '**/*.h', + '**/*.hpp', + '**/*.c', + '**/*.cpp', + '**/*.wgsl', + '**/*.tmpl', + 'ggml/src/ggml-webgpu/wgsl-shaders/embed_wgsl.py' + ] + +concurrency: + group: ${{ github.workflow }}-${{ github.head_ref && github.ref || github.run_id }} + cancel-in-progress: true + +env: + GGML_NLOOP: 3 + GGML_N_THREADS: 1 + LLAMA_ARG_LOG_COLORS: 1 + LLAMA_ARG_LOG_PREFIX: 1 + LLAMA_ARG_LOG_TIMESTAMPS: 1 + +jobs: + ubuntu-webgpu: + runs-on: ubuntu-24.04-arm + + steps: + - name: Clone + id: checkout + uses: actions/checkout@v6 + + - name: ccache + uses: ggml-org/ccache-action@v1.2.21 + with: + key: webgpu-ubuntu-24.04-arm-wasm + evict-old-files: 1d + save: ${{ github.event_name == 'push' && github.ref == 'refs/heads/master' }} + + - name: Install Emscripten + run: | + git clone https://github.com/emscripten-core/emsdk.git + cd emsdk + ./emsdk install latest + ./emsdk activate latest + + - name: Fetch emdawnwebgpu + run: | + DAWN_TAG="v20260317.182325" + EMDAWN_PKG="emdawnwebgpu_pkg-${DAWN_TAG}.zip" + echo "Downloading ${EMDAWN_PKG}" + curl -L -o emdawn.zip \ + "https://github.com/google/dawn/releases/download/${DAWN_TAG}/${EMDAWN_PKG}" + unzip emdawn.zip + + - name: Build WASM WebGPU + run: | + source emsdk/emsdk_env.sh + emcmake cmake -B build-wasm \ + -G "Ninja" \ + -DCMAKE_BUILD_TYPE=Release \ + -DGGML_WEBGPU=ON \ + -DGGML_OPENMP=OFF \ + -DLLAMA_OPENSSL=OFF \ + -DEMDAWNWEBGPU_DIR=emdawnwebgpu_pkg + + time cmake --build build-wasm --config Release --target test-backend-ops -j $(nproc) diff --git a/.github/workflows/build-webgpu.yml b/.github/workflows/build-webgpu.yml index 0f5ade7af651..ed73c185aa53 100644 --- a/.github/workflows/build-webgpu.yml +++ b/.github/workflows/build-webgpu.yml @@ -13,7 +13,9 @@ on: '**/*.hpp', '**/*.c', '**/*.cpp', - '**/*.wgsl' + '**/*.wgsl', + '**/*.tmpl', + 'ggml/src/ggml-webgpu/wgsl-shaders/embed_wgsl.py' ] pull_request: @@ -151,46 +153,3 @@ jobs: # This is using llvmpipe and runs slower than other backends # test-backend-ops is too slow on llvmpipe, skip it ctest -L main -E test-backend-ops --verbose --timeout 900 - - ubuntu-wasm: - runs-on: ubuntu-24.04-arm - - steps: - - name: Clone - id: checkout - uses: actions/checkout@v6 - - - name: ccache - uses: ggml-org/ccache-action@v1.2.21 - with: - key: webgpu-ubuntu-24.04-arm-wasm - evict-old-files: 1d - save: ${{ github.event_name == 'push' && github.ref == 'refs/heads/master' }} - - - name: Install Emscripten - run: | - git clone https://github.com/emscripten-core/emsdk.git - cd emsdk - ./emsdk install latest - ./emsdk activate latest - - - name: Fetch emdawnwebgpu - run: | - DAWN_TAG="v20260317.182325" - EMDAWN_PKG="emdawnwebgpu_pkg-${DAWN_TAG}.zip" - echo "Downloading ${EMDAWN_PKG}" - curl -L -o emdawn.zip \ - "https://github.com/google/dawn/releases/download/${DAWN_TAG}/${EMDAWN_PKG}" - unzip emdawn.zip - - - name: Build WASM WebGPU - run: | - source emsdk/emsdk_env.sh - emcmake cmake -B build-wasm \ - -G "Ninja" \ - -DCMAKE_BUILD_TYPE=Release \ - -DGGML_WEBGPU=ON \ - -DLLAMA_OPENSSL=OFF \ - -DEMDAWNWEBGPU_DIR=emdawnwebgpu_pkg - - time cmake --build build-wasm --config Release --target test-backend-ops -j $(nproc) From ff067f76dd8e9e05f0528056f1274adf01a54d70 Mon Sep 17 00:00:00 2001 From: Pascal Date: Sun, 26 Jul 2026 06:51:10 +0200 Subject: [PATCH 005/190] ui: fix context gauge card regressions and land at the conversation end (#26099) The context gauge card starts monitoring like the dial does, because its own processing state instance only follows the live stream while its monitoring flag is set. It also gets back the text-sm and ring classes the removed hover card wrapper used to inject, which restores its layout. Routing to a conversation now lands at the bottom instantly and keeps the pin one frame at a time until the page height settles, since content-visibility size realizations and syntax highlight passes grow the page without DOM mutations. --- .../ContextGaugePopup.svelte | 9 ++++- .../app/chat/ChatScreen/ChatScreen.svelte | 40 ++++++++++++++++++- tools/ui/src/lib/constants/auto-scroll.ts | 6 +++ 3 files changed, 52 insertions(+), 3 deletions(-) diff --git a/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormContextGauge/ContextGaugePopup.svelte b/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormContextGauge/ContextGaugePopup.svelte index 81acdbaf7296..af9ad010e3a3 100644 --- a/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormContextGauge/ContextGaugePopup.svelte +++ b/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormContextGauge/ContextGaugePopup.svelte @@ -13,6 +13,13 @@ const gauge = useContextGauge(); + // The gauge hook wraps a processing state instance that only follows the + // live stream while its own monitoring flag is set, so the card instance + // starts monitoring like the dial does. + $effect(() => { + gauge.startMonitoring(); + }); + let cardEl = $state(null); // Any press outside the card and outside the dial closes the card. @@ -44,7 +51,7 @@
{ + if (autoScroll.userScrolledUp) return; + if (activeConversation()?.id !== id) return; + autoScroll.scrollToBottom(); + const height = container.scrollHeight; + stableFrames = height === lastHeight ? stableFrames + 1 : 0; + lastHeight = height; + if (stableFrames >= LANDING_STABLE_FRAMES) return; + if (performance.now() - started > LANDING_SETTLE_MAX_MS) return; + requestAnimationFrame(settle); + }; + requestAnimationFrame(settle); + } + function handleSendLikeScroll() { if (!isMobile.current) { autoScroll.enable(); @@ -246,6 +281,7 @@ {#if !isEmpty} { handleSendLikeScroll(); }} diff --git a/tools/ui/src/lib/constants/auto-scroll.ts b/tools/ui/src/lib/constants/auto-scroll.ts index c629228eb485..67c5f930108d 100644 --- a/tools/ui/src/lib/constants/auto-scroll.ts +++ b/tools/ui/src/lib/constants/auto-scroll.ts @@ -1,4 +1,10 @@ export const AUTO_SCROLL_INTERVAL = 100; +// Conversation landing: the page keeps growing after the first bottom pin +// without DOM mutations (content-visibility size realizations, syntax +// highlight passes), so the pin repeats every frame until the height holds +// for this many consecutive frames, bounded by the time cap below. +export const LANDING_STABLE_FRAMES = 10; +export const LANDING_SETTLE_MAX_MS = 1000; // Chat main view: tight threshold because scroll-here events come from // discrete assistant-message appends. export const AUTO_SCROLL_AT_BOTTOM_THRESHOLD = 10; From 42fc243060709331ff9b158a9ed2cbe37219ae83 Mon Sep 17 00:00:00 2001 From: yzyyzyhhh <96101183+happyyzy@users.noreply.github.com> Date: Sun, 26 Jul 2026 23:01:08 +0800 Subject: [PATCH 006/190] opencl: fix fused RMS norm mul view offset (#26085) --- ggml/src/ggml-opencl/ggml-opencl.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ggml/src/ggml-opencl/ggml-opencl.cpp b/ggml/src/ggml-opencl/ggml-opencl.cpp index b6079d08086c..a05d18ee30af 100644 --- a/ggml/src/ggml-opencl/ggml-opencl.cpp +++ b/ggml/src/ggml-opencl/ggml-opencl.cpp @@ -12772,7 +12772,7 @@ static void ggml_opencl_op_rms_norm_fused(ggml_backend_t backend, ggml_tensor * ggml_tensor_extra_cl * extrad = (ggml_tensor_extra_cl *)dst->extra; cl_ulong offset0 = extra0->offset + src0->view_offs; - cl_ulong offset1 = extra1->offset + src0->view_offs; + cl_ulong offset1 = extra1->offset + src1->view_offs; cl_ulong offsetd = extrad->offset + dst->view_offs; ggml_backend_opencl_context *backend_ctx = (ggml_backend_opencl_context *)backend->context; From b1d4c6552489d27eb732c2804f189d9cc6ba99bd Mon Sep 17 00:00:00 2001 From: timkhronos Date: Sun, 26 Jul 2026 19:43:45 +0200 Subject: [PATCH 007/190] model: Add MiniMax-M3 (MSA: MiniMax Sparse Attention) support (#24908) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Add preliminary MiniMax-M3 support Text-only port that re-uses existing components: MiniMax-M2 style GQA with per-head QK-norm and partial rotary, DeepSeek-V3 style leading-dense and routed/shared experts, and swigluoai activation. Sparse attention is not yet supported (dense fallback); vision tower and MTP heads are dropped. * MiniMax-M3 vision tower (mmproj + clip graph) * Delete m3_vision_ref.py * Update clip.cpp * MSA * Update constants.py * Update minimax.py * Cache creation. Working withotu flash attention * Added flash attention for sparse layers * Decomposed slow cpu OP into GPU + CPU ops. Massive speedup over long ctx * Rewrote indexer op to be cuda native. Modified flash attention to match per group block picking * Implement sparse attention calc out of stock ops. * Fix a cache allocation and cont issue * Fixed -fa auto crash, flagged debug spots * Delete vocab.json * Delete model.safetensors.index.json * Delete generation_config.json * Delete Minimax directory * Handled multi stream case to fall back on Dense Attention * Development scaffolding cleanup. No functional change to the decode or 4-way paths. Full debug harness remains at <8136a9c68ed7a5eb009aa67bba3fda8062f4648f> for reproducing the selection-parity validation. * Remove redundant comment from minimax-m3.cpp * Changed 3 Gelu Ops for vision into Gelu_erf ops * Assert that n_kv is multiple of 128 * Rename MSA index tensors to indexer convention Note: All GGUFs generated before this change will need to be regenerated. * Fix incorrect Assert * Review driven changes (#3) * Remove comment from conversion minimax.py Co-authored-by: Sigbjørn Skjæret <1629204+CISC@users.noreply.github.com> * Remove whitespaces from constants.py Co-authored-by: Sigbjørn Skjæret <1629204+CISC@users.noreply.github.com> * Tighten comment in minimax.py Co-authored-by: Sigbjørn Skjæret <1629204+CISC@users.noreply.github.com> * inherit MiniMax-M3 from MiniMax-M2 * drop dead text_config fallbacks * Add indexer writer methods * Reuse LLM_FFN_SWIGLU_OAI_MOE * Remove duplicate indexer setters, add only block_size/local_blocks, follow value naming convention * Fix conversion error /gguf_writer.py Co-authored-by: Sigbjørn Skjæret <1629204+CISC@users.noreply.github.com> * Update gguf-py/gguf/gguf_writer.py Co-authored-by: Sigbjørn Skjæret <1629204+CISC@users.noreply.github.com> * Update gguf-py/gguf/tensor_mapping.py Co-authored-by: Sigbjørn Skjæret <1629204+CISC@users.noreply.github.com> * Update conversion/minimax.py Co-authored-by: Sigbjørn Skjæret <1629204+CISC@users.noreply.github.com> * Update conversion/minimax.py Co-authored-by: Sigbjørn Skjæret <1629204+CISC@users.noreply.github.com> * Remove whitespace in src/llama-kv-cache.cpp Co-authored-by: Sigbjørn Skjæret <1629204+CISC@users.noreply.github.com> * Remove Whitespace in Update src/llama-model.h Co-authored-by: Sigbjørn Skjæret <1629204+CISC@users.noreply.github.com> * Remove whitespace in src/llama-hparams.h Co-authored-by: Sigbjørn Skjæret <1629204+CISC@users.noreply.github.com> * remove multimodal code upon maintainer request. Will be made as a separate PR * Whitespace clean in tensor_mapping.py * Log cache size on launch, block ctx shift, support prompt caching Log indexer cache size on launch Disallow ctx shift Support prompt caching * Update minimax-m3.cpp * Optimize implementation, add multi stream support. Fully rewrote minimax-m3.cpp for speed and buffer size gains: Unified the 4-way + decode, 1 FA call per layer instead of 4, with the groups mapped onto ne[3] Custom CPU op now emits block-level mask, expanded on GPU, which causes CPU to GPU transfer to shrinks at prefill Decode: ~25 nodes/layer vs ~50, no per-group concats/conts Unified selection semantics, so both regimes rank bs + local bias (position-anchored local force), which means prefill/decode can no longer disagree on selection can_reuse on the MSA bias input. Graph reuse at decode restored (was rebuilding the full graph every token) In-place mask adds, shrinking compute buffer ~6.8 to ~4.2 GiB at ub2048/62k Multi-stream: MSA now runs with -np N when kv_unified=false. Decode stays batched across streams (still 1 FA call), prefill loops per stream. dense fallback only for --kv-unified + multi-seq Measured effect on expert offload bound setup: decode 6.2(4WAY)–7.15(MSA_decode) -> 7.7~7.8 t/s, flat from 5k to 60k+. prefill around 10% faster. buffer about 20% smaller, multi-user support. * set default cache type to F32 * Fix potential DSA double indexer cache allocation bug, only allocate in-cache k_idx for archs that opt in * remove F16 downcasts in MSA attention, force F32 indexer score accum * Add Minimax eos to llama vocab * Guard edge case where idx cache can become stale after a tail trim * Update llama-kv-cache.h * Update llama-kv-cache.cpp * Update llama-kv-cache.cpp * Update llama-kv-cache.h * Update llama-kv-cache.cpp * Review driven changes * style fix * indexer hparams are required * fix tests * fix lint --------- Co-authored-by: Daniel Han Co-authored-by: Sigbjørn Skjæret <1629204+CISC@users.noreply.github.com> Co-authored-by: Xuan Son Nguyen --- conversion/__init__.py | 2 + conversion/base.py | 4 +- conversion/minimax.py | 37 ++- gguf-py/gguf/constants.py | 38 +++ gguf-py/gguf/gguf_writer.py | 6 + gguf-py/gguf/tensor_mapping.py | 15 +- src/llama-arch.cpp | 10 + src/llama-arch.h | 6 + src/llama-context.cpp | 3 +- src/llama-graph.cpp | 13 +- src/llama-hparams.cpp | 10 + src/llama-hparams.h | 8 + src/llama-kv-cache.cpp | 292 ++++++++++++++++- src/llama-kv-cache.h | 10 + src/llama-model-saver.cpp | 2 + src/llama-model.cpp | 4 + src/llama-model.h | 7 + src/llama-vocab.cpp | 1 + src/models/minimax-m3.cpp | 562 +++++++++++++++++++++++++++++++++ src/models/models.h | 23 ++ tests/test-llama-archs.cpp | 14 +- 21 files changed, 1044 insertions(+), 23 deletions(-) create mode 100644 src/models/minimax-m3.cpp diff --git a/conversion/__init__.py b/conversion/__init__.py index 7936f1159cb8..0b08e6e57008 100644 --- a/conversion/__init__.py +++ b/conversion/__init__.py @@ -158,6 +158,8 @@ "MiniCPMForCausalLM": "minicpm", "MiniCPMV4_6ForConditionalGeneration": "minicpm", "MiniMaxM2ForCausalLM": "minimax", + "MiniMaxM3SparseForCausalLM": "minimax", + "MiniMaxM3SparseForConditionalGeneration": "minimax", "Ministral3ForCausalLM": "mistral3", "Mistral3ForConditionalGeneration": "mistral3", "MistralForCausalLM": "llama", diff --git a/conversion/base.py b/conversion/base.py index 051b8b4e59fd..a7cd3fd904aa 100644 --- a/conversion/base.py +++ b/conversion/base.py @@ -1156,7 +1156,7 @@ def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Ca or "projector." in name or "pre_mm_projector_norm" in name \ or "image_newline" in name or "view_seperator" in name \ or "patch_embed" in name or "patch_embedding" in name \ - or "patch_merger." in name or "model.connector." in name: + or "patch_merger." in name or "patch_merge_mlp." in name or "model.connector." in name: return None return super().filter_tensors(item) @@ -1203,7 +1203,7 @@ def set_gguf_parameters(self): self.gguf_writer.add_embedding_length(n_embd) logger.info(f"gguf: embedding length = {n_embd}") - if (n_ff := self.find_hparam(["prefix_dense_intermediate_size", "intermediate_size", "n_inner", "hidden_dim"], optional=True)) is not None: + if (n_ff := self.find_hparam(["prefix_dense_intermediate_size", "dense_intermediate_size", "intermediate_size", "n_inner", "hidden_dim"], optional=True)) is not None: self.gguf_writer.add_feed_forward_length(n_ff) logger.info(f"gguf: feed forward length = {n_ff}") diff --git a/conversion/minimax.py b/conversion/minimax.py index 4857775cbfb9..cbbdfe3ae82d 100644 --- a/conversion/minimax.py +++ b/conversion/minimax.py @@ -23,7 +23,7 @@ def set_gguf_parameters(self): def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None): # merge expert weights - if 'experts' in name: + if "block_sparse_moe.experts." in name: n_experts = self.find_hparam(["num_local_experts", "num_experts"]) assert bid is not None @@ -52,3 +52,38 @@ def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None): return yield from super().modify_tensors(data_torch, name, bid) + + +@ModelBase.register("MiniMaxM3SparseForCausalLM", "MiniMaxM3SparseForConditionalGeneration") +class MiniMaxM3Model(MiniMaxM2Model): + model_arch = gguf.MODEL_ARCH.MINIMAXM3 + + def set_gguf_parameters(self): + super().set_gguf_parameters() + + self.gguf_writer.add_expert_shared_count(self.find_hparam(["n_shared_experts"])) + self.gguf_writer.add_expert_weights_scale(self.find_hparam(["routed_scaling_factor"])) + self.gguf_writer.add_expert_weights_norm(True) + + sac = self.find_hparam(["sparse_attention_config"]) + self.gguf_writer.add_indexer_head_count(sac["sparse_num_index_heads"]) + self.gguf_writer.add_indexer_key_length(sac["sparse_index_dim"]) + self.gguf_writer.add_indexer_top_k(sac["sparse_topk_blocks"]) + self.gguf_writer.add_indexer_block_size(sac["sparse_block_size"]) + self.gguf_writer.add_indexer_local_blocks(sac["sparse_local_block"]) + + moe_layer_freq = self.find_hparam(["moe_layer_freq"]) + n_dense = 0 + for v in moe_layer_freq: + if v == 0: + n_dense += 1 + else: + break + self.gguf_writer.add_leading_dense_block_count(n_dense) + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None): + # Gemma-style (1 + w) RMSNorm: bake the +1 in so llama.cpp can use plain RMSNorm + if name.endswith("norm.weight"): + data_torch = data_torch + 1.0 + + yield from super().modify_tensors(data_torch, name, bid) diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index d55253e0eb4b..66d50cca2684 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -200,6 +200,8 @@ class Indexer: HEAD_COUNT = "{arch}.attention.indexer.head_count" KEY_LENGTH = "{arch}.attention.indexer.key_length" TOP_K = "{arch}.attention.indexer.top_k" + BLOCK_SIZE = "{arch}.attention.indexer.block_size" # MSA + LOCAL_BLOCKS = "{arch}.attention.indexer.local_blocks" # MSA TYPES = "{arch}.attention.indexer.types" class HyperConnection: @@ -528,6 +530,7 @@ class MODEL_ARCH(IntEnum): APERTUS = auto() COGVLM = auto() MINIMAXM2 = auto() + MINIMAXM3 = auto() RND1 = auto() PANGU_EMBED = auto() MISTRAL3 = auto() @@ -774,6 +777,9 @@ class MODEL_TENSOR(IntEnum): INDEXER_PROJ = auto() INDEXER_ATTN_K = auto() INDEXER_ATTN_Q_B = auto() + INDEXER_Q_PROJ = auto() + INDEXER_K_PROJ = auto() + INDEXER_Q_NORM = auto() INDEXER_COMPRESSOR_WKV = auto() INDEXER_COMPRESSOR_WGATE = auto() INDEXER_COMPRESSOR_APE = auto() @@ -1110,6 +1116,7 @@ class MODEL_TENSOR(IntEnum): MODEL_ARCH.GROVEMOE: "grovemoe", MODEL_ARCH.APERTUS: "apertus", MODEL_ARCH.MINIMAXM2: "minimax-m2", + MODEL_ARCH.MINIMAXM3: "minimax-m3", MODEL_ARCH.COGVLM: "cogvlm", MODEL_ARCH.RND1: "rnd1", MODEL_ARCH.PANGU_EMBED: "pangu-embedded", @@ -1355,6 +1362,9 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.INDEXER_PROJ: "blk.{bid}.indexer.proj", MODEL_TENSOR.INDEXER_ATTN_K: "blk.{bid}.indexer.attn_k", MODEL_TENSOR.INDEXER_ATTN_Q_B: "blk.{bid}.indexer.attn_q_b", + MODEL_TENSOR.INDEXER_Q_PROJ: "blk.{bid}.indexer.q_proj", + MODEL_TENSOR.INDEXER_K_PROJ: "blk.{bid}.indexer.k_proj", + MODEL_TENSOR.INDEXER_Q_NORM: "blk.{bid}.indexer.q_norm", MODEL_TENSOR.INDEXER_COMPRESSOR_WKV: "blk.{bid}.indexer_compressor_kv", MODEL_TENSOR.INDEXER_COMPRESSOR_WGATE: "blk.{bid}.indexer_compressor_gate", MODEL_TENSOR.INDEXER_COMPRESSOR_APE: "blk.{bid}.indexer_compressor_ape", @@ -4163,6 +4173,34 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.FFN_UP_EXP, MODEL_TENSOR.FFN_EXP_PROBS_B, ], + MODEL_ARCH.MINIMAXM3: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_Q_NORM, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_K_NORM, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE_INP, + MODEL_TENSOR.FFN_EXP_PROBS_B, + MODEL_TENSOR.FFN_GATE_EXP, + MODEL_TENSOR.FFN_DOWN_EXP, + MODEL_TENSOR.FFN_UP_EXP, + MODEL_TENSOR.FFN_GATE_SHEXP, + MODEL_TENSOR.FFN_DOWN_SHEXP, + MODEL_TENSOR.FFN_UP_SHEXP, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.INDEXER_Q_PROJ, + MODEL_TENSOR.INDEXER_K_PROJ, + MODEL_TENSOR.INDEXER_Q_NORM, + MODEL_TENSOR.INDEXER_K_NORM, + ], MODEL_ARCH.COGVLM: [ MODEL_TENSOR.TOKEN_EMBD, MODEL_TENSOR.OUTPUT_NORM, diff --git a/gguf-py/gguf/gguf_writer.py b/gguf-py/gguf/gguf_writer.py index bb21596701d4..ba08f8d65004 100644 --- a/gguf-py/gguf/gguf_writer.py +++ b/gguf-py/gguf/gguf_writer.py @@ -793,6 +793,12 @@ def add_indexer_key_length(self, length: int) -> None: def add_indexer_top_k(self, top_k: int) -> None: self.add_uint32(Keys.Attention.Indexer.TOP_K.format(arch=self.arch), top_k) + def add_indexer_block_size(self, block_size: int) -> None: + self.add_uint32(Keys.Attention.Indexer.BLOCK_SIZE.format(arch=self.arch), block_size) + + def add_indexer_local_blocks(self, local_blocks: int) -> None: + self.add_uint32(Keys.Attention.Indexer.LOCAL_BLOCKS.format(arch=self.arch), local_blocks) + def add_indexer_types(self, value: Sequence[bool]) -> None: key = Keys.Attention.Indexer.TYPES.format(arch=self.arch) self.add_array(key, value) diff --git a/gguf-py/gguf/tensor_mapping.py b/gguf-py/gguf/tensor_mapping.py index b5707f11f5c4..59623accfdb6 100644 --- a/gguf-py/gguf/tensor_mapping.py +++ b/gguf-py/gguf/tensor_mapping.py @@ -1264,7 +1264,8 @@ class TensorNameMap: ), MODEL_TENSOR.INDEXER_K_NORM: ( - "model.layers.{bid}.self_attn.indexer.k_norm", # DSA + "model.layers.{bid}.self_attn.indexer.k_norm", # DSA + "model.layers.{bid}.self_attn.index_k_norm", # MSA ), MODEL_TENSOR.INDEXER_PROJ: ( @@ -1279,6 +1280,18 @@ class TensorNameMap: "model.layers.{bid}.self_attn.indexer.wq_b", # DSA ), + MODEL_TENSOR.INDEXER_Q_PROJ: ( + "model.layers.{bid}.self_attn.index_q_proj", # MSA + ), + + MODEL_TENSOR.INDEXER_K_PROJ: ( + "model.layers.{bid}.self_attn.index_k_proj", # MSA + ), + + MODEL_TENSOR.INDEXER_Q_NORM: ( + "model.layers.{bid}.self_attn.index_q_norm", # MSA + ), + ############################################################################ # TODO: these do not belong to block_mappings_cfg - move them to mappings_cfg MODEL_TENSOR.ENC_OUTPUT_NORM: ( diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index 9aa3dace5ce0..39bf2c79590b 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -127,6 +127,7 @@ static const std::map LLM_ARCH_NAMES = { { LLM_ARCH_GROVEMOE, "grovemoe" }, { LLM_ARCH_APERTUS, "apertus" }, { LLM_ARCH_MINIMAX_M2, "minimax-m2" }, + { LLM_ARCH_MINIMAX_M3, "minimax-m3" }, { LLM_ARCH_COGVLM, "cogvlm" }, { LLM_ARCH_RND1, "rnd1" }, { LLM_ARCH_PANGU_EMBED, "pangu-embedded" }, @@ -253,6 +254,8 @@ static const std::map LLM_KV_NAMES = { { LLM_KV_ATTENTION_INDEXER_HEAD_COUNT, "%s.attention.indexer.head_count" }, { LLM_KV_ATTENTION_INDEXER_KEY_LENGTH, "%s.attention.indexer.key_length" }, { LLM_KV_ATTENTION_INDEXER_TOP_K, "%s.attention.indexer.top_k" }, + { LLM_KV_ATTENTION_INDEXER_BLOCK_SIZE, "%s.attention.indexer.block_size" }, + { LLM_KV_ATTENTION_INDEXER_LOCAL_BLOCKS, "%s.attention.indexer.local_blocks" }, { LLM_KV_ATTENTION_INDEXER_TYPES, "%s.attention.indexer.types" }, { LLM_KV_ATTENTION_OUTPUT_GROUP_COUNT, "%s.attention.output_group_count" }, { LLM_KV_ATTENTION_OUTPUT_LORA_RANK, "%s.attention.output_lora_rank" }, @@ -597,6 +600,9 @@ static const std::map LLM_TENSOR_NAMES = { { LLM_TENSOR_INDEXER_PROJ, "blk.%d.indexer.proj" }, { LLM_TENSOR_INDEXER_ATTN_K, "blk.%d.indexer.attn_k" }, { LLM_TENSOR_INDEXER_ATTN_Q_B, "blk.%d.indexer.attn_q_b" }, + { LLM_TENSOR_INDEXER_Q_PROJ, "blk.%d.indexer.q_proj" }, + { LLM_TENSOR_INDEXER_K_PROJ, "blk.%d.indexer.k_proj" }, + { LLM_TENSOR_INDEXER_Q_NORM, "blk.%d.indexer.q_norm" }, { LLM_TENSOR_INDEXER_COMPRESSOR_WKV, "blk.%d.indexer_compressor_kv" }, { LLM_TENSOR_INDEXER_COMPRESSOR_WGATE, "blk.%d.indexer_compressor_gate" }, { LLM_TENSOR_INDEXER_COMPRESSOR_APE, "blk.%d.indexer_compressor_ape" }, @@ -832,6 +838,9 @@ static const std::map LLM_TENSOR_INFOS = { {LLM_TENSOR_INDEXER_PROJ, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, {LLM_TENSOR_INDEXER_ATTN_K, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, {LLM_TENSOR_INDEXER_ATTN_Q_B, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_INDEXER_Q_PROJ, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_INDEXER_K_PROJ, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_INDEXER_Q_NORM, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, {LLM_TENSOR_INDEXER_COMPRESSOR_WKV, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, {LLM_TENSOR_INDEXER_COMPRESSOR_WGATE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, {LLM_TENSOR_INDEXER_COMPRESSOR_APE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_GET_ROWS}}, @@ -1001,6 +1010,7 @@ bool llm_arch_supports_sm_tensor(const llm_arch & arch) { case LLM_ARCH_LFM2: case LLM_ARCH_LFM2MOE: case LLM_ARCH_MINIMAX_M2: + case LLM_ARCH_MINIMAX_M3: case LLM_ARCH_MISTRAL4: case LLM_ARCH_KIMI_LINEAR: return false; diff --git a/src/llama-arch.h b/src/llama-arch.h index 39c55a66a943..2e3916a0beee 100644 --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -146,6 +146,7 @@ enum llm_arch { LLM_ARCH_TALKIE, LLM_ARCH_MELLUM, LLM_ARCH_EAGLE3, + LLM_ARCH_MINIMAX_M3, LLM_ARCH_DFLASH, LLM_ARCH_UNKNOWN, }; @@ -258,6 +259,8 @@ enum llm_kv { LLM_KV_ATTENTION_INDEXER_HEAD_COUNT, LLM_KV_ATTENTION_INDEXER_KEY_LENGTH, LLM_KV_ATTENTION_INDEXER_TOP_K, + LLM_KV_ATTENTION_INDEXER_BLOCK_SIZE, + LLM_KV_ATTENTION_INDEXER_LOCAL_BLOCKS, LLM_KV_ATTENTION_INDEXER_TYPES, LLM_KV_ATTENTION_OUTPUT_GROUP_COUNT, LLM_KV_ATTENTION_OUTPUT_LORA_RANK, @@ -597,6 +600,9 @@ enum llm_tensor { LLM_TENSOR_INDEXER_PROJ, LLM_TENSOR_INDEXER_ATTN_K, LLM_TENSOR_INDEXER_ATTN_Q_B, + LLM_TENSOR_INDEXER_Q_PROJ, + LLM_TENSOR_INDEXER_K_PROJ, + LLM_TENSOR_INDEXER_Q_NORM, LLM_TENSOR_INDEXER_COMPRESSOR_WKV, LLM_TENSOR_INDEXER_COMPRESSOR_WGATE, LLM_TENSOR_INDEXER_COMPRESSOR_APE, diff --git a/src/llama-context.cpp b/src/llama-context.cpp index eed041eef4e7..c512477c0eab 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -2338,7 +2338,8 @@ uint32_t llama_context::graph_max_nodes(uint32_t n_tokens) const { model.arch == LLM_ARCH_KIMI_LINEAR || model.arch == LLM_ARCH_QWEN35 || model.arch == LLM_ARCH_QWEN35MOE || - model.arch == LLM_ARCH_DEEPSEEK4) { + model.arch == LLM_ARCH_DEEPSEEK4 || + model.arch == LLM_ARCH_MINIMAX_M3) { return std::max(n_tokens * 40, 32u * model.n_tensors()); } uint32_t res = std::max(1024u, 8u*model.n_tensors()); diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index c8ecb0a2854c..6d1c8f4e42a8 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -1709,6 +1709,17 @@ ggml_tensor * llm_graph_context::build_ffn( cur = ggml_swiglu(ctx0, cur); cb(cur, "ffn_swiglu", il); } break; + case LLM_FFN_SWIGLU_OAI_MOE: + if (gate && type_gate == LLM_FFN_PAR) { + // same alpha/limit constants as gpt-oss + const float alpha = 1.702f; + const float limit = 7.0f; + cur = ggml_swiglu_oai(ctx0, cur, tmp, alpha, limit); + cb(cur, "ffn_swiglu_oai", il); + type_gate = LLM_FFN_SEQ; + } else { + GGML_ABORT("LLM_FFN_SWIGLU_OAI_MOE requires a parallel gate"); + } break; case LLM_FFN_GEGLU: { cur = ggml_geglu(ctx0, cur); @@ -2668,7 +2679,7 @@ ggml_tensor * llm_graph_context::build_attn( ggml_build_forward_expand(gf, mctx_cur->cpy_v(ctx0, v_cur, v_idxs, il)); } - const auto & kq_mask = inp->get_kq_mask(); + ggml_tensor * kq_mask = inp->get_kq_mask(); ggml_tensor * q = q_cur; ggml_tensor * k = mctx_cur->get_k(ctx0, il); diff --git a/src/llama-hparams.cpp b/src/llama-hparams.cpp index 846d4c69a626..50af97f358c3 100644 --- a/src/llama-hparams.cpp +++ b/src/llama-hparams.cpp @@ -180,6 +180,16 @@ uint32_t llama_hparams::n_embd_v_gqa_max() const { return val; } +uint32_t llama_hparams::n_embd_k_idx(uint32_t il) const { + if (!indexer_kv || indexer_head_size == 0) { + return 0; // arch without a MSA indexer + } + if (il < n_layer_dense_lead) { + return 0; // leading dense layers carry no indexer + } + return indexer_head_size; // 128 +} + uint32_t llama_hparams::n_embd_r() const { if (wkv_head_size != 0) { // for RWKV models diff --git a/src/llama-hparams.h b/src/llama-hparams.h index 747754fc0d0b..727df6ca21e2 100644 --- a/src/llama-hparams.h +++ b/src/llama-hparams.h @@ -226,6 +226,11 @@ struct llama_hparams { uint32_t indexer_n_head = 0; uint32_t indexer_head_size = 0; uint32_t indexer_top_k = 0; + // MSA + uint32_t indexer_block_size = 0; + uint32_t indexer_local_blocks = 0; + // MSA stores its indexer keys in the main KV cache (k_idx tensors); + bool indexer_kv = false; // Indexer is "full" (1) or "shared" (0) // Shared indexers reuse top-k from previous full layer @@ -350,6 +355,9 @@ struct llama_hparams { uint32_t n_embd_k_gqa_max() const; uint32_t n_embd_v_gqa_max() const; + // dimension of the single-head MSA indexer key stream + uint32_t n_embd_k_idx(uint32_t il = 0) const; + // dimension of the rolling state embeddings // corresponds to Mamba's conv_states size or RWKV's token_shift states size uint32_t n_embd_r() const; diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index e25464c597ac..44cb1668dacf 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -112,7 +112,7 @@ llama_kv_cache::llama_kv_cache( auto it = ctx_map.find(buft); if (it == ctx_map.end()) { ggml_init_params params = { - /*.mem_size =*/ size_t(2u*(1 + n_stream)*n_layer*ggml_tensor_overhead()), + /*.mem_size =*/ size_t(3u*(1 + n_stream)*n_layer*ggml_tensor_overhead()), //Reserve tensor metadata for up to 3 tensors per layer (K, V, and optional K_idx), plus one view per tensor per stream. /*.mem_buffer =*/ NULL, /*.no_alloc =*/ true, }; @@ -242,9 +242,25 @@ llama_kv_cache::llama_kv_cache( v_stream.push_back(has_v ? ggml_view_2d(ctx, v, n_embd_v_gqa, kv_size, v->nb[1], s*v->nb[2]) : nullptr); } + const uint32_t n_embd_k_idx = hparams.n_embd_k_idx(il); + ggml_tensor * k_idx = n_embd_k_idx > 0 + ? ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_embd_k_idx, kv_size, n_stream) + : nullptr; + if (k_idx) { + ggml_format_name(k_idx, "cache_k_idx_l%d", il); + msa_strict_slots = (n_stream == n_seq_max); + } + + std::vector k_idx_stream; + for (uint32_t s = 0; s < n_stream; ++s) { + k_idx_stream.push_back(k_idx + ? ggml_view_2d(ctx, k_idx, n_embd_k_idx, kv_size, k_idx->nb[1], s*k_idx->nb[2]) + : nullptr); + } + map_layer_ids[il] = layers.size(); - layers.push_back({ il, k, v, k_stream, v_stream, }); + layers.push_back({ il, k, v, k_idx, k_stream, v_stream, k_idx_stream }); } if (reuse) { @@ -293,13 +309,24 @@ llama_kv_cache::llama_kv_cache( } { - const size_t memory_size_k = size_k_bytes(); - const size_t memory_size_v = size_v_bytes(); + const size_t memory_size_k = size_k_bytes(); + const size_t memory_size_v = size_v_bytes(); + const size_t memory_size_k_idx = size_k_idx_bytes(); + const size_t memory_size_total = memory_size_k + memory_size_v + memory_size_k_idx; + + constexpr float mib = 1024.0f * 1024.0f; + + const std::string k_log = format(", K (%s): %7.2f MiB", ggml_type_name(type_k), (float) memory_size_k / mib); + const std::string v_log = format(", V (%s): %7.2f MiB", ggml_type_name(type_v), (float) memory_size_v / mib); + + std::string k_idx_log; + if (memory_size_k_idx > 0) { + k_idx_log = format(", K_idx (%s): %7.2f MiB", ggml_type_name(GGML_TYPE_F32), (float) memory_size_k_idx / mib); + } - LLAMA_LOG_INFO("%s: size = %7.2f MiB (%6u cells, %3d layers, %2u/%u seqs), K (%s): %7.2f MiB, V (%s): %7.2f MiB\n", __func__, - (float)(memory_size_k + memory_size_v) / (1024.0f * 1024.0f), kv_size, (int) layers.size(), n_seq_max, n_stream, - ggml_type_name(type_k), (float)memory_size_k / (1024.0f * 1024.0f), - ggml_type_name(type_v), (float)memory_size_v / (1024.0f * 1024.0f)); + LLAMA_LOG_INFO("%s: size = %7.2f MiB (%6u cells, %3d layers, %2u/%u seqs)%s%s%s\n", __func__, + (float) memory_size_total / mib, kv_size, (int) layers.size(), n_seq_max, n_stream, + k_log.c_str(), v_log.c_str(), k_idx_log.c_str()); } // TODO: refactor [TAG_KV_CACHE_SHARE_CELLS] @@ -392,6 +419,39 @@ bool llama_kv_cache::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1) { p1 = std::numeric_limits::max(); } + // empty range - nothing to remove + if (p0 >= p1) { + return true; + } + + // MSA anchors block selection to absolute cache slots (slot == position). Tail trim and full removal preserve this invariant, but removing a prefix + // or middle range would free slots while later cells survive, desynchronizing the indexer cache. Reject such removals before modifying the cache. + if (msa_strict_slots) { + for (llama_seq_id sid = 0; sid < (llama_seq_id) seq_to_stream.size(); ++sid) { + if (seq_id >= 0 && sid != seq_id) { + continue; + } + + const auto & cells = v_cells[seq_to_stream[sid]]; + + const llama_pos pmin = cells.seq_pos_min(sid); + const llama_pos pmax = cells.seq_pos_max(sid); + + if (pmin < 0) { + continue; // empty sequence + } + + const bool overlaps = p0 <= pmax && p1 > pmin; // the range removes something + const bool leaves_tail = p1 <= pmax; // cells beyond the range survive + + if (overlaps && leaves_tail) { + LLAMA_LOG_WARN("%s: MSA: partial (non-suffix) removal [%d, %d) for seq %d is not supported " + "(block selection is anchored to cache slots) - rejected\n", __func__, p0, p1, sid); + return false; + } + } + } + if (seq_id >= 0) { auto & cells = v_cells[seq_to_stream[seq_id]]; auto & head = v_heads[seq_to_stream[seq_id]]; @@ -846,6 +906,10 @@ bool llama_kv_cache::update(llama_context * lctx, bool do_shift, const stream_co if (layer.v_stream[ssrc]) { ggml_backend_tensor_copy(layer.v_stream[ssrc], layer.v_stream[sdst]); } + if (layer.k_idx_stream[ssrc]) { + GGML_ASSERT(layer.k_idx_stream[sdst]); + ggml_backend_tensor_copy(layer.k_idx_stream[ssrc], layer.k_idx_stream[sdst]); + } } } } @@ -994,6 +1058,44 @@ llama_kv_cache::slot_info llama_kv_cache::find_slot(const llama_ubatch & ubatch, const auto & cells = v_cells[seq_to_stream[seq_id]]; + if (n_tokens > cells.size()) { + LLAMA_LOG_ERROR("%s: n_tokens = %d > size = %u\n", __func__, n_tokens, cells.size()); + return { }; + } + + // MSA block selection assumes slot == logical position (append-only streams). + if (msa_strict_slots) { + for (uint32_t ii = 0; ii < n_tokens; ++ii) { + const llama_pos pos = ubatch.pos[s*n_tokens + ii]; + + if (pos < 0 || (uint64_t) pos >= cells.size()) { + LLAMA_LOG_WARN("%s: MSA: position %d is outside the cache range [0, %u)\n", + __func__, pos, cells.size()); + return { }; + } + + const uint32_t idx = (uint32_t) pos; + + if (!cells.is_empty(idx)) { + LLAMA_LOG_WARN("%s: MSA: required slot %u is already occupied (stream %u)\n", + __func__, idx, seq_to_stream[seq_id]); + return { }; + } + + // strictly increasing positions, rules out duplicates and, for contiguous requests, is tightened to exact adjacency + if (!res.idxs[s].empty() && (cont ? idx != res.idxs[s].back() + 1 + : idx <= res.idxs[s].back())) { + LLAMA_LOG_WARN("%s: MSA: token positions are not %s within the ubatch\n", + __func__, cont ? "contiguous" : "strictly increasing"); + return { }; + } + + res.idxs[s].push_back(idx); + } + + continue; + } + uint32_t head_cur = v_heads[seq_to_stream[seq_id]]; // if we have enough unused cells before the current head -> @@ -1002,11 +1104,6 @@ llama_kv_cache::slot_info llama_kv_cache::find_slot(const llama_ubatch & ubatch, head_cur = 0; } - if (n_tokens > cells.size()) { - LLAMA_LOG_ERROR("%s: n_tokens = %d > size = %u\n", __func__, n_tokens, cells.size()); - return { }; - } - uint32_t n_tested = 0; // for continuous slots, we test that all tokens in the ubatch fit, starting from the current head @@ -1113,6 +1210,15 @@ void llama_kv_cache::apply_ubatch(const slot_info & sinfo, const llama_ubatch & const auto idx = sinfo.idxs[s][ii]; + if (msa_strict_slots && (llama_pos) idx != ubatch.pos[i]) { + LLAMA_LOG_ERROR("%s: MSA slot/position invariant violated: " + "writing pos %d into cell %u (stream %u). The indexer cache " + "would desync and block selection would silently corrupt. " + "This is a bug, please report it with reproduction steps.\n", + __func__, ubatch.pos[i], idx, sinfo.strm[s]); + GGML_ABORT("MSA: slot != pos"); + } + if (!cells.is_empty(idx)) { assert(cells.seq_count(idx) == 1); @@ -1156,7 +1262,8 @@ void llama_kv_cache::apply_ubatch(const slot_info & sinfo, const llama_ubatch & LLAMA_LOG_DEBUG("%s: purging positions [%d, %d] of sequence %d from KV cache\n", __func__, cells.seq_pos_min(s), seq_pos_max_rm[s], s); - seq_rm(s, cells.seq_pos_min(s), seq_pos_max_rm[s] + 1); + // under MSA strict slots this path should be unreachable, since strict MSA placement never selects occupied cells + GGML_ASSERT(seq_rm(s, cells.seq_pos_min(s), seq_pos_max_rm[s] + 1)); } } @@ -1176,6 +1283,12 @@ bool llama_kv_cache::get_can_shift() const { if (hparams.n_pos_per_embd() > 1) { return false; } + // shifting would leave k_idx stale + for (const auto & layer : layers) { + if (layer.k_idx) { + return false; + } + } return true; } @@ -1292,6 +1405,23 @@ ggml_tensor * llama_kv_cache::get_v(ggml_context * ctx, int32_t il, uint32_t n_k ggml_row_size(v->type, kv_size*n_embd_v_gqa)*sinfo.s0); } +ggml_tensor * llama_kv_cache::get_k_idx(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo) const { + const int32_t ikv = map_layer_ids.at(il); + auto * k_idx = layers[ikv].k_idx; + GGML_ASSERT(k_idx); + + const uint64_t kv_size = get_size(); + const int64_t n_idx = k_idx->ne[0]; // 128 + const uint32_t ns = sinfo.s1 - sinfo.s0 + 1; + + return ggml_view_4d(ctx, k_idx, + n_idx, 1, n_kv, ns, + ggml_row_size(k_idx->type, n_idx), // nb1 (single head) + ggml_row_size(k_idx->type, n_idx), // nb2 (per cell) + ggml_row_size(k_idx->type, n_idx*kv_size), // nb3 (per stream) + ggml_row_size(k_idx->type, n_idx*kv_size)*sinfo.s0); +} + ggml_tensor * llama_kv_cache::cpy_k(ggml_context * ctx, ggml_tensor * k_cur, ggml_tensor * k_idxs, int32_t il, const slot_info & sinfo) const { GGML_UNUSED(sinfo); @@ -1393,6 +1523,28 @@ ggml_tensor * llama_kv_cache::build_input_k_idxs(ggml_context * ctx, const llama return k_idxs; } +ggml_tensor * llama_kv_cache::cpy_k_idx(ggml_context * ctx, ggml_tensor * k_idx_cur, ggml_tensor * k_idxs, int32_t il, const slot_info & sinfo) const { + GGML_UNUSED(sinfo); + const int32_t ikv = map_layer_ids.at(il); + ggml_tensor * k_idx = layers[ikv].k_idx; + GGML_ASSERT(k_idx && "cpy_k_idx on a layer with no indexer cache"); + + const int64_t n_embd_head = k_idx_cur->ne[0]; // 128 + const int64_t n_head = k_idx_cur->ne[1]; // 1 + const int64_t n_tokens = k_idx_cur->ne[2]; + const int64_t n_embd_gqa = n_embd_head*n_head; // 128 + + GGML_ASSERT(ggml_row_size(k_idx_cur->type, n_embd_head) == k_idx_cur->nb[1]); + k_idx_cur = ggml_view_2d(ctx, k_idx_cur, n_embd_gqa, n_tokens, k_idx_cur->nb[2], 0); + + const int64_t n_stream = k_idx->ne[2]; + if (n_stream > 1) { + const int64_t kv_size = get_size(); + k_idx = ggml_reshape_2d(ctx, k_idx, n_embd_gqa, kv_size*n_stream); + } + return ggml_set_rows(ctx, k_idx, k_idx_cur, k_idxs); // same k_idxs as the K store +} + ggml_tensor * llama_kv_cache::build_input_v_idxs(ggml_context * ctx, const llama_ubatch & ubatch) const { const uint32_t n_tokens = ubatch.n_tokens; @@ -1827,6 +1979,18 @@ size_t llama_kv_cache::size_v_bytes() const { return size_v_bytes; } +size_t llama_kv_cache::size_k_idx_bytes() const { + size_t size_k_idx_bytes = 0; + + for (const auto & layer : layers) { + if (layer.k_idx) { + size_k_idx_bytes += ggml_nbytes(layer.k_idx); + } + } + + return size_k_idx_bytes; +} + ggml_tensor * llama_kv_cache::build_rope_shift( const llama_cparams & cparams, ggml_context * ctx, @@ -2139,6 +2303,36 @@ void llama_kv_cache::state_write_data(llama_io_write_i & io, const cell_ranges_t } } + if (size_k_idx_bytes() > 0) { + const uint32_t has_k_idx_u32 = 1; + io.write(&has_k_idx_u32, sizeof(has_k_idx_u32)); + + for (const auto & layer : layers) { + const uint32_t layer_has_k_idx = layer.k_idx ? 1 : 0; + io.write(&layer_has_k_idx, sizeof(layer_has_k_idx)); + + if (!layer_has_k_idx) { + continue; + } + + GGML_ASSERT(layer.k_idx_stream[cr.strm]); + + const int32_t k_idx_type_i = (int32_t) layer.k_idx->type; + io.write(&k_idx_type_i, sizeof(k_idx_type_i)); + + const uint64_t k_idx_size_row = ggml_row_size(layer.k_idx->type, layer.k_idx->ne[0]); + io.write(&k_idx_size_row, sizeof(k_idx_size_row)); + + for (const auto & range : cr.data) { + const size_t range_size = range.second - range.first; + const size_t buf_size = range_size * k_idx_size_row; + const size_t offset = range.first * k_idx_size_row; + + io.write_tensor(layer.k_idx_stream[cr.strm], offset, buf_size); + } + } + } + if (!v_trans) { for (const auto & layer : layers) { const uint32_t il = layer.il; @@ -2387,6 +2581,68 @@ bool llama_kv_cache::state_read_data(llama_io_read_i & io, uint32_t strm, uint32 } } + if (size_k_idx_bytes() > 0) { + uint32_t has_k_idx_u32 = 0; + io.read(&has_k_idx_u32, sizeof(has_k_idx_u32)); + + if (has_k_idx_u32 != 1) { + LLAMA_LOG_ERROR("%s: missing k_idx data in KV cache state\n", __func__); + return false; + } + + for (const auto & layer : layers) { + uint32_t layer_has_k_idx = 0; + io.read(&layer_has_k_idx, sizeof(layer_has_k_idx)); + + const uint32_t expected_layer_has_k_idx = layer.k_idx ? 1 : 0; + + if (layer_has_k_idx != expected_layer_has_k_idx) { + LLAMA_LOG_ERROR( + "%s: mismatched k_idx state for layer: got %u, expected %u\n", + __func__, layer_has_k_idx, expected_layer_has_k_idx); + return false; + } + + if (!layer_has_k_idx) { + continue; + } + + GGML_ASSERT(layer.k_idx_stream[strm]); + + int32_t k_idx_type_i = -1; + io.read(&k_idx_type_i, sizeof(k_idx_type_i)); + + if (k_idx_type_i != (int32_t) layer.k_idx->type) { + LLAMA_LOG_ERROR( + "%s: mismatched k_idx type: got %d, expected %d\n", + __func__, k_idx_type_i, (int32_t) layer.k_idx->type); + return false; + } + + uint64_t k_idx_size_row = 0; + io.read(&k_idx_size_row, sizeof(k_idx_size_row)); + + const uint64_t expected_k_idx_size_row = ggml_row_size(layer.k_idx->type, layer.k_idx->ne[0]); + + if (k_idx_size_row != expected_k_idx_size_row) { + LLAMA_LOG_ERROR( + "%s: mismatched k_idx row size: got %zu, expected %zu\n", + __func__, (size_t) k_idx_size_row, (size_t) expected_k_idx_size_row); + return false; + } + + if (cell_count) { + if (sinfo.is_contiguous()) { + io.read_tensor(layer.k_idx_stream[strm], sinfo.head() * k_idx_size_row, cell_count * k_idx_size_row); + } else { + for (uint32_t i = 0; i < cell_count; ++i) { + io.read_tensor(layer.k_idx_stream[strm], sinfo.idxs[0][i] * k_idx_size_row, k_idx_size_row); + } + } + } + } + } + if (!this->v_trans) { for (const auto & layer : layers) { const uint32_t il = layer.il; @@ -2588,6 +2844,10 @@ ggml_tensor * llama_kv_cache_context::get_v(ggml_context * ctx, int32_t il) cons return kv->get_v(ctx, il, n_kv, sinfos[i_cur]); } +ggml_tensor * llama_kv_cache_context::get_k_idx(ggml_context * ctx, int32_t il) const { + return kv->get_k_idx(ctx, il, n_kv, sinfos[i_cur]); +} + ggml_tensor * llama_kv_cache_context::cpy_k(ggml_context * ctx, ggml_tensor * k_cur, ggml_tensor * k_idxs, int32_t il) const { return kv->cpy_k(ctx, k_cur, k_idxs, il, sinfos[i_cur]); } @@ -2596,6 +2856,10 @@ ggml_tensor * llama_kv_cache_context::cpy_v(ggml_context * ctx, ggml_tensor * v_ return kv->cpy_v(ctx, v_cur, v_idxs, il, sinfos[i_cur]); } +ggml_tensor * llama_kv_cache_context::cpy_k_idx(ggml_context * ctx, ggml_tensor * k_idx_cur, ggml_tensor * k_idxs, int32_t il) const { + return kv->cpy_k_idx(ctx, k_idx_cur, k_idxs, il, sinfos[i_cur]); +} + ggml_tensor * llama_kv_cache_context::build_input_k_idxs(ggml_context * ctx, const llama_ubatch & ubatch) const { return kv->build_input_k_idxs(ctx, ubatch); } diff --git a/src/llama-kv-cache.h b/src/llama-kv-cache.h index 531d99dbdec1..d5a92f4405b5 100644 --- a/src/llama-kv-cache.h +++ b/src/llama-kv-cache.h @@ -173,10 +173,12 @@ class llama_kv_cache : public llama_memory_i { // get views of the current state of the cache ggml_tensor * get_k(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo) const; ggml_tensor * get_v(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo) const; + ggml_tensor * get_k_idx(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo) const; // store k_cur and v_cur in the cache based on the provided head location ggml_tensor * cpy_k(ggml_context * ctx, ggml_tensor * k_cur, ggml_tensor * k_idxs, int32_t il, const slot_info & sinfo) const; ggml_tensor * cpy_v(ggml_context * ctx, ggml_tensor * v_cur, ggml_tensor * v_idxs, int32_t il, const slot_info & sinfo) const; + ggml_tensor * cpy_k_idx(ggml_context * ctx, ggml_tensor * k_idx_cur, ggml_tensor * k_idxs, int32_t il, const slot_info & sinfo) const; // // preparation API @@ -228,9 +230,11 @@ class llama_kv_cache : public llama_memory_i { ggml_tensor * k; ggml_tensor * v; + ggml_tensor * k_idx; // MSA single-head indexer keys, F32 std::vector k_stream; std::vector v_stream; + std::vector k_idx_stream; }; bool v_trans = true; // the value tensor is transposed @@ -259,6 +263,9 @@ class llama_kv_cache : public llama_memory_i { // env: LLAMA_KV_CACHE_DEBUG int debug = 0; + // set when a k_idx (indexer) cache exists and the stream layout supports MSA (single seq, or one stream per seq) + bool msa_strict_slots = false; + // this is the SWA type of the cache - not to be confused with the model SWA type const llama_swa_type swa_type = LLAMA_SWA_TYPE_NONE; @@ -291,6 +298,7 @@ class llama_kv_cache : public llama_memory_i { size_t size_k_bytes() const; size_t size_v_bytes() const; + size_t size_k_idx_bytes() const; ggml_tensor * build_rope_shift( const llama_cparams & cparams, @@ -370,6 +378,7 @@ class llama_kv_cache_context : public llama_memory_context_i { // get views of the current state of the cache ggml_tensor * get_k(ggml_context * ctx, int32_t il) const; ggml_tensor * get_v(ggml_context * ctx, int32_t il) const; + ggml_tensor * get_k_idx(ggml_context * ctx, int32_t il) const; // store k_cur and v_cur in the cache based on the provided head location // note: the heads in k_cur and v_cur should be laid out contiguously in memory @@ -379,6 +388,7 @@ class llama_kv_cache_context : public llama_memory_context_i { // - v_idxs [n_tokens] or [n_tokens*n_embd_v_gqa] depending if V cache is transposed ggml_tensor * cpy_k(ggml_context * ctx, ggml_tensor * k_cur, ggml_tensor * k_idxs, int32_t il) const; ggml_tensor * cpy_v(ggml_context * ctx, ggml_tensor * v_cur, ggml_tensor * v_idxs, int32_t il) const; + ggml_tensor * cpy_k_idx(ggml_context * ctx, ggml_tensor * k_idx_cur, ggml_tensor * k_idxs, int32_t il) const; // create destination indices for each head of the current batch for where it would be written in the KV cache // the indices address the global KV cache (not per stream) - this is not relevant for the user of this API, but diff --git a/src/llama-model-saver.cpp b/src/llama-model-saver.cpp index d26e2ff7af62..3812c594e795 100644 --- a/src/llama-model-saver.cpp +++ b/src/llama-model-saver.cpp @@ -281,6 +281,8 @@ void llama_model_saver::add_kv_from_model() { add_kv(LLM_KV_ATTENTION_INDEXER_HEAD_COUNT, hparams.indexer_n_head); add_kv(LLM_KV_ATTENTION_INDEXER_KEY_LENGTH, hparams.indexer_head_size); add_kv(LLM_KV_ATTENTION_INDEXER_TOP_K, hparams.indexer_top_k); + add_kv(LLM_KV_ATTENTION_INDEXER_BLOCK_SIZE, hparams.indexer_block_size); + add_kv(LLM_KV_ATTENTION_INDEXER_LOCAL_BLOCKS, hparams.indexer_local_blocks); add_kv(LLM_KV_ATTENTION_INDEXER_TYPES, hparams.is_indexer_full_impl, true); add_kv(LLM_KV_ATTENTION_RECURRENT_LAYERS, hparams.is_recr_impl, true); diff --git a/src/llama-model.cpp b/src/llama-model.cpp index b100f6018150..51796921081f 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -285,6 +285,8 @@ static llama_model * llama_model_mapping(llm_arch arch, const llama_model_params return new llama_model_apertus(params); case LLM_ARCH_MINIMAX_M2: return new llama_model_minimax_m2(params); + case LLM_ARCH_MINIMAX_M3: + return new llama_model_minimax_m3(params); case LLM_ARCH_COGVLM: return new llama_model_cogvlm(params); case LLM_ARCH_PANGU_EMBED: @@ -818,6 +820,7 @@ const char * llm_type_name(llm_type type) { case LLM_TYPE_122B_A10B: return "122B.A10B"; case LLM_TYPE_196B_A11B: return "196B.A11B"; case LLM_TYPE_230B_A10B: return "230B.A10B"; + case LLM_TYPE_428B_A23B: return "428B.A23B"; case LLM_TYPE_235B_A22B: return "235B.A22B"; case LLM_TYPE_300B_A47B: return "300B.A47B"; case LLM_TYPE_310B_A15B: return "310B.A15B"; @@ -2550,6 +2553,7 @@ llama_rope_type llama_model_rope_type(const llama_model * model) { case LLM_ARCH_GROVEMOE: case LLM_ARCH_APERTUS: case LLM_ARCH_MINIMAX_M2: + case LLM_ARCH_MINIMAX_M3: case LLM_ARCH_COGVLM: case LLM_ARCH_PANGU_EMBED: case LLM_ARCH_AFMOE: diff --git a/src/llama-model.h b/src/llama-model.h index 45b054cedf1d..36d0480e5eb7 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -134,6 +134,7 @@ enum llm_type { LLM_TYPE_122B_A10B, // Qwen3.5 LLM_TYPE_196B_A11B, // Step3.5-Flash LLM_TYPE_230B_A10B, // Minimax M2 + LLM_TYPE_428B_A23B, // Minimax M3 LLM_TYPE_235B_A22B, LLM_TYPE_300B_A47B, // Ernie MoE big LLM_TYPE_310B_A15B, // /MiMo-V2-Flash @@ -515,6 +516,12 @@ struct llama_layer { struct ggml_tensor * indexer_attn_k = nullptr; struct ggml_tensor * indexer_attn_q_b = nullptr; // note: for lora a/b, not bias + // MSA + struct ggml_tensor * index_q_proj = nullptr; + struct ggml_tensor * index_k_proj = nullptr; + struct ggml_tensor * index_q_norm = nullptr; + struct ggml_tensor * index_k_norm = nullptr; + // gemma4 layer output scale, reused for talkie embedding skip scale struct ggml_tensor * out_scale = nullptr; diff --git a/src/llama-vocab.cpp b/src/llama-vocab.cpp index 7b312d1d88e1..9164a4dd888d 100644 --- a/src/llama-vocab.cpp +++ b/src/llama-vocab.cpp @@ -2809,6 +2809,7 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) { || t.first == "" // gemma4 || t.first == "<|tool_response>" // gemma4 || t.first == "<｜end▁of▁sentence｜>" // deepseek-ocr + || t.first == "[e~[" // minimax-m2/m3 ) { special_eog_ids.insert(t.second); if ((attr & LLAMA_TOKEN_ATTR_CONTROL) == 0) { diff --git a/src/models/minimax-m3.cpp b/src/models/minimax-m3.cpp new file mode 100644 index 000000000000..6068fc6b87a2 --- /dev/null +++ b/src/models/minimax-m3.cpp @@ -0,0 +1,562 @@ +#include "models.h" +#include "llama-kv-cache.h" +#include +#include +#include +#include + +// MiniMax-M3: MiniMax-M2 style GQA (per-head QK-norm, partial rotary) with +// DeepSeek-V3 leading-dense + routed/shared experts (sigmoid gating, routed scaling), +// swigluoai activation, and MiniMax Sparse Attention (MSA). MTP is not in released model weights. +// Notes: Blocks are anchored to absolute KV cache slots. + +void llama_model_minimax_m3::load_arch_hparams(llama_model_loader & ml) { + ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps); + ml.get_key(LLM_KV_LEADING_DENSE_BLOCK_COUNT, hparams.n_layer_dense_lead, false); + ml.get_key(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp); + ml.get_key(LLM_KV_EXPERT_SHARED_COUNT, hparams.n_expert_shared); + ml.get_key(LLM_KV_EXPERT_WEIGHTS_SCALE, hparams.expert_weights_scale, false); + ml.get_key(LLM_KV_EXPERT_WEIGHTS_NORM, hparams.expert_weights_norm, false); + ml.get_key(LLM_KV_EXPERT_GATING_FUNC, hparams.expert_gating_func); + ml.get_key(LLM_KV_ATTENTION_INDEXER_HEAD_COUNT, hparams.indexer_n_head); + ml.get_key(LLM_KV_ATTENTION_INDEXER_KEY_LENGTH, hparams.indexer_head_size); + ml.get_key(LLM_KV_ATTENTION_INDEXER_TOP_K, hparams.indexer_top_k); + ml.get_key(LLM_KV_ATTENTION_INDEXER_BLOCK_SIZE, hparams.indexer_block_size); + ml.get_key(LLM_KV_ATTENTION_INDEXER_LOCAL_BLOCKS, hparams.indexer_local_blocks); + msa_p = { (int) hparams.indexer_block_size, (int) hparams.indexer_top_k, (int) hparams.indexer_local_blocks }; + hparams.indexer_kv = true; + + switch (hparams.n_layer()) { + case 60: type = LLM_TYPE_428B_A23B; break; + default: type = LLM_TYPE_UNKNOWN; + } +} + +void llama_model_minimax_m3::load_arch_tensors(llama_model_loader &) { + LLAMA_LOAD_LOCALS; + const int64_t n_expert_shared = hparams.n_expert_shared; + const int64_t n_ff_exp = hparams.n_ff_exp; + + tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0); + + // output + output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0); + output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_vocab}, 0); + + for (int i = 0; i < n_layer; ++i) { + auto & layer = layers[i]; + + create_tensor_qkv(layer, i, n_embd, n_embd_head_k * n_head, n_embd_gqa, n_embd_gqa, 0); + layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), { n_embd_head_k * n_head, n_embd }, 0); + + layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0); + // per-head QK-norm: a single head_dim vector applied to every head + layer.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", i), {n_embd_head_k}, 0); + layer.attn_k_norm = create_tensor(tn(LLM_TENSOR_ATTN_K_NORM, "weight", i), {n_embd_head_k}, 0); + + layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0); + + if (i < (int) hparams.n_layer_dense_lead) { + // leading dense layers + layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd, n_ff}, 0); + layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), { n_ff, n_embd}, 0); + layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, n_ff}, 0); + } else { + // routed experts + layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, 0); + layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert}, 0); + layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", i), {n_embd, n_ff_exp, n_expert}, 0); + layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff_exp, n_embd, n_expert}, 0); + layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), {n_embd, n_ff_exp, n_expert}, 0); + + // shared expert + layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", i), {n_embd, n_ff_exp * n_expert_shared}, 0); + layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), { n_ff_exp * n_expert_shared, n_embd}, 0); + layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", i), {n_embd, n_ff_exp * n_expert_shared}, 0); + + // indexer + layer.index_q_proj = create_tensor(tn(LLM_TENSOR_INDEXER_Q_PROJ, "weight", i), {n_embd, hparams.indexer_n_head * hparams.indexer_head_size}, 0); + layer.index_k_proj = create_tensor(tn(LLM_TENSOR_INDEXER_K_PROJ, "weight", i), {n_embd, hparams.indexer_head_size}, 0); + layer.index_q_norm = create_tensor(tn(LLM_TENSOR_INDEXER_Q_NORM, "weight", i), {hparams.indexer_head_size}, 0); + layer.index_k_norm = create_tensor(tn(LLM_TENSOR_INDEXER_K_NORM, "weight", i), {hparams.indexer_head_size}, 0); + } + } +} + +std::unique_ptr llama_model_minimax_m3::build_arch_graph(const llm_graph_params & params) const { + return std::make_unique(*this, params); +} + +// per-query local-force bias for MSA selection +// local window always wins a slot +class llm_graph_input_msa_local : public llm_graph_input_i { +public: + llm_graph_input_msa_local(int blk, int local, int64_t nblk) : blk(blk), local(local), nblk(nblk) {} + + void set_input(const llama_ubatch * ubatch) override { + if (!bias || !ubatch->pos) { + return; + } + const int64_t n_tokens = ubatch->n_tokens; + std::vector data((size_t) nblk * n_tokens, 0.0f); + for (int64_t i = 0; i < n_tokens; ++i) { + const int64_t L = ubatch->pos[i] / blk; + for (int l = 0; l < local && L - l >= 0; ++l) { + if (L - l < nblk) { + data[(size_t) i * nblk + (L - l)] = 1e30f; + } + } + } + ggml_backend_tensor_set(bias, data.data(), 0, data.size() * sizeof(float)); + } + + // valid as long as the bias tensor dims still match the new ubatch/cache window + bool can_reuse(const llm_graph_params & params) override { + const auto * mctx = static_cast(params.mctx); + + bool res = true; + res &= bias->ne[1] == params.ubatch.n_tokens; + res &= bias->ne[0] * blk == (int64_t) mctx->get_n_kv(); + return res; + } + + ggml_tensor * bias = nullptr; + int blk; + int local; + int64_t nblk; +}; + +// pooled score of a block with no visible token: -inf from the mask, or -FLT_MAX from the +// max-pool identity when every element of the block is -inf +static inline bool msa_score_masked(float x) { return x <= -1e30f; } + +// MSA block selection (batch regime) +// CPU custom op, the token-level expansion and the combination with the causal mask happen on the GPU. +static void msa_block_mask_op(struct ggml_tensor * dst, int ith, int nth, void * userdata) { + const struct ggml_tensor * bs = dst->src[0]; + const struct ggml_tensor * bias = dst->src[1]; + const msa_params * p = (const msa_params *) userdata; + + const int nblk = (int) bs->ne[0]; + const int Hd = (int) bs->ne[1]; + const int S = (int) bs->ne[2]; + + GGML_ASSERT(bs->type == GGML_TYPE_F32 && ggml_is_contiguous(bs)); + GGML_ASSERT(bias->type == GGML_TYPE_F32 && ggml_is_contiguous(bias)); + GGML_ASSERT(dst->type == GGML_TYPE_F16 && ggml_is_contiguous(dst)); + GGML_ASSERT(dst->ne[0] == nblk && dst->ne[1] == S && dst->ne[2] == Hd); + GGML_ASSERT(bias->ne[0] == nblk && bias->ne[1] == S); + + const int topk = p->topk_blocks < nblk ? p->topk_blocks : nblk; + + const ggml_fp16_t f16_zero = ggml_fp32_to_fp16(0.0f); + const ggml_fp16_t f16_ninf = ggml_fp32_to_fp16(-INFINITY); + + std::vector rank(nblk); + std::vector valid(nblk); + std::vector ord(nblk); + + ggml_fp16_t * out = (ggml_fp16_t *) dst->data; + + for (int i = ith; i < S; i += nth) { + const float * bias_col = (const float *) bias->data + (size_t) i * nblk; + for (int h = 0; h < Hd; ++h) { + const float * bs_col = (const float *) bs->data + ((size_t) i * Hd + h) * nblk; + + for (int bk = 0; bk < nblk; ++bk) { + // a block is selectable if it has a visible token or is locally forced + valid[bk] = !msa_score_masked(bs_col[bk]) || bias_col[bk] > 0.0f; + rank [bk] = bias_col[bk] > 0.0f ? bias_col[bk] : bs_col[bk]; + ord [bk] = bk; + } + + std::partial_sort(ord.begin(), ord.begin() + topk, ord.end(), + [&](int a, int b) { return rank[a] > rank[b]; }); + + ggml_fp16_t * dst_col = out + ((size_t) h * S + i) * nblk; + for (int bk = 0; bk < nblk; ++bk) { + dst_col[bk] = f16_ninf; + } + for (int t = 0; t < topk; ++t) { + const int bk = ord[t]; + if (!valid[bk]) { + break; // sorted desc: first invalid -> fewer than topk selectable blocks + } + dst_col[bk] = f16_zero; + } + } + } +} + +// One FA call for all GQA groups (and at multi-stream decode, all streams) by mapping them onto the FA sequence dim (ne[3]) +ggml_tensor * llama_model_minimax_m3::graph::build_attn_msa_fa( + ggml_tensor * q_cur, // [D, HQ, T] + ggml_tensor * k, // [D, n_keys, 1, C] + ggml_tensor * v, // [D, n_keys, 1, C] + ggml_tensor * mask, // [n_keys, R, 1, C] f16, contiguous + int64_t Gp, float kq_scale, int il) const { + + const int64_t D = q_cur->ne[0]; + const int64_t HQ = q_cur->ne[1]; + const int64_t T = q_cur->ne[2]; + const int64_t C = k->ne[3]; + const int64_t R = HQ*T/(Gp*C); + GGML_ASSERT(Gp*C*R == HQ*T); + GGML_ASSERT(mask->type == GGML_TYPE_F16); + + // [D, HQ, T] -> [D, Gp, C, R] -> [D, R, Gp, C] + // batch (C=HKV, R=T): channel = group + // decode (C=HKV*ns, R=1): channel = (group, stream), group innermost + ggml_tensor * q = ggml_reshape_4d(ctx0, q_cur, D, Gp, C, R); + q = ggml_permute(ctx0, q, 0, 2, 3, 1); + + ggml_tensor * o = ggml_flash_attn_ext(ctx0, q, k, v, mask, kq_scale, + hparams.f_max_alibi_bias, 0.0f); + ggml_flash_attn_ext_set_prec(o, GGML_PREC_F32); + cb(o, "msa_fattn", il); + + // [D, Gp, R, C] -> [D, Gp, C, R] -> [n_embd, T] + o = ggml_permute(ctx0, o, 0, 1, 3, 2); + if (!ggml_is_contiguous(o)) { + o = ggml_cont(ctx0, o); // no-op layout at decode (R == 1), copy at batch + } + return ggml_reshape_2d(ctx0, o, D*HQ, T); +} + +llama_model_minimax_m3::graph::graph(const llama_model & model, const llm_graph_params & params) : llm_graph_context(params) { + const int64_t n_embd_head = hparams.n_embd_head_v(); + const auto & mm = static_cast(model); + + GGML_ASSERT(n_embd_head == hparams.n_embd_head_k()); + // partial rotary: head_dim != n_rot, so don't assert n_embd_head == n_rot + + ggml_tensor * cur; + ggml_tensor * inpL; + + inpL = build_inp_embd(model.tok_embd); + + ggml_tensor * inp_pos = build_inp_pos(); + auto inp_attn = build_attn_inp_kv(); + + // MSA calls ggml_flash_attn_ext directly and assumes the non-transposed V layout that + // llama.cpp only provides when flash attention is enabled. Block selection is anchored + // to absolute KV cache slots, which equal positions only for append-only per-stream + // caches either a single sequence, or multiple sequences with kv_unified == false (each + // stream then has its own slot space). A unified cache with multiple sequences + // interleaves slots and would silently break block anchoring so it falls back to dense. + const bool fa_on = cparams.flash_attn; + const bool streams_ok = cparams.n_seq_max == 1 || !cparams.kv_unified; + const bool msa_enabled = fa_on && streams_ok; + + static bool warned_no_fa = false; + if (!fa_on && !warned_no_fa) { + LLAMA_LOG_WARN("%s: flash attention disabled; MSA requires it -> running DENSE attention " + "(output may be degraded). Enable flash attention for MSA.\n", __func__); + warned_no_fa = true; + } + static bool warned_unified = false; + if (fa_on && !streams_ok && !warned_unified) { + LLAMA_LOG_WARN("%s: unified KV cache with n_seq_max > 1; MSA needs per-sequence streams " + "-> running DENSE attention. Output may be degraded. Drop --kv-unified to enable MSA.\n", __func__); + warned_unified = true; + } + + // hoisted per-graph MSA state (shared by every sparse layer) + llm_graph_input_msa_local * msa_loc = nullptr; + ggml_tensor * msa_kqm = nullptr; + ggml_tensor * msa_mf = nullptr; + int64_t n_kv = 0, nblk = 0, ns = 1, n_tps = 0; + bool msa_decode = false; // gather (1 token per stream) vs mask + const int blk = mm.msa_p.blk; + const int64_t Hd = hparams.indexer_n_head; // one indexer head per GQA group + + if (msa_enabled) { + msa_kqm = inp_attn->get_kq_mask(); + n_kv = msa_kqm->ne[0]; + n_tps = msa_kqm->ne[1]; // tokens per stream + ns = msa_kqm->ne[3]; // streams in this ubatch + GGML_ASSERT(msa_kqm->type == GGML_TYPE_F16 && "MSA requires the FA (f16) mask"); + GGML_ASSERT(n_tps*ns == n_tokens); + GGML_ASSERT(n_kv % blk == 0 && + "MSA: KV/mask n_kv must be a multiple of indexer.block_size (128); " + "the flash-attention KV padding must be a multiple of the block size. " + "A non-multiple would silently drop the partial tail block."); + nblk = n_kv / blk; + msa_decode = n_tps == 1; + + msa_mf = ggml_cast(ctx0, msa_kqm, GGML_TYPE_F32); + + auto loc = std::make_unique(blk, mm.msa_p.local, nblk); + loc->bias = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, nblk, n_tokens); // stream-grouped tokens + ggml_set_input(loc->bias); + msa_loc = (llm_graph_input_msa_local *) res->add_input(std::move(loc)); + } + + ggml_tensor * inp_out_ids = build_inp_out_ids(); + + for (int il = 0; il < n_layer; ++il) { + ggml_tensor * inpSA = inpL; + + // self-attention + { + cur = build_norm(inpL, model.layers[il].attn_norm, NULL, LLM_NORM_RMS, il); + cb(cur, "attn_norm", il); + + auto [Qcur, Kcur, Vcur] = build_qkv(model.layers[il], cur, + n_embd_head, n_head, n_head_kv, il); + + // per-head QK RMSNorm (weights already include Gemma's +1) + Qcur = build_norm(Qcur, model.layers[il].attn_q_norm, NULL, LLM_NORM_RMS, il); + cb(Qcur, "Qcur_normed", il); + Kcur = build_norm(Kcur, model.layers[il].attn_k_norm, NULL, LLM_NORM_RMS, il); + cb(Kcur, "Kcur_normed", il); + + // partial rotary: only the first n_rot dims are rotated + Qcur = ggml_rope_ext( + ctx0, Qcur, inp_pos, nullptr, + n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow); + Kcur = ggml_rope_ext( + ctx0, Kcur, inp_pos, nullptr, + n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow); + + cb(Qcur, "Qcur", il); + cb(Kcur, "Kcur", il); + cb(Vcur, "Vcur", il); + + const bool is_sparse = msa_enabled && il >= (int) hparams.n_layer_dense_lead; + + if (!is_sparse) { + cur = build_attn(inp_attn, model.layers[il].wo, NULL, model.layers[il].wo_s, + Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, + 1.0f/sqrtf(float(n_embd_head)), il); + } else { + const int64_t n_idx_dim = hparams.indexer_head_size; // 128 + + GGML_ASSERT(!inp_attn->self_k_rot && !inp_attn->self_v_rot && "MSA: attn-rot not supported"); + + // Index Branch, project, norm, partial RoPE, cache + ggml_tensor * iq = build_lora_mm(model.layers[il].index_q_proj, cur); + ggml_tensor * ik = build_lora_mm(model.layers[il].index_k_proj, cur); + iq = ggml_reshape_3d(ctx0, iq, n_idx_dim, Hd, n_tokens); + ik = ggml_reshape_3d(ctx0, ik, n_idx_dim, 1, n_tokens); + iq = build_norm(iq, model.layers[il].index_q_norm, NULL, LLM_NORM_RMS, il); // +1 baked + ik = build_norm(ik, model.layers[il].index_k_norm, NULL, LLM_NORM_RMS, il); + iq = ggml_rope_ext(ctx0, iq, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, + freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow); + ik = ggml_rope_ext(ctx0, ik, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, + freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow); + + const auto * mctx_cur = inp_attn->mctx; + ggml_build_forward_expand(gf, mctx_cur->cpy_k_idx(ctx0, ik, inp_attn->get_k_idxs(), il)); + ggml_tensor * ik_kv = mctx_cur->get_k_idx(ctx0, il); + + // Main branch: store K/V, take cache views + ggml_build_forward_expand(gf, Qcur); + ggml_build_forward_expand(gf, Kcur); + ggml_build_forward_expand(gf, Vcur); + ggml_build_forward_expand(gf, mctx_cur->cpy_k(ctx0, Kcur, inp_attn->get_k_idxs(), il)); + ggml_build_forward_expand(gf, mctx_cur->cpy_v(ctx0, Vcur, inp_attn->get_v_idxs(), il)); + ggml_tensor * k = mctx_cur->get_k(ctx0, il); + ggml_tensor * v = mctx_cur->get_v(ctx0, il); + GGML_ASSERT(!(v->nb[1] > v->nb[2]) && "MSA assumes v_trans=false (FA on)"); + + const int64_t D = k->ne[0]; + const int64_t HKV = k->ne[1]; + const int64_t Gp = n_head/HKV; + GGML_ASSERT(HKV == Hd && "MSA: one indexer head per GQA group"); + GGML_ASSERT(k->ne[3] == ns); + const int K = mm.msa_p.topk_blocks < (int) nblk ? mm.msa_p.topk_blocks : (int) nblk; + + const float kq_scale = 1.0f/sqrtf(float(n_embd_head)); + + if (msa_decode) { + // decode: batched over streams top-k + gather, one grouped FA + // scores: per-stream batched matmul over the stream dim (ne[3]). + // the cache views are not contiguous across streams (stride = kv_size, not n_kv) + ggml_tensor * ikv4 = ggml_view_4d(ctx0, ik_kv, n_idx_dim, n_kv, 1, ns, + ik_kv->nb[2], ik_kv->nb[3], ik_kv->nb[3], 0); + ggml_tensor * iq4 = ggml_reshape_4d(ctx0, iq, n_idx_dim, Hd, 1, ns); + ggml_tensor * sc = ggml_mul_mat(ctx0, ikv4, iq4); + ggml_mul_mat_set_prec(sc, GGML_PREC_F32); + sc = ggml_add_inplace(ctx0, sc, msa_mf); + ggml_tensor * bs = ggml_pool_2d(ctx0, sc, GGML_OP_POOL_MAX, blk, 1, blk, 1, 0, 0); + cb(bs, "msa_bs", il); + + ggml_tensor * bsf = ggml_add(ctx0, bs, + ggml_reshape_4d(ctx0, msa_loc->bias, nblk, 1, 1, ns)); + ggml_tensor * idx = ggml_top_k(ctx0, bsf, K); + + // token idx: tj[t,k,h,s] = blk*idx[k,h,s] + t (for the mask gather) + // row idx: tr[t,k,h,s] = tj*HKV + h (for the per-stream K/V gather) + ggml_tensor * a = ggml_scale(ctx0, ggml_cast(ctx0, idx, GGML_TYPE_F32), (float) blk); + a = ggml_reshape_4d(ctx0, a, 1, K, Hd, ns); + ggml_tensor * tj = ggml_add(ctx0, + ggml_repeat_4d(ctx0, a, blk, K, Hd, ns), + ggml_reshape_3d(ctx0, ggml_arange(ctx0, 0.0f, (float) blk, 1.0f), blk, 1, 1)); + ggml_tensor * tr = ggml_add(ctx0, + ggml_scale(ctx0, tj, (float) HKV), + ggml_reshape_3d(ctx0, ggml_arange(ctx0, 0.0f, (float) HKV, 1.0f), 1, 1, Hd)); + + ggml_tensor * tokj = ggml_cast(ctx0, ggml_reshape_2d(ctx0, tj, (int64_t) blk*K*Hd, ns), GGML_TYPE_I32); + ggml_tensor * tokr = ggml_cast(ctx0, ggml_reshape_2d(ctx0, tr, (int64_t) blk*K*Hd, ns), GGML_TYPE_I32); + + ggml_tensor * k3 = ggml_view_3d(ctx0, k, D, HKV*n_kv, ns, k->nb[1], k->nb[3], 0); + ggml_tensor * v3 = ggml_view_3d(ctx0, v, D, HKV*n_kv, ns, v->nb[1], v->nb[3], 0); + ggml_tensor * m3 = ggml_reshape_3d(ctx0, msa_kqm, 1, n_kv, ns); + + ggml_tensor * kg = ggml_get_rows(ctx0, k3, tokr); + ggml_tensor * vg = ggml_get_rows(ctx0, v3, tokr); + ggml_tensor * mg = ggml_get_rows(ctx0, m3, tokj); + + // fold (group, stream) onto the FA channel dim + const ggml_type kt = ggml_is_quantized(k->type) ? GGML_TYPE_F16 : k->type; + const ggml_type vt = ggml_is_quantized(v->type) ? GGML_TYPE_F16 : v->type; + ggml_tensor * kfa = ggml_reshape_4d(ctx0, kg, D, (int64_t) blk*K, 1, Hd*ns); + ggml_tensor * vfa = ggml_reshape_4d(ctx0, vg, D, (int64_t) blk*K, 1, Hd*ns); + if (kfa->type != kt) { kfa = ggml_cast(ctx0, kfa, kt); } + if (vfa->type != vt) { vfa = ggml_cast(ctx0, vfa, vt); } + // the FA mask must be F16 + ggml_tensor * mfa = ggml_cast(ctx0, ggml_reshape_4d(ctx0, mg, (int64_t) blk*K, 1, 1, Hd*ns), GGML_TYPE_F16); + + cur = build_attn_msa_fa(Qcur, kfa, vfa, mfa, Gp, kq_scale, il); + } else { + // batch: per-stream loop + std::vector outs(ns); + for (int64_t st = 0; st < ns; ++st) { + ggml_tensor * iq_s = ggml_view_3d(ctx0, iq, n_idx_dim, Hd, n_tps, + iq->nb[1], iq->nb[2], st*n_tps*iq->nb[2]); + ggml_tensor * ik_s = ggml_view_2d(ctx0, ik_kv, n_idx_dim, n_kv, + ik_kv->nb[2], st*ik_kv->nb[3]); + ggml_tensor * mf_s = ggml_view_3d(ctx0, msa_mf, n_kv, 1, n_tps, + msa_mf->nb[1], msa_mf->nb[1], st*msa_mf->nb[3]); + ggml_tensor * km_s = ggml_view_3d(ctx0, msa_kqm, n_kv, n_tps, 1, + msa_kqm->nb[1], msa_kqm->nb[3], st*msa_kqm->nb[3]); + ggml_tensor * bias_s = ggml_view_2d(ctx0, msa_loc->bias, nblk, n_tps, + msa_loc->bias->nb[1], st*n_tps*msa_loc->bias->nb[1]); + ggml_tensor * q_s = ggml_view_3d(ctx0, Qcur, D, n_head, n_tps, + Qcur->nb[1], Qcur->nb[2], st*n_tps*Qcur->nb[2]); + ggml_tensor * k_s = ggml_view_4d(ctx0, k, D, HKV, n_kv, 1, + k->nb[1], k->nb[2], k->nb[3], st*k->nb[3]); + ggml_tensor * v_s = ggml_view_4d(ctx0, v, D, HKV, n_kv, 1, + v->nb[1], v->nb[2], v->nb[3], st*v->nb[3]); + + // block scores: bs = maxpool_blk(idx_q * idx_k^T + causal mask) + // scores are unscaled, only the top-k ordering matters + ggml_tensor * sc = ggml_mul_mat(ctx0, ik_s, + ggml_reshape_2d(ctx0, iq_s, n_idx_dim, Hd*n_tps)); + // indexer scores run in F32 + ggml_mul_mat_set_prec(sc, GGML_PREC_F32); + sc = ggml_reshape_3d(ctx0, sc, n_kv, Hd, n_tps); + sc = ggml_add_inplace(ctx0, sc, mf_s); + ggml_tensor * bs = ggml_pool_2d(ctx0, sc, GGML_OP_POOL_MAX, blk, 1, blk, 1, 0, 0); + cb(bs, "msa_bs", il); + + // block-level 0/-inf keep mask on the CPU, tiny transfer + ggml_tensor * srcs[2] = { bs, bias_s }; + ggml_tensor * bm = ggml_custom_4d(ctx0, GGML_TYPE_F16, + nblk, n_tps, Hd, 1, + srcs, 2, msa_block_mask_op, GGML_N_TASKS_MAX, + const_cast(&mm.msa_p)); + cb(bm, "msa_block_mask", il); + + // expand block -> token granularity on the GPU (j = bk*blk + t), + // then combine with the causal mask in place + ggml_tensor * bmx = ggml_repeat_4d(ctx0, + ggml_reshape_3d(ctx0, bm, 1, nblk, n_tps*Hd), + blk, nblk, n_tps*Hd, 1); + bmx = ggml_reshape_3d(ctx0, bmx, n_kv, n_tps, Hd); + ggml_tensor * mask4 = ggml_add_inplace(ctx0, bmx, km_s); + mask4 = ggml_reshape_4d(ctx0, mask4, n_kv, n_tps, 1, Hd); + cb(mask4, "msa_mask4", il); + + // cache views with groups on ne[3]; + ggml_tensor * kfa = ggml_permute(ctx0, k_s, 0, 3, 1, 2); + ggml_tensor * vfa = ggml_permute(ctx0, v_s, 0, 3, 1, 2); + + outs[st] = build_attn_msa_fa(q_s, kfa, vfa, mask4, Gp, kq_scale, il); + } + cur = outs[0]; + for (int64_t st = 1; st < ns; ++st) { + cur = ggml_concat(ctx0, cur, outs[st], 1); + } + } + + cb(cur, "kqv_out", il); + if (model.layers[il].wo) { + cur = build_lora_mm(model.layers[il].wo, cur, model.layers[il].wo_s); + } + } + } + + if (il == n_layer - 1 && inp_out_ids) { + cur = ggml_get_rows(ctx0, cur, inp_out_ids); + inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids); + } + + ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpSA); + cb(ffn_inp, "ffn_inp", il); + + cur = build_norm(ffn_inp, model.layers[il].ffn_norm, NULL, LLM_NORM_RMS, il); + cb(cur, "ffn_norm", il); + + if ((uint32_t) il < hparams.n_layer_dense_lead) { + // leading dense FFN (swigluoai) + cur = build_ffn(cur, + model.layers[il].ffn_up, NULL, NULL, + model.layers[il].ffn_gate, NULL, NULL, + model.layers[il].ffn_down, NULL, NULL, + NULL, + LLM_FFN_SWIGLU_OAI_MOE, LLM_FFN_PAR, il); + cb(cur, "ffn_out", il); + } else { + // routed experts (swigluoai MoE) + ggml_tensor * moe_out = build_moe_ffn(cur, + model.layers[il].ffn_gate_inp, + model.layers[il].ffn_up_exps, + model.layers[il].ffn_gate_exps, + model.layers[il].ffn_down_exps, + model.layers[il].ffn_exp_probs_b, + n_expert, n_expert_used, + LLM_FFN_SWIGLU_OAI_MOE, hparams.expert_weights_norm, + hparams.expert_weights_scale, + (llama_expert_gating_func_type) hparams.expert_gating_func, + il); + cb(moe_out, "ffn_moe_out", il); + + // shared expert (swigluoai) + ggml_tensor * ffn_shexp = build_ffn(cur, + model.layers[il].ffn_up_shexp, NULL, NULL, + model.layers[il].ffn_gate_shexp, NULL, NULL, + model.layers[il].ffn_down_shexp, NULL, NULL, + NULL, + LLM_FFN_SWIGLU_OAI_MOE, LLM_FFN_PAR, il); + cb(ffn_shexp, "ffn_shexp", il); + + cur = ggml_add(ctx0, moe_out, ffn_shexp); + cb(cur, "ffn_out", il); + } + + cur = ggml_add(ctx0, cur, ffn_inp); + + cur = build_cvec(cur, il); + cb(cur, "l_out", il); + + // input for next layer + inpL = cur; + } + + cur = inpL; + + cur = build_norm(cur, model.output_norm, NULL, LLM_NORM_RMS, -1); + cb(cur, "result_norm", -1); + res->t_embd = cur; + + // lm_head + cur = build_lora_mm(model.output, cur, model.output_s); + cb(cur, "result_output", -1); + res->t_logits = cur; + + ggml_build_forward_expand(gf, cur); +} diff --git a/src/models/models.h b/src/models/models.h index 76daa8cc1994..916459e12782 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -1902,6 +1902,29 @@ struct llama_model_minimax_m2 : public llama_model_base { std::unique_ptr build_arch_graph(const llm_graph_params & params) const override; }; +struct msa_params { + int blk; + int topk_blocks; + int local; +}; + +struct llama_model_minimax_m3 : public llama_model_base { + llama_model_minimax_m3(const struct llama_model_params & params) : llama_model_base(params) {} + void load_arch_hparams(llama_model_loader & ml) override; + void load_arch_tensors(llama_model_loader & ml) override; + msa_params msa_p; + struct graph : public llm_graph_context { + graph(const llama_model & model, const llm_graph_params & params); + + ggml_tensor * build_attn_msa_fa( + ggml_tensor * q_cur, // [D, HQ, S] f32 + ggml_tensor * k, // [D, n_keys, 1, C] C = HKV or HKV*n_stream + ggml_tensor * v, // [D, n_keys, 1, C] + ggml_tensor * mask, // [n_keys, R, 1, C] f16, R = HQ*T/(Gp*C) + int64_t Gp, float kq_scale, int il) const; + }; + std::unique_ptr build_arch_graph(const llm_graph_params & params) const override; +}; struct llama_model_cogvlm : public llama_model_base { llama_model_cogvlm(const struct llama_model_params & params) : llama_model_base(params) {} diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index 86c3051c5fe5..d02e65c9ead0 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -168,6 +168,9 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { ms.add_kv(LLM_KV_ROPE_DIMENSION_COUNT, uint32_t(64)); ms.add_kv(LLM_KV_ATTENTION_KEY_LENGTH_MLA, uint32_t(192)); ms.add_kv(LLM_KV_ATTENTION_VALUE_LENGTH_MLA, uint32_t(128)); + } else if (arch == LLM_ARCH_MINIMAX_M3) { + // partial rotary: n_rot must not exceed the indexer key length (64) + ms.add_kv(LLM_KV_ROPE_DIMENSION_COUNT, uint32_t(64)); } ms.add_kv(LLM_KV_ATTENTION_CLAMP_KQV, 1.0f); ms.add_kv(LLM_KV_ATTENTION_LAYERNORM_EPS, 1e-5f); @@ -198,9 +201,13 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { ms.add_kv(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, uint32_t(2)); } - ms.add_kv(LLM_KV_ATTENTION_INDEXER_HEAD_COUNT, uint32_t(1)); - ms.add_kv(LLM_KV_ATTENTION_INDEXER_KEY_LENGTH, uint32_t(64)); - ms.add_kv(LLM_KV_ATTENTION_INDEXER_TOP_K, uint32_t(8)); + // MSA requires one indexer head per GQA (KV) head, unlike the DSA archs where the + // indexer head count is independent of the main attention head count. + ms.add_kv(LLM_KV_ATTENTION_INDEXER_HEAD_COUNT, arch == LLM_ARCH_MINIMAX_M3 ? n_head : uint32_t(1)); + ms.add_kv(LLM_KV_ATTENTION_INDEXER_KEY_LENGTH, uint32_t(64)); + ms.add_kv(LLM_KV_ATTENTION_INDEXER_TOP_K, uint32_t(8)); + ms.add_kv(LLM_KV_ATTENTION_INDEXER_BLOCK_SIZE, uint32_t(4)); + ms.add_kv(LLM_KV_ATTENTION_INDEXER_LOCAL_BLOCKS, uint32_t(1)); ms.add_kv(LLM_KV_ROPE_DIMENSION_SECTIONS, std::vector({n_embd_head/4, n_embd_head/4, n_embd_head/4, n_embd_head/4})); ms.add_kv(LLM_KV_TOKENIZER_MODEL, "no_vocab"); // ms.add_kv(LLM_KV_DENSE_2_FEAT_OUT, n_embd); @@ -355,6 +362,7 @@ static bool moe_mandatory(const llm_arch arch) { case LLM_ARCH_LLADA_MOE: case LLM_ARCH_GROVEMOE: case LLM_ARCH_MINIMAX_M2: + case LLM_ARCH_MINIMAX_M3: case LLM_ARCH_RND1: case LLM_ARCH_PADDLEOCR: case LLM_ARCH_MIMO2: From af285020e909aac47ed77842cf17bc08c33a32f7 Mon Sep 17 00:00:00 2001 From: Eric Hartford Date: Sun, 26 Jul 2026 14:43:51 -0400 Subject: [PATCH 008/190] mtmd: add GLM-5.2-Vision (#26126) Co-authored-by: Eric Hartford --- conversion/__init__.py | 1 + conversion/kimivl.py | 16 ++++++++++++++++ tools/mtmd/mtmd.cpp | 14 +++++++++++--- 3 files changed, 28 insertions(+), 3 deletions(-) diff --git a/conversion/__init__.py b/conversion/__init__.py index 0b08e6e57008..c5ecc68cfda2 100644 --- a/conversion/__init__.py +++ b/conversion/__init__.py @@ -269,6 +269,7 @@ "Gemma4UnifiedForConditionalGeneration": "gemma", "Glm4vForConditionalGeneration": "qwen3vl", "Glm4vMoeForConditionalGeneration": "qwen3vl", + "Glm5vForConditionalGeneration": "kimivl", "GlmOcrForConditionalGeneration": "qwen3vl", "GlmasrModel": "ultravox", "Granite4VisionForConditionalGeneration": "granite", diff --git a/conversion/kimivl.py b/conversion/kimivl.py index 63b8a079b722..5ff3c39ca9c1 100644 --- a/conversion/kimivl.py +++ b/conversion/kimivl.py @@ -152,3 +152,19 @@ def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iter name = name.replace(".proj.2.", ".proj.linear_2.") yield from super().modify_tensors(data_torch, name, bid) + + +@ModelBase.register("Glm5vForConditionalGeneration") +class Glm5vModel(KimiK25Model): + """GLM-5.2-Vision MoonViT3d encoder and projector + + Uses the same vision encoder and projector as Kimi-K2.5, so it reuses the + kimik25 projector type. The image begin/end tokens differ, but they are + resolved at runtime from the text model vocab. + """ + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + if name.startswith("mm_projector.linear_"): + name = name.replace("mm_projector.linear_", "mm_projector.proj.linear_", 1) + + yield from super().modify_tensors(data_torch, name, bid) diff --git a/tools/mtmd/mtmd.cpp b/tools/mtmd/mtmd.cpp index e10ccf186b14..5915b4cba967 100644 --- a/tools/mtmd/mtmd.cpp +++ b/tools/mtmd/mtmd.cpp @@ -555,9 +555,17 @@ struct mtmd_context { } break; case PROJECTOR_TYPE_KIMIK25: { - // <|media_begin|> ... (image embeddings) ... <|media_end|> - img_beg = "<|media_begin|>"; - img_end = "<|media_end|>"; + // GLM-5.2-V reuses the Kimi-K2.5 vision encoder and projector, but marks + // images with its own tokens, so decide based on the text model vocab + if (lookup_token("<|begin_of_image|>") != LLAMA_TOKEN_NULL) { + // <|begin_of_image|> ... (image embeddings) ... <|end_of_image|> + img_beg = "<|begin_of_image|>"; + img_end = "<|end_of_image|>"; + } else { + // <|media_begin|> ... (image embeddings) ... <|media_end|> + img_beg = "<|media_begin|>"; + img_end = "<|media_end|>"; + } image_preproc = std::make_unique(ctx_v); } break; case PROJECTOR_TYPE_LIGHTONOCR: From d2a818231effb12b7b20b80b3b8c7756a9a33a04 Mon Sep 17 00:00:00 2001 From: Xuan-Son Nguyen Date: Sun, 26 Jul 2026 20:54:25 +0200 Subject: [PATCH 009/190] common: add `subproc.h` wrapper, disabled on android/ios (#26102) * add common/subproc.h|cpp * add compile flag LLAMA_SUBPROCESS * disabled by default on android and ios * test-jinja: use common subproc * mtmd: disable video if subproc is not set * disable subproc on wasm * make is_created atomic * migrate server-mcp --- CMakeLists.txt | 9 +++ common/CMakeLists.txt | 6 ++ common/subproc.cpp | 143 +++++++++++++++++++++++++++++++++ common/subproc.h | 59 ++++++++++++++ tests/test-jinja.cpp | 25 +++--- tools/mtmd/CMakeLists.txt | 9 ++- tools/server/server-mcp.cpp | 34 +++----- tools/server/server-models.cpp | 66 ++++----------- tools/server/server-tools.cpp | 30 +++---- 9 files changed, 272 insertions(+), 109 deletions(-) create mode 100644 common/subproc.cpp create mode 100644 common/subproc.h diff --git a/CMakeLists.txt b/CMakeLists.txt index 81f23d7e70b7..3df1d82dbe09 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -84,6 +84,14 @@ else() set(LLAMA_TOOLS_INSTALL_DEFAULT ${LLAMA_STANDALONE}) endif() +# subprocess spawning isn't a supported/sandbox-friendly operation on mobile OSes or in WASM +if (CMAKE_SYSTEM_NAME STREQUAL "iOS" OR CMAKE_SYSTEM_NAME STREQUAL "Android" OR ANDROID + OR CMAKE_SYSTEM_NAME STREQUAL "Emscripten" OR EMSCRIPTEN) + set(LLAMA_SUBPROCESS_DEFAULT OFF) +else() + set(LLAMA_SUBPROCESS_DEFAULT ON) +endif() + # # option list # @@ -117,6 +125,7 @@ option(LLAMA_TESTS_INSTALL "llama: install tests" ON) # 3rd party libs option(LLAMA_OPENSSL "llama: use openssl to support HTTPS" ON) +option(LLAMA_SUBPROCESS "llama-common: use subprocess, required by server tools and server router mode" ${LLAMA_SUBPROCESS_DEFAULT}) option(LLAMA_LLGUIDANCE "llama-common: include LLGuidance library for structured output in common utils" OFF) diff --git a/common/CMakeLists.txt b/common/CMakeLists.txt index 99688f53b87b..799d227519f9 100644 --- a/common/CMakeLists.txt +++ b/common/CMakeLists.txt @@ -100,6 +100,8 @@ add_library(${TARGET} sampling.h speculative.cpp speculative.h + subproc.cpp + subproc.h trie.cpp trie.h unicode.cpp @@ -127,6 +129,10 @@ set_target_properties(${TARGET} PROPERTIES target_include_directories(${TARGET} PUBLIC . ../vendor) target_compile_features (${TARGET} PUBLIC cxx_std_17) +if (LLAMA_SUBPROCESS) + target_compile_definitions(${TARGET} PUBLIC LLAMA_SUBPROCESS) +endif() + if (BUILD_SHARED_LIBS) set_target_properties(${TARGET} PROPERTIES POSITION_INDEPENDENT_CODE ON) diff --git a/common/subproc.cpp b/common/subproc.cpp new file mode 100644 index 000000000000..6d37f59002b6 --- /dev/null +++ b/common/subproc.cpp @@ -0,0 +1,143 @@ +#include "subproc.h" + +bool common_subproc::is_supported() { +#ifdef LLAMA_SUBPROCESS + return true; +#else + return false; +#endif +} + +#ifdef LLAMA_SUBPROCESS + +static std::vector to_cstr_vec(const std::vector & v) { + std::vector r; + r.reserve(v.size() + 1); + for (const auto & s : v) { + r.push_back(const_cast(s.c_str())); + } + r.push_back(nullptr); + return r; +} + +common_subproc::~common_subproc() { + if (is_created) { + subprocess_destroy(&proc); + is_created = false; + } +} + +bool common_subproc::create( + const std::vector & args, + int options, + const std::vector & env, + const char * cwd) { + auto argv = to_cstr_vec(args); + + int result; + if (env.empty() && cwd == nullptr) { + result = subprocess_create(argv.data(), options, &proc); + } else { + auto envp = to_cstr_vec(env); + result = subprocess_create_ex(argv.data(), options, env.empty() ? nullptr : envp.data(), cwd, &proc); + } + + is_created = result == 0; + return is_created; +} + +bool common_subproc::has_handle() const { + if (!is_created) { + return false; + } +#if defined(_WIN32) + return proc.hProcess != nullptr; +#else + return proc.child > 0; +#endif +} + +bool common_subproc::alive() { + return is_created && subprocess_alive(&proc); +} + +FILE * common_subproc::stdin_file() { + return is_created ? subprocess_stdin(&proc) : nullptr; +} + +FILE * common_subproc::stdout_file() { + return is_created ? subprocess_stdout(&proc) : nullptr; +} + +FILE * common_subproc::stderr_file() { + return is_created ? subprocess_stderr(&proc) : nullptr; +} + +void common_subproc::close_stdin() { + if (is_created && proc.stdin_file) { + fclose(proc.stdin_file); + proc.stdin_file = nullptr; + } +} + +void common_subproc::terminate() { + if (has_handle()) { + subprocess_terminate(&proc); + } +} + +int common_subproc::join() { + int exit_code = -1; + if (is_created) { + subprocess_join(&proc, &exit_code); + subprocess_destroy(&proc); + is_created = false; + } + return exit_code; +} + +#else // !LLAMA_SUBPROCESS + +common_subproc::~common_subproc() = default; + +bool common_subproc::create( + const std::vector &, + int, + const std::vector &, + const char *) { + (void)(proc); + (void)(is_created); + return false; +} + +bool common_subproc::has_handle() const { + return false; +} + +bool common_subproc::alive() { + return false; +} + +FILE * common_subproc::stdin_file() { + return nullptr; +} + +FILE * common_subproc::stdout_file() { + return nullptr; +} + +FILE * common_subproc::stderr_file() { + return nullptr; +} + +void common_subproc::close_stdin() { +} + +void common_subproc::terminate() { +} + +int common_subproc::join() { + return -1; +} + +#endif // LLAMA_SUBPROCESS diff --git a/common/subproc.h b/common/subproc.h new file mode 100644 index 000000000000..89b69ee262fb --- /dev/null +++ b/common/subproc.h @@ -0,0 +1,59 @@ +#pragma once + +#include +#include +#include +#include + +#ifdef LLAMA_SUBPROCESS +#include +#else +// dummy values to allow compilation when subprocess is disabled +struct subprocess_s {}; +static constexpr int subprocess_option_no_window = 0; +static constexpr int subprocess_option_combined_stdout_stderr = 0; +static constexpr int subprocess_option_inherit_environment = 0; +static constexpr int subprocess_option_search_user_path = 0; +#endif + +// RAII-style wrapper around https://github.com/sheredom/subprocess.h, +// exposing method calls instead of free functions operating on subprocess_s. +struct common_subproc { + common_subproc() = default; + ~common_subproc(); + + common_subproc(const common_subproc &) = delete; + common_subproc & operator=(const common_subproc &) = delete; + + // spawn a child process; if env is non-empty it replaces the child's environment + // (do not combine with subprocess_option_inherit_environment) + bool create( + const std::vector & args, + int options, + const std::vector & env = {}, + const char * cwd = nullptr); + + bool alive(); + + // true if LLAMA_SUBPROCESS was enabled at build time; when false, create() always fails + static bool is_supported(); + + FILE * stdin_file(); + FILE * stdout_file(); + FILE * stderr_file(); + + // close stdin and detach it from the process, so a later join()/destroy() won't double-close it; + // use this after writing all input to signal EOF to the child while it's still running + void close_stdin(); + + void terminate(); + + // wait for the process to exit, release the underlying handle and return its exit code + int join(); + +private: + subprocess_s proc {}; + std::atomic is_created{false}; + + bool has_handle() const; +}; diff --git a/tests/test-jinja.cpp b/tests/test-jinja.cpp index 90bdbc445d52..1ac5b57decca 100644 --- a/tests/test-jinja.cpp +++ b/tests/test-jinja.cpp @@ -4,7 +4,7 @@ #include #include -#include +#include "subproc.h" #include "jinja/runtime.h" #include "jinja/parser.h" @@ -2135,21 +2135,20 @@ static void test_template_py(testing & t, const std::string & name, const std::s const char * python_executable = "python3"; #endif - const char * command_line[] = {python_executable, "-c", py_script.c_str(), NULL}; + std::vector args = {python_executable, "-c", py_script, }; - struct subprocess_s subprocess; + common_subproc subprocess; int options = subprocess_option_combined_stdout_stderr | subprocess_option_no_window | subprocess_option_inherit_environment | subprocess_option_search_user_path; - int result = subprocess_create(command_line, options, &subprocess); - if (result != 0) { - t.log("Failed to create subprocess, error code: " + std::to_string(result)); + if (!subprocess.create(args, options)) { + t.log("Failed to create subprocess"); t.assert_true("subprocess creation", false); return; } - FILE * p_stdin = subprocess_stdin(&subprocess); + FILE * p_stdin = subprocess.stdin_file(); // Write input std::string input = merged.dump(); @@ -2157,24 +2156,22 @@ static void test_template_py(testing & t, const std::string & name, const std::s if (written != input.size()) { t.log("Failed to write complete input to subprocess stdin"); t.assert_true("subprocess stdin write", false); - subprocess_destroy(&subprocess); + subprocess.close_stdin(); + subprocess.join(); return; } fflush(p_stdin); - fclose(p_stdin); // Close stdin to signal EOF to the Python process - subprocess.stdin_file = nullptr; + subprocess.close_stdin(); // Close stdin to signal EOF to the Python process // Read output std::string output; char buffer[1024]; - FILE * p_stdout = subprocess_stdout(&subprocess); + FILE * p_stdout = subprocess.stdout_file(); while (fgets(buffer, sizeof(buffer), p_stdout)) { output += buffer; } - int process_return; - subprocess_join(&subprocess, &process_return); - subprocess_destroy(&subprocess); + int process_return = subprocess.join(); if (process_return != 0) { t.log("Python script failed with exit code: " + std::to_string(process_return)); diff --git a/tools/mtmd/CMakeLists.txt b/tools/mtmd/CMakeLists.txt index ea684d9f156d..d0329ca56743 100644 --- a/tools/mtmd/CMakeLists.txt +++ b/tools/mtmd/CMakeLists.txt @@ -1,8 +1,15 @@ # mtmd -set(MTMD_VIDEO ON CACHE BOOL "enable video support in mtmd (requires ffmpeg binary in PATH)") +set(MTMD_VIDEO_HELP "enable video support in mtmd (requires ffmpeg binary in PATH)") + +set(MTMD_VIDEO ON CACHE BOOL "${MTMD_VIDEO_HELP}") # TODO: add MTMD_VIDEO_METHOD in the future to select between ffmpeg and other backends +if (MTMD_VIDEO AND NOT LLAMA_SUBPROCESS) + message(STATUS "Disabling MTMD_VIDEO because LLAMA_SUBPROCESS is OFF") + set(MTMD_VIDEO OFF CACHE BOOL "${MTMD_VIDEO_HELP}" FORCE) +endif() + find_package(Threads REQUIRED) add_library(mtmd diff --git a/tools/server/server-mcp.cpp b/tools/server/server-mcp.cpp index aa87afeb065a..93db6164d34e 100644 --- a/tools/server/server-mcp.cpp +++ b/tools/server/server-mcp.cpp @@ -1,6 +1,6 @@ #include "server-mcp.h" -#include +#include "subproc.h" #include #include @@ -341,7 +341,7 @@ json server_mcp_transport::call_tool(const std::string & tool_name, // struct server_mcp_stdio::process_handle { - subprocess_s sp; + common_subproc sp; FILE * in = nullptr; // child stdin FILE * out = nullptr; // child stdout FILE * err = nullptr; // child stderr @@ -483,30 +483,15 @@ bool server_mcp_stdio::start() { envp_s = mcp_build_env(config.env); } - auto to_ptrs = [](std::vector & v) { - std::vector p; - p.reserve(v.size() + 1); - for (auto & s : v) { - p.push_back(s.c_str()); - } - p.push_back(nullptr); - return p; - }; - auto argv = to_ptrs(argv_s); - auto envp = to_ptrs(envp_s); - auto handle = std::make_unique(); - int rc = subprocess_create_ex(argv.data(), options, - config.env.empty() ? nullptr : envp.data(), - config.cwd.empty() ? nullptr : config.cwd.c_str(), - &handle->sp); - if (rc != 0) { + bool ok = handle->sp.create(argv_s, options, envp_s, config.cwd.empty() ? nullptr : config.cwd.c_str()); + if (!ok) { SRV_WRN("MCP '%s': failed to spawn '%s'\n", config.name.c_str(), config.command.c_str()); return false; } - handle->in = subprocess_stdin(&handle->sp); - handle->out = subprocess_stdout(&handle->sp); - handle->err = subprocess_stderr(&handle->sp); + handle->in = handle->sp.stdin_file(); + handle->out = handle->sp.stdout_file(); + handle->err = handle->sp.stderr_file(); proc = std::move(handle); running.store(true); @@ -654,14 +639,13 @@ void server_mcp_stdio::join_pumps() { to_server.close_write(); // wake the writer if it waits for a message from_server.close_write(); // wake any caller waiting for a reply - subprocess_terminate(&proc->sp); // child death unblocks the blocked fread/fwrite + proc->sp.terminate(); // child death unblocks the blocked fread/fwrite if (writer.joinable()) writer.join(); if (reader.joinable()) reader.join(); if (errlog.joinable()) errlog.join(); - subprocess_join(&proc->sp, nullptr); // reap the child: destroy() never waits, so the pid would stay a zombie for the process lifetime - subprocess_destroy(&proc->sp); // safe now: no thread touches the FILE* anymore + proc->sp.join(); // reap the child: never waiting would leave the pid a zombie for the process lifetime proc.reset(); } diff --git a/tools/server/server-models.cpp b/tools/server/server-models.cpp index ba63788146fe..923b3533e9cb 100644 --- a/tools/server/server-models.cpp +++ b/tools/server/server-models.cpp @@ -8,10 +8,10 @@ #include "preset.h" #include "download.h" #include "http.h" +#include "subproc.h" #include // TODO: remove this once we use HTTP client from download.h #include -#include #include #include @@ -49,43 +49,24 @@ extern char **environ; #define CHILD_ADDR "127.0.0.1" struct server_subproc { - std::optional sproc; // empty while in DOWNLOADING state + common_subproc sproc; // not yet spawned while in DOWNLOADING state std::atomic stopped{false}; // set to cancel a download or signal child process exit - subprocess_s & get() { - GGML_ASSERT(sproc.has_value() && "subprocess not initialized"); - return sproc.value(); - } - bool is_alive() { - return sproc.has_value() && subprocess_alive(&sproc.value()); + return sproc.alive(); } void request_exit() { - if (sproc.has_value()) { - FILE * stdin_file = subprocess_stdin(&sproc.value()); - if (stdin_file) { - fprintf(stdin_file, "%s\n", CMD_ROUTER_TO_CHILD_EXIT); - fflush(stdin_file); - } + FILE * stdin_file = sproc.stdin_file(); + if (stdin_file) { + fprintf(stdin_file, "%s\n", CMD_ROUTER_TO_CHILD_EXIT); + fflush(stdin_file); } stopped.store(true, std::memory_order_relaxed); } void terminate() { - if (!sproc.has_value()) { - return; - } -#if defined(_WIN32) - if (sproc->hProcess == NULL) { - return; - } -#else - if (sproc->child <= 0) { - return; - } -#endif - subprocess_terminate(&sproc.value()); + sproc.terminate(); } }; @@ -711,18 +692,6 @@ std::optional server_models::get_meta(const std::string & nam return std::nullopt; } -// helper to convert vector to char ** -// pointers are only valid as long as the original vector is valid -static std::vector to_char_ptr_array(const std::vector & vec) { - std::vector result; - result.reserve(vec.size() + 1); - for (const auto & s : vec) { - result.push_back(const_cast(s.c_str())); - } - result.push_back(nullptr); - return result; -} - std::vector server_models::get_all_meta() { std::unique_lock lk(mutex); if (need_reload) { @@ -845,15 +814,10 @@ void server_models::load(const std::string & name, const load_options & opts) { } inst.meta.args = child_args; // save for debugging - std::vector argv = to_char_ptr_array(child_args); - std::vector envp = to_char_ptr_array(child_env); - // TODO @ngxson : maybe separate stdout and stderr in the future // so that we can use stdout for commands and stderr for logging int options = subprocess_option_no_window | subprocess_option_combined_stdout_stderr; - inst.subproc->sproc.emplace(); - int result = subprocess_create_ex(argv.data(), options, envp.data(), nullptr, &inst.subproc->get()); - if (result != 0) { + if (!inst.subproc->sproc.create(child_args, options, child_env)) { throw std::runtime_error("failed to spawn server instance"); } } @@ -867,8 +831,8 @@ void server_models::load(const std::string & name, const load_options & opts) { stop_timeout = inst.meta.stop_timeout, child_mode = opts.mode ]() { - FILE * stdin_file = subprocess_stdin(&child_proc->get()); - FILE * stdout_file = subprocess_stdout(&child_proc->get()); // combined stdout/stderr + FILE * stdin_file = child_proc->sproc.stdin_file(); + FILE * stdout_file = child_proc->sproc.stdout_file(); // combined stdout/stderr std::thread log_thread([&]() { // read stdout/stderr and forward to main server log @@ -942,9 +906,7 @@ void server_models::load(const std::string & name, const load_options & opts) { } // get the exit code - int exit_code = 0; - subprocess_join(&child_proc->get(), &exit_code); - subprocess_destroy(&child_proc->get()); + int exit_code = child_proc->sproc.join(); // update status and exit code if (child_mode == SERVER_CHILD_MODE_DOWNLOAD) { @@ -1567,6 +1529,10 @@ static std::optional resolve_child_for_conv( } void server_models_routes::init_routes() { + if (!common_subproc::is_supported()) { + throw std::runtime_error("subprocess is not enabled on this build"); + } + this->get_router_props = [this](const server_http_req & req) { std::string name = req.get_param("model"); if (name.empty()) { diff --git a/tools/server/server-tools.cpp b/tools/server/server-tools.cpp index 2af44e49b9d4..90b7e2a9f0ee 100644 --- a/tools/server/server-tools.cpp +++ b/tools/server/server-tools.cpp @@ -1,6 +1,6 @@ #include "server-tools.h" -#include +#include "subproc.h" #include #include @@ -138,15 +138,14 @@ class tools_io_basic : public tools_io { const std::function & on_chunk = nullptr) const override { exec_result res; - subprocess_s proc; - auto argv = to_cstr_vec(args); + common_subproc proc; int options = subprocess_option_no_window | subprocess_option_combined_stdout_stderr | subprocess_option_inherit_environment | subprocess_option_search_user_path; - if (subprocess_create(argv.data(), options, &proc) != 0) { + if (!proc.create(args, options)) { res.output = "failed to spawn process"; return res; } @@ -159,14 +158,14 @@ class tools_io_basic : public tools_io { while (!done.load()) { if (std::chrono::steady_clock::now() >= deadline) { timed_out.store(true); - subprocess_terminate(&proc); + proc.terminate(); return; } std::this_thread::sleep_for(std::chrono::milliseconds(100)); } }); - FILE * f = subprocess_stdout(&proc); + FILE * f = proc.stdout_file(); std::string output; bool truncated = false; if (f) { @@ -177,7 +176,7 @@ class tools_io_basic : public tools_io { if (output.size() + len <= max_output) { output.append(buf, len); if (on_chunk && !on_chunk(std::string(buf, len))) { - subprocess_terminate(&proc); + proc.terminate(); break; } } else { @@ -195,8 +194,7 @@ class tools_io_basic : public tools_io { timeout_thread.join(); } - subprocess_join(&proc, &res.exit_code); - subprocess_destroy(&proc); + res.exit_code = proc.join(); res.output = output; res.timed_out = timed_out.load(); @@ -207,16 +205,6 @@ class tools_io_basic : public tools_io { } private: - static std::vector to_cstr_vec(const std::vector & v) { - std::vector r; - r.reserve(v.size() + 1); - for (const auto & s : v) { - r.push_back(const_cast(s.c_str())); - } - r.push_back(nullptr); - return r; - } - static const std::unordered_set & junk_dir_names() { static const std::unordered_set names = { ".git", ".svn", ".hg", "node_modules", "__pycache__", @@ -1203,6 +1191,10 @@ static std::vector> build_tools() { void server_tools::setup(const std::vector & enabled_tools, server_mcp & mcp_mgr) { if (!enabled_tools.empty()) { + if (!common_subproc::is_supported()) { + throw std::runtime_error("subprocess is not enabled on this build"); + } + std::unordered_set enabled_set(enabled_tools.begin(), enabled_tools.end()); auto all_tools = build_tools(); From 55b7d6c4c7e518005bc320cf0139264411deac7f Mon Sep 17 00:00:00 2001 From: Pascal Date: Sun, 26 Jul 2026 23:32:58 +0200 Subject: [PATCH 010/190] ui: detect the conversation import format from file contents (#26121) * ui: detect the conversation import format from file contents iOS resolves every accept entry to a UTI and has none for ".jsonl", so the picker greyed out exported conversations. Drop the accept filter and pick the parser from the file contents: ZIP magic bytes, then a first "session" record for JSONL, otherwise the legacy JSON format. Also remove the unused importConversations() picker and an orphan doc comment, and cover each format with unit tests. * ui: report what a conversation import actually wrote The import summary echoed the selection back, so re-importing conversations already in the database claimed success while nothing was written and only a console warning said otherwise. Return the imported and skipped conversations from the database layer, list the written ones in the summary, and count the rest in a toast. * ui: name the literals of the JSONL conversation format Introduce SessionRecordType and SESSION_HARNESS, and reuse the existing NEWLINE constant, so the record format lives in one place. This also covers the writer side, which predates the import path under review and carried the same literals: an enum stated by the reader alone lets the two sides drift. Values are unchanged, so an export stays byte identical. --- .../SettingsChatImportExportTab.svelte | 16 ++- .../src/lib/constants/conversation-import.ts | 3 + tools/ui/src/lib/constants/index.ts | 1 + tools/ui/src/lib/constants/message-export.ts | 3 + .../lib/enums/conversation-import.enums.ts | 9 ++ tools/ui/src/lib/enums/index.ts | 2 + tools/ui/src/lib/services/database.service.ts | 14 +- .../ui/src/lib/stores/conversations.svelte.ts | 135 +++++++----------- .../src/lib/utils/modality-file-validation.ts | 6 - .../conversation-import-db.svelte.test.ts | 68 +++++++++ .../ui/tests/unit/conversation-import.test.ts | 112 +++++++++++++++ 11 files changed, 270 insertions(+), 99 deletions(-) create mode 100644 tools/ui/src/lib/constants/conversation-import.ts create mode 100644 tools/ui/src/lib/enums/conversation-import.enums.ts create mode 100644 tools/ui/tests/client/conversation-import-db.svelte.test.ts create mode 100644 tools/ui/tests/unit/conversation-import.test.ts diff --git a/tools/ui/src/lib/components/app/settings/SettingsChat/SettingsChatImportExportTab.svelte b/tools/ui/src/lib/components/app/settings/SettingsChat/SettingsChatImportExportTab.svelte index a86d68584c01..57dbba30bc39 100644 --- a/tools/ui/src/lib/components/app/settings/SettingsChat/SettingsChatImportExportTab.svelte +++ b/tools/ui/src/lib/components/app/settings/SettingsChat/SettingsChatImportExportTab.svelte @@ -159,8 +159,10 @@ try { const input = document.createElement('input'); + // No `accept` filter: iOS resolves each entry to a UTI and has none for + // `.jsonl`, which greys out exported conversations in the file picker. + // `parseImportFile` detects the format from the file contents instead. input.type = HtmlInputType.FILE; - input.accept = `${FileExtensionText.JSON},${FileExtensionText.JSONL},${FileExtensionText.ZIP}`; input.onchange = async (e) => { const file = (e.target as HTMLInputElement)?.files?.[0]; @@ -199,9 +201,17 @@ .snapshot(fullImportData) .filter((item) => selectedIds.has(item.conv.id)); - await conversationsStore.importConversationsData(selectedData); + const { imported, skipped } = await conversationsStore.importConversationsData(selectedData); - importedConversations = selectedConversations; + // A conversation already in the database is left untouched, so the summary + // lists what was written and the toast accounts for the rest. + if (skipped.length > 0) { + toast.info( + `Skipped ${skipped.length} conversation${skipped.length === 1 ? '' : 's'} already in your library` + ); + } + + importedConversations = imported; showImportSummary = true; showExportSummary = false; showImportDialog = false; diff --git a/tools/ui/src/lib/constants/conversation-import.ts b/tools/ui/src/lib/constants/conversation-import.ts new file mode 100644 index 000000000000..ed500440a427 --- /dev/null +++ b/tools/ui/src/lib/constants/conversation-import.ts @@ -0,0 +1,3 @@ +// First bytes of every ZIP local file header ("PK"). Import detects an archive +// from these bytes rather than from the filename, which the OS may not preserve. +export const ZIP_MAGIC = [0x50, 0x4b]; diff --git a/tools/ui/src/lib/constants/index.ts b/tools/ui/src/lib/constants/index.ts index a80e0cb63379..100432c18c0d 100644 --- a/tools/ui/src/lib/constants/index.ts +++ b/tools/ui/src/lib/constants/index.ts @@ -13,6 +13,7 @@ export * from './storage'; export * from './attachment-menu'; export * from './auto-scroll'; export * from './context-gauge-popup'; +export * from './conversation-import'; export * from './binary-detection'; export * from './built-in-tools'; export * from './cache'; diff --git a/tools/ui/src/lib/constants/message-export.ts b/tools/ui/src/lib/constants/message-export.ts index 79fa36f9141d..fc4dbe259c28 100644 --- a/tools/ui/src/lib/constants/message-export.ts +++ b/tools/ui/src/lib/constants/message-export.ts @@ -7,6 +7,9 @@ export const EXPORT_CONV_NAME_SUFFIX_MAX_LENGTH = 20; // Characters to keep in the ISO timestamp. 19 keeps 2026-01-01T00:00:00 export const ISO_TIMESTAMP_SLICE_LENGTH = 19; +// Producer marker carried by the session record of a JSONL export +export const SESSION_HARNESS = 'llama.app'; + // Replacements for making the conversation title filename-friendly export const NON_ALPHANUMERIC_REGEX = /[^a-z0-9]/gi; export const EXPORT_CONV_NONALNUM_REPLACEMENT = '_'; diff --git a/tools/ui/src/lib/enums/conversation-import.enums.ts b/tools/ui/src/lib/enums/conversation-import.enums.ts new file mode 100644 index 000000000000..eef47c5cc1c4 --- /dev/null +++ b/tools/ui/src/lib/enums/conversation-import.enums.ts @@ -0,0 +1,9 @@ +/** + * Discriminator of a record line in the JSONL conversation format. A session + * record opens a conversation and carries its properties; every following + * message record belongs to it. + */ +export enum SessionRecordType { + SESSION = 'session', + MESSAGE = 'message' +} diff --git a/tools/ui/src/lib/enums/index.ts b/tools/ui/src/lib/enums/index.ts index 2f70e063d4e0..ee14293fc952 100644 --- a/tools/ui/src/lib/enums/index.ts +++ b/tools/ui/src/lib/enums/index.ts @@ -27,6 +27,8 @@ export { ReasoningFormat } from './chat.enums'; +export { SessionRecordType } from './conversation-import.enums'; + export { ReasoningEffort } from './reasoning-effort.enums'; export { diff --git a/tools/ui/src/lib/services/database.service.ts b/tools/ui/src/lib/services/database.service.ts index 4fb70e29af7a..bc65caaca386 100644 --- a/tools/ui/src/lib/services/database.service.ts +++ b/tools/ui/src/lib/services/database.service.ts @@ -554,12 +554,13 @@ export class DatabaseService { * Skips conversations that already exist. * * @param data - Array of { conv, messages } objects + * @returns The conversations written to the database and the ones skipped */ static async importConversations( data: { conv: DatabaseConversation; messages: DatabaseMessage[] }[] - ): Promise<{ imported: number; skipped: number }> { - let importedCount = 0; - let skippedCount = 0; + ): Promise<{ imported: DatabaseConversation[]; skipped: DatabaseConversation[] }> { + const imported: DatabaseConversation[] = []; + const skipped: DatabaseConversation[] = []; return await db.transaction( 'rw', @@ -570,8 +571,7 @@ export class DatabaseService { const existing = await db[IDXDB_TABLES.conversations].get(conv.id); if (existing) { - console.warn(`Conversation "${conv.name}" already exists, skipping...`); - skippedCount++; + skipped.push(conv); continue; } @@ -580,10 +580,10 @@ export class DatabaseService { await db[IDXDB_TABLES.messages].put(msg); } - importedCount++; + imported.push(conv); } - return { imported: importedCount, skipped: skippedCount }; + return { imported, skipped }; } ); } diff --git a/tools/ui/src/lib/stores/conversations.svelte.ts b/tools/ui/src/lib/stores/conversations.svelte.ts index bc2feefd3d95..e467c8fad973 100644 --- a/tools/ui/src/lib/stores/conversations.svelte.ts +++ b/tools/ui/src/lib/stores/conversations.svelte.ts @@ -30,11 +30,11 @@ import type { McpServerOverride } from '$lib/types/database'; import { zipSync, unzipSync, strToU8, strFromU8 } from 'fflate'; import { MessageRole, - HtmlInputType, FileExtensionText, MimeTypeText, MimeTypeApplication, - ReasoningEffort + ReasoningEffort, + SessionRecordType } from '$lib/enums'; import { ISO_DATE_TIME_SEPARATOR, @@ -47,7 +47,10 @@ import { ISO_TIME_SEPARATOR_REPLACEMENT, NON_ALPHANUMERIC_REGEX, MULTIPLE_UNDERSCORE_REGEX, - REASONING_EFFORT_DEFAULT_LOCALSTORAGE_KEY + REASONING_EFFORT_DEFAULT_LOCALSTORAGE_KEY, + NEWLINE, + SESSION_HARNESS, + ZIP_MAGIC } from '$lib/constants'; import { ROUTES } from '$lib/constants/routes'; @@ -914,30 +917,35 @@ class ConversationsStore { /** * Serializes a session (a conversation with its messages) as JSONL. - * The first line is the session header (a `type: 'session'` record carrying the - * conversation properties); each subsequent line is a single message. + * The first line is the session header (a `SessionRecordType.SESSION` record + * carrying the conversation properties); each subsequent line is a single message. * @param data - The exported conversation payload * @returns The JSONL string (one record per line) */ serializeSessionToJsonl(data: ExportedConversation): string { const { conv, messages } = data; - const sessionLine = JSON.stringify({ type: 'session', harness: 'llama.app', ...conv }); + const sessionLine = JSON.stringify({ + type: SessionRecordType.SESSION, + harness: SESSION_HARNESS, + ...conv + }); const messageLines = messages.map((message: DatabaseMessage) => { // `toolCalls` is stored as a JSON string; drop it when empty, otherwise parse it. const { toolCalls, ...rest } = message; const normalized = toolCalls ? { ...rest, toolCalls: JSON.parse(toolCalls) } : rest; - return JSON.stringify({ type: 'message', message: normalized }); + return JSON.stringify({ type: SessionRecordType.MESSAGE, message: normalized }); }); - return [sessionLine, ...messageLines].join('\n'); + return [sessionLine, ...messageLines].join(NEWLINE); } /** * Parses the JSONL session format produced by {@link serializeSessionToJsonl}. - * A `type: 'session'` line starts a new session; following `type: 'message'` - * lines are appended to it. Supports multiple sessions in a single file. + * A `SessionRecordType.SESSION` line starts a new session; following + * `SessionRecordType.MESSAGE` lines are appended to it. Supports multiple + * sessions in a single file. * @param text - The JSONL file contents * @returns The parsed conversations with their messages */ @@ -945,20 +953,20 @@ class ConversationsStore { const sessions: ExportedConversation[] = []; let current: ExportedConversation | null = null; - for (const line of text.split('\n')) { + for (const line of text.split(NEWLINE)) { const trimmed = line.trim(); if (!trimmed) continue; const record = JSON.parse(trimmed); - if (record.type === 'session') { + if (record.type === SessionRecordType.SESSION) { // Drop the discriminator and harness marker; the rest is the conversation. const conv = { ...record }; delete conv.type; delete conv.harness; current = { conv: conv as DatabaseConversation, messages: [] }; sessions.push(current); - } else if (record.type === 'message') { + } else if (record.type === SessionRecordType.MESSAGE) { if (!current) { throw new Error('Invalid JSONL: message record before any session record'); } @@ -977,27 +985,47 @@ class ConversationsStore { } /** - * Parses an import file into conversations, accepting the current `.jsonl` and - * `.zip` formats as well as the legacy `.json` format. + * Reports whether the text is the JSONL session format, whose first non-empty + * line is a `SessionRecordType.SESSION` record. A legacy JSON export starts + * with an array or an object that has no such discriminator. + * @param text - The file contents + */ + private isSessionsJsonl(text: string): boolean { + const trimmed = text.trimStart(); + const lineEnd = trimmed.indexOf(NEWLINE); + const firstLine = lineEnd === -1 ? trimmed : trimmed.slice(0, lineEnd); + + try { + return JSON.parse(firstLine).type === SessionRecordType.SESSION; + } catch { + // Not a standalone JSON record, so not the JSONL format. + return false; + } + } + + /** + * Parses an import file into conversations, accepting the current JSONL and + * ZIP formats as well as the legacy JSON format. The format comes from the + * contents, so an import works whatever the file is named. * @param file - The user-selected file * @returns The parsed conversations with their messages */ async parseImportFile(file: File): Promise { - const name = file.name.toLowerCase(); + const bytes = new Uint8Array(await file.arrayBuffer()); - if (name.endsWith(FileExtensionText.ZIP)) { - const entries = unzipSync(new Uint8Array(await file.arrayBuffer())); + if (ZIP_MAGIC.every((byte, index) => bytes[index] === byte)) { + const entries = unzipSync(bytes); const sessions: ExportedConversation[] = []; - for (const [entryName, bytes] of Object.entries(entries)) { + for (const [entryName, entryBytes] of Object.entries(entries)) { if (!entryName.toLowerCase().endsWith(FileExtensionText.JSONL)) continue; - sessions.push(...this.parseSessionsJsonl(strFromU8(bytes))); + sessions.push(...this.parseSessionsJsonl(strFromU8(entryBytes))); } return sessions; } - const text = await file.text(); + const text = strFromU8(bytes); - if (name.endsWith(FileExtensionText.JSONL)) { + if (this.isSessionsJsonl(text)) { return this.parseSessionsJsonl(text); } @@ -1103,73 +1131,14 @@ class ConversationsStore { this.downloadConversationFile({ conv: conversation, messages }); } - /** - * Imports conversations from a JSON file - * Opens file picker and processes the selected file - * @returns The list of imported conversations - */ - async importConversations(): Promise { - return new Promise((resolve, reject) => { - const input = document.createElement('input'); - input.type = HtmlInputType.FILE; - input.accept = FileExtensionText.JSON; - - input.onchange = async (e) => { - const file = (e.target as HTMLInputElement)?.files?.[0]; - - if (!file) { - reject(new Error('No file selected')); - return; - } - - try { - const text = await file.text(); - const parsedData = JSON.parse(text); - let importedData: ExportedConversations; - - if (Array.isArray(parsedData)) { - importedData = parsedData; - } else if ( - parsedData && - typeof parsedData === 'object' && - 'conv' in parsedData && - 'messages' in parsedData - ) { - importedData = [parsedData]; - } else { - throw new Error('Invalid file format'); - } - - const result = await DatabaseService.importConversations(importedData); - toast.success(`Imported ${result.imported} conversation(s), skipped ${result.skipped}`); - - await this.loadConversations(); - - const importedConversations = ( - Array.isArray(importedData) ? importedData : [importedData] - ).map((item) => item.conv); - - resolve(importedConversations); - } catch (err: unknown) { - const message = err instanceof Error ? err.message : 'Unknown error'; - console.error('Failed to import conversations:', err); - toast.error('Import failed', { description: message }); - reject(new Error(`Import failed: ${message}`)); - } - }; - - input.click(); - }); - } - /** * Imports conversations from provided data (without file picker) * @param data - Array of conversation data with messages - * @returns Import result with counts + * @returns The conversations written to the database and the ones skipped */ async importConversationsData( data: ExportedConversations - ): Promise<{ imported: number; skipped: number }> { + ): Promise<{ imported: DatabaseConversation[]; skipped: DatabaseConversation[] }> { const result = await DatabaseService.importConversations(data); await this.loadConversations(); return result; diff --git a/tools/ui/src/lib/utils/modality-file-validation.ts b/tools/ui/src/lib/utils/modality-file-validation.ts index bf78a70088cf..bfdee75ce323 100644 --- a/tools/ui/src/lib/utils/modality-file-validation.ts +++ b/tools/ui/src/lib/utils/modality-file-validation.ts @@ -161,9 +161,3 @@ export function generateModalityErrorMessage( return message; } - -/** - * Generate file input accept string based on model modalities - * @param capabilities - The modality capabilities to check against - * @returns Accept string for HTML file input element - */ diff --git a/tools/ui/tests/client/conversation-import-db.svelte.test.ts b/tools/ui/tests/client/conversation-import-db.svelte.test.ts new file mode 100644 index 000000000000..2a27be9b0cb7 --- /dev/null +++ b/tools/ui/tests/client/conversation-import-db.svelte.test.ts @@ -0,0 +1,68 @@ +import { afterEach, describe, expect, it } from 'vitest'; +import { DatabaseService } from '$lib/services/database.service'; +import { MessageRole, MessageType } from '$lib/enums'; +import type { ExportedConversation } from '$lib/types/database'; + +function makeSession(id: string): ExportedConversation { + return { + conv: { id, currNode: `${id}-msg`, lastModified: 0, name: `Chat ${id}` }, + messages: [ + { + id: `${id}-msg`, + convId: id, + type: MessageType.TEXT, + timestamp: 0, + role: MessageRole.USER, + content: `hello from ${id}`, + parent: null, + children: [] + } + ] + } as unknown as ExportedConversation; +} + +afterEach(async () => { + const conversations = await DatabaseService.getAllConversations(); + await DatabaseService.bulkDeleteConversations(conversations.map((conv) => conv.id)); +}); + +/** + * An import leaves a conversation already in the database untouched, so the + * caller needs to know what was written to report it instead of echoing the + * selection back at the user. + */ +describe('DatabaseService.importConversations', () => { + it('reports the conversations it wrote', async () => { + const { imported, skipped } = await DatabaseService.importConversations([ + makeSession('a'), + makeSession('b') + ]); + + expect(imported.map((conv) => conv.id)).toEqual(['a', 'b']); + expect(skipped).toEqual([]); + expect(await DatabaseService.getConversationMessages('a')).toHaveLength(1); + }); + + it('reports an existing conversation as skipped and leaves it untouched', async () => { + await DatabaseService.importConversations([makeSession('a')]); + await DatabaseService.updateConversation('a', { name: 'Renamed locally' }); + + const { imported, skipped } = await DatabaseService.importConversations([makeSession('a')]); + + expect(imported).toEqual([]); + expect(skipped.map((conv) => conv.id)).toEqual(['a']); + expect((await DatabaseService.getConversation('a'))?.name).toBe('Renamed locally'); + }); + + it('imports the new conversations of a partially known selection', async () => { + await DatabaseService.importConversations([makeSession('a')]); + + const { imported, skipped } = await DatabaseService.importConversations([ + makeSession('a'), + makeSession('b') + ]); + + expect(imported.map((conv) => conv.id)).toEqual(['b']); + expect(skipped.map((conv) => conv.id)).toEqual(['a']); + }); +}); diff --git a/tools/ui/tests/unit/conversation-import.test.ts b/tools/ui/tests/unit/conversation-import.test.ts new file mode 100644 index 000000000000..4565ed5b5d44 --- /dev/null +++ b/tools/ui/tests/unit/conversation-import.test.ts @@ -0,0 +1,112 @@ +import { beforeAll, describe, expect, it } from 'vitest'; +import { zipSync, strToU8 } from 'fflate'; +import { MessageRole, MessageType } from '$lib/enums'; +import { NEWLINE } from '$lib/constants'; +import type { ExportedConversation } from '$lib/types/database'; + +let conversationsStore: typeof import('$lib/stores/conversations.svelte').conversationsStore; + +// node env unit project has no DOM, install a minimal localStorage backed by a +// Map before the store module reads it. Transforming the store takes seconds, +// so import it once for the whole file. +beforeAll(async () => { + const store = new Map(); + const polyfill: Storage = { + get length() { + return store.size; + }, + clear: () => store.clear(), + getItem: (k) => (store.has(k) ? store.get(k)! : null), + key: (i) => Array.from(store.keys())[i] ?? null, + removeItem: (k) => { + store.delete(k); + }, + setItem: (k, v) => { + store.set(k, String(v)); + } + }; + (globalThis as unknown as { localStorage: Storage }).localStorage = polyfill; + + ({ conversationsStore } = await import('$lib/stores/conversations.svelte')); +}, 30000); + +function makeSession(id: string): ExportedConversation { + return { + conv: { id, currNode: `${id}-msg`, lastModified: 0, name: `Chat ${id}` }, + messages: [ + { + id: `${id}-msg`, + convId: id, + type: MessageType.TEXT, + timestamp: 0, + role: MessageRole.USER, + content: `hello from ${id}`, + parent: null, + children: [] + } + ] + } as unknown as ExportedConversation; +} + +/** + * `parseImportFile` detects the format from the file contents. iOS has no UTI + * for `.jsonl`, so the picker cannot filter on it and the filename carries no + * guarantee: a JSONL export must import under any name. + */ +describe('conversationsStore.parseImportFile', () => { + it('imports a JSONL export whose name has no meaningful extension', async () => { + const jsonl = conversationsStore.serializeSessionToJsonl(makeSession('a')); + + const sessions = await conversationsStore.parseImportFile(new File([jsonl], 'export')); + + expect(sessions).toHaveLength(1); + expect(sessions[0].conv.id).toBe('a'); + expect(sessions[0].messages[0].content).toBe('hello from a'); + }); + + it('imports several sessions from one JSONL file', async () => { + const jsonl = [makeSession('a'), makeSession('b')] + .map((session) => conversationsStore.serializeSessionToJsonl(session)) + .join(NEWLINE); + + const sessions = await conversationsStore.parseImportFile(new File([jsonl], 'export.txt')); + + expect(sessions.map((session) => session.conv.id)).toEqual(['a', 'b']); + }); + + it('imports a ZIP archive whose name has no meaningful extension', async () => { + const zipped = zipSync({ + 'a.jsonl': strToU8(conversationsStore.serializeSessionToJsonl(makeSession('a'))), + 'b.jsonl': strToU8(conversationsStore.serializeSessionToJsonl(makeSession('b'))), + 'notes.txt': strToU8('ignored') + }); + + const sessions = await conversationsStore.parseImportFile(new File([zipped], 'archive')); + + expect(sessions.map((session) => session.conv.id).sort()).toEqual(['a', 'b']); + }); + + it('imports the legacy JSON array format', async () => { + const json = JSON.stringify([makeSession('a')], null, 2); + + const sessions = await conversationsStore.parseImportFile(new File([json], 'export.jsonl')); + + expect(sessions).toHaveLength(1); + expect(sessions[0].conv.id).toBe('a'); + }); + + it('imports the legacy JSON single object format', async () => { + const json = JSON.stringify(makeSession('a')); + + const sessions = await conversationsStore.parseImportFile(new File([json], 'export')); + + expect(sessions).toHaveLength(1); + expect(sessions[0].conv.id).toBe('a'); + }); + + it('rejects a file that holds neither format', async () => { + await expect( + conversationsStore.parseImportFile(new File(['not an export'], 'export.jsonl')) + ).rejects.toThrow(); + }); +}); From 7657a6c26a7c74480db23893c3b5cc68172dacd3 Mon Sep 17 00:00:00 2001 From: Bartowski <3266127+bartowski1182@users.noreply.github.com> Date: Sun, 26 Jul 2026 18:02:56 -0400 Subject: [PATCH 011/190] Keep Minimax's indexer tensors at F32 for speed and accuracy (#26144) * Keep Minimax's indexer tensors at F32 for speed and accuracy * name -> new_name --- conversion/minimax.py | 5 +++++ src/llama-quant.cpp | 4 ++++ 2 files changed, 9 insertions(+) diff --git a/conversion/minimax.py b/conversion/minimax.py index cbbdfe3ae82d..e82e393a3829 100644 --- a/conversion/minimax.py +++ b/conversion/minimax.py @@ -58,6 +58,11 @@ def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None): class MiniMaxM3Model(MiniMaxM2Model): model_arch = gguf.MODEL_ARCH.MINIMAXM3 + def tensor_force_quant(self, name, new_name, bid, n_dims): + if ".indexer." in new_name: + return gguf.GGMLQuantizationType.F32 + return super().tensor_force_quant(name, new_name, bid, n_dims) + def set_gguf_parameters(self): super().set_gguf_parameters() diff --git a/src/llama-quant.cpp b/src/llama-quant.cpp index caf7733a5bff..7c0bac07d096 100644 --- a/src/llama-quant.cpp +++ b/src/llama-quant.cpp @@ -326,6 +326,10 @@ static bool tensor_allows_quantization(const llama_model_quantize_params * param quantize &= name.find("ssm_conv1d") == std::string::npos; quantize &= name.find("shortconv.conv.weight") == std::string::npos; + // do not quantize MiniMax's indexer projection weights, they are tiny + quantize &= name.find("indexer.k_proj.weight") == std::string::npos; + quantize &= name.find("indexer.q_proj.weight") == std::string::npos; + // do not quantize RWKV's small yet 2D weights quantize &= name.find("time_mix_first.weight") == std::string::npos; quantize &= name.find("time_mix_w0.weight") == std::string::npos; From d4d057b6dd6a9df2a44479a68d737154600e6162 Mon Sep 17 00:00:00 2001 From: Piero Evangelista Date: Sun, 26 Jul 2026 18:03:06 -0400 Subject: [PATCH 012/190] ui: fix system message edit box not expanding to fit content (#26006) The in-conversation system-message edit textarea had a fixed min-height and no auto-resize, so long messages were crammed to 2 lines. Reused existing autoResizeTextarea helper on input and when the editor opens, and added max-height to cap growth. --- .../ChatMessageSystem/ChatMessageSystem.svelte | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageSystem/ChatMessageSystem.svelte b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageSystem/ChatMessageSystem.svelte index 36798e2283a4..24b3be4c5ff7 100644 --- a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageSystem/ChatMessageSystem.svelte +++ b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageSystem/ChatMessageSystem.svelte @@ -7,7 +7,7 @@ import { getMessageEditContext } from '$lib/contexts'; import { KeyboardKey, MessageRole } from '$lib/enums'; import { config } from '$lib/stores/settings.svelte'; - import { isIMEComposing } from '$lib/utils'; + import { autoResizeTextarea, isIMEComposing } from '$lib/utils'; interface Props { class?: string; @@ -91,6 +91,11 @@ resizeObserver.disconnect(); }; }); + $effect(() => { + if (editCtx.isEditing && textareaElement) { + autoResizeTextarea(textareaElement); + } + }); function toggleExpand() { isExpanded = !isExpanded; @@ -105,11 +110,15 @@ {#if editCtx.isEditing}
From 0d47ea7427463093e69128bf2c2f9cd06b3ee5b3 Mon Sep 17 00:00:00 2001 From: Xuan-Son Nguyen Date: Mon, 27 Jul 2026 00:22:02 +0200 Subject: [PATCH 013/190] mtmd: fix android build (#26150) --- tools/mtmd/mtmd-helper.cpp | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/tools/mtmd/mtmd-helper.cpp b/tools/mtmd/mtmd-helper.cpp index 84422c89f3ab..90451d02ebd8 100644 --- a/tools/mtmd/mtmd-helper.cpp +++ b/tools/mtmd/mtmd-helper.cpp @@ -640,6 +640,7 @@ bool mtmd_helper_support_video(mtmd_context * ctx) { #ifdef MTMD_VIDEO return mtmd_support_vision(ctx); #else + GGML_UNUSED(ctx); return false; #endif } @@ -1007,6 +1008,9 @@ mtmd_helper_video * mtmd_helper_video_init( return ctx; #else + GGML_UNUSED(mctx); + GGML_UNUSED(path); + GGML_UNUSED(params); LOG_ERR("%s: video is not supported in this build (MTMD_VIDEO is set to OFF)\n", __func__); return nullptr; #endif @@ -1039,6 +1043,10 @@ mtmd_helper_video * mtmd_helper_video_init_from_buf( return ctx; #else + GGML_UNUSED(mctx); + GGML_UNUSED(buf); + GGML_UNUSED(len); + GGML_UNUSED(params); LOG_ERR("%s: video is not supported in this build (MTMD_VIDEO is set to OFF)\n", __func__); return nullptr; #endif @@ -1050,6 +1058,7 @@ void mtmd_helper_video_free(mtmd_helper_video * ctx) { ctx->stop_ffmpeg(); delete ctx; #else + GGML_UNUSED(ctx); LOG_ERR("%s: video is not supported in this build (MTMD_VIDEO is set to OFF)\n", __func__); #endif } @@ -1058,6 +1067,7 @@ mtmd_helper_video_info mtmd_helper_video_get_info(const mtmd_helper_video * ctx) #ifdef MTMD_VIDEO return ctx->info; #else + GGML_UNUSED(ctx); GGML_ASSERT(false && "video is not supported in this build (MTMD_VIDEO is set to OFF)"); #endif } @@ -1068,6 +1078,9 @@ int32_t mtmd_helper_video_read_next(mtmd_helper_video * ctx, if (!ctx) return -2; return ctx->read_next(out_bitmap, out_text); #else + GGML_UNUSED(ctx); + GGML_UNUSED(out_bitmap); + GGML_UNUSED(out_text); GGML_ASSERT(false && "video is not supported in this build (MTMD_VIDEO is set to OFF)"); #endif } From 3d1c3a8975f970a8e5f99ea648733087b52124c5 Mon Sep 17 00:00:00 2001 From: timkhronos Date: Mon, 27 Jul 2026 01:44:41 +0200 Subject: [PATCH 014/190] mtmd: Add Vision Support for Minimax-M3 (#25113) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Add preliminary MiniMax-M3 support Text-only port that re-uses existing components: MiniMax-M2 style GQA with per-head QK-norm and partial rotary, DeepSeek-V3 style leading-dense and routed/shared experts, and swigluoai activation. Sparse attention is not yet supported (dense fallback); vision tower and MTP heads are dropped. * MiniMax-M3 vision tower (mmproj + clip graph) * Delete m3_vision_ref.py * Update clip.cpp * MSA * Update constants.py * Update minimax.py * Cache creation. Working withotu flash attention * Added flash attention for sparse layers * Decomposed slow cpu OP into GPU + CPU ops. Massive speedup over long ctx * Rewrote indexer op to be cuda native. Modified flash attention to match per group block picking * Implement sparse attention calc out of stock ops. * Fix a cache allocation and cont issue * Fixed -fa auto crash, flagged debug spots * Delete vocab.json * Delete model.safetensors.index.json * Delete generation_config.json * Delete Minimax directory * Handled multi stream case to fall back on Dense Attention * Development scaffolding cleanup. No functional change to the decode or 4-way paths. Full debug harness remains at <8136a9c68ed7a5eb009aa67bba3fda8062f4648f> for reproducing the selection-parity validation. * Remove redundant comment from minimax-m3.cpp * Changed 3 Gelu Ops for vision into Gelu_erf ops * Assert that n_kv is multiple of 128 * Rename MSA index tensors to indexer convention Note: All GGUFs generated before this change will need to be regenerated. * Fix incorrect Assert * Review driven changes (#3) * Remove comment from conversion minimax.py Co-authored-by: Sigbjørn Skjæret <1629204+CISC@users.noreply.github.com> * Remove whitespaces from constants.py Co-authored-by: Sigbjørn Skjæret <1629204+CISC@users.noreply.github.com> * Tighten comment in minimax.py Co-authored-by: Sigbjørn Skjæret <1629204+CISC@users.noreply.github.com> * inherit MiniMax-M3 from MiniMax-M2 * drop dead text_config fallbacks * Add indexer writer methods * Reuse LLM_FFN_SWIGLU_OAI_MOE * Remove duplicate indexer setters, add only block_size/local_blocks, follow value naming convention * Fix conversion error /gguf_writer.py Co-authored-by: Sigbjørn Skjæret <1629204+CISC@users.noreply.github.com> * Update gguf-py/gguf/gguf_writer.py Co-authored-by: Sigbjørn Skjæret <1629204+CISC@users.noreply.github.com> * Update gguf-py/gguf/tensor_mapping.py Co-authored-by: Sigbjørn Skjæret <1629204+CISC@users.noreply.github.com> * Update conversion/minimax.py Co-authored-by: Sigbjørn Skjæret <1629204+CISC@users.noreply.github.com> * Update conversion/minimax.py Co-authored-by: Sigbjørn Skjæret <1629204+CISC@users.noreply.github.com> * Remove whitespace in src/llama-kv-cache.cpp Co-authored-by: Sigbjørn Skjæret <1629204+CISC@users.noreply.github.com> * Remove Whitespace in Update src/llama-model.h Co-authored-by: Sigbjørn Skjæret <1629204+CISC@users.noreply.github.com> * Remove whitespace in src/llama-hparams.h Co-authored-by: Sigbjørn Skjæret <1629204+CISC@users.noreply.github.com> * Update minimax_m3.cpp Rewrite code comment based on feedback and to better reflect the actual architecture, and reuse existing build_vit * Rename minimax_m3.cpp to minimax-m3.cpp * Update CMakeLists.txt * Remove debug code from clip.cpp * Update clip.cpp * Update comments in tools/mtmd/models/minimax-m3.cpp * Permute Q/K at conversion, drop precomputed sin/cos * Log cache size on launch, block ctx shift, support prompt caching Log indexer cache size on launch Disallow ctx shift Support prompt caching * Update minimax-m3.cpp * Optimize implementation, add multi stream support. Fully rewrote minimax-m3.cpp for speed and buffer size gains: Unified the 4-way + decode, 1 FA call per layer instead of 4, with the groups mapped onto ne[3] Custom CPU op now emits block-level mask, expanded on GPU, which causes CPU to GPU transfer to shrinks at prefill Decode: ~25 nodes/layer vs ~50, no per-group concats/conts Unified selection semantics, so both regimes rank bs + local bias (position-anchored local force), which means prefill/decode can no longer disagree on selection can_reuse on the MSA bias input. Graph reuse at decode restored (was rebuilding the full graph every token) In-place mask adds, shrinking compute buffer ~6.8 to ~4.2 GiB at ub2048/62k Multi-stream: MSA now runs with -np N when kv_unified=false. Decode stays batched across streams (still 1 FA call), prefill loops per stream. dense fallback only for --kv-unified + multi-seq Measured effect on expert offload bound setup: decode 6.2(4WAY)–7.15(MSA_decode) -> 7.7~7.8 t/s, flat from 5k to 60k+. prefill around 10% faster. buffer about 20% smaller, multi-user support. * set default cache type to F32 * Fix potential DSA double indexer cache allocation bug, only allocate in-cache k_idx for archs that opt in * remove F16 downcasts in MSA attention, force F32 indexer score accum * Add Minimax eos to llama vocab * Guard edge case where idx cache can become stale after a tail trim * Update llama-kv-cache.h * Update llama-kv-cache.cpp * Update llama-kv-cache.cpp * Update llama-kv-cache.h * Change resize Pad to none, resize alg to Bicubic Pillow * Review driven changes * Update llama-kv-cache.cpp * rm unrotated pos_t * fused rope w + pad * rename merge --> merger for consistency * add review skill for mtmd * graph should use hparams n_merge * fix lint --------- Co-authored-by: Daniel Han Co-authored-by: Sigbjørn Skjæret <1629204+CISC@users.noreply.github.com> Co-authored-by: Xuan Son Nguyen --- conversion/__init__.py | 1 + conversion/minimax.py | 77 ++++++++++++++++++++++++++++- gguf-py/gguf/constants.py | 7 +++ gguf-py/gguf/tensor_mapping.py | 8 +++ skills/code-review/SKILL.md | 9 ++++ tools/mtmd/CMakeLists.txt | 1 + tools/mtmd/clip-impl.h | 4 ++ tools/mtmd/clip-model.h | 4 ++ tools/mtmd/clip.cpp | 49 +++++++++++++++++++ tools/mtmd/models/minimax-m3.cpp | 84 ++++++++++++++++++++++++++++++++ tools/mtmd/models/models.h | 6 +++ tools/mtmd/mtmd.cpp | 7 +++ 12 files changed, 256 insertions(+), 1 deletion(-) create mode 100644 tools/mtmd/models/minimax-m3.cpp diff --git a/conversion/__init__.py b/conversion/__init__.py index c5ecc68cfda2..b2bb7e5161eb 100644 --- a/conversion/__init__.py +++ b/conversion/__init__.py @@ -288,6 +288,7 @@ "LlavaForConditionalGeneration": "llava", "MERaLiON2ForConditionalGeneration": "ultravox", "MiMoV2ForCausalLM": "mimo", + "MiniMaxM3SparseForConditionalGeneration": "minimax", "MiniCPMV4_6ForConditionalGeneration": "minicpm", "Mistral3ForConditionalGeneration": "llava", "NemotronH_Nano_VL_V2": "nemotron", diff --git a/conversion/minimax.py b/conversion/minimax.py index e82e393a3829..c2175cc93267 100644 --- a/conversion/minimax.py +++ b/conversion/minimax.py @@ -7,7 +7,7 @@ if TYPE_CHECKING: from torch import Tensor -from .base import ModelBase, TextModel, gguf +from .base import ModelBase, TextModel, MmprojModel, gguf @ModelBase.register("MiniMaxM2ForCausalLM") @@ -92,3 +92,78 @@ def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None): data_torch = data_torch + 1.0 yield from super().modify_tensors(data_torch, name, bid) + + +@ModelBase.register("MiniMaxM3SparseForConditionalGeneration", "MiniMaxM3VLForConditionalGeneration") +class MiniMaxM3VisionModel(MmprojModel): + @classmethod + def filter_tensors(cls, item): + name, gen = item + # keep only the vision-side tensors; text / mtp / sparse-index are dropped + if not name.startswith(("vision_tower.", "multi_modal_projector.", "patch_merge_mlp.")): + return None + return super().filter_tensors((name, gen)) + + def set_gguf_parameters(self): + super().set_gguf_parameters() + assert self.hparams_vision is not None + + self.gguf_writer.add_clip_projector_type(gguf.VisionProjectorType.MINIMAXM3) + self.gguf_writer.add_vision_use_gelu(True) + + # the ViT carries its own LayerNorm eps (text tower uses a different one) + self.gguf_writer.add_vision_attention_layernorm_eps( + self.hparams_vision.get("layer_norm_eps", 1e-5) + ) + + comp = self.hparams_vision.get("img_token_compression_config", {}) + merge_size = comp.get("spatial_merge_size", 2) + self.gguf_writer.add_vision_spatial_merge_size(int(merge_size)) + + def modify_tensors(self, data_torch, name, bid): + assert self.hparams_vision is not None + + # Conv3d patch embed -> Conv2d slices + if name == "vision_tower.vision_model.embeddings.patch_embedding.weight": + if data_torch.ndim != 5: + raise ValueError(f"unexpected patch_embedding rank {data_torch.ndim} for {name}") + kt = data_torch.shape[2] + base = gguf.TENSOR_NAMES[gguf.MODEL_TENSOR.V_ENC_EMBD_PATCH] + for t in range(kt): + suffix = ".weight" if t == 0 else f".weight.{t}" + yield (base + suffix, data_torch[:, :, t, ...]) + return + + # Permute ViT q/k. HF [Ta Ha Wa | Tb Hb Wb | pad] reorder to [Ta Tb | Ha Hb | Wa Wb | pad]. + for new_name, tensor in super().modify_tensors(data_torch, name, bid): + if ".attn_q." in new_name or ".attn_k." in new_name: + tensor = self._permute_vit_qk(tensor, new_name) + yield new_name, tensor + + def _permute_vit_qk(self, t: "Tensor", new_name: str) -> "Tensor": + assert self.hparams_vision is not None + n_head = self.hparams_vision["num_attention_heads"] + d_head = t.shape[0] // n_head + axis_dim = 2 * ((2 * (d_head // 2) // 3) // 2) + ah = axis_dim // 2 + half = 3 * ah + perm = [] + perm += list(range(0, ah)) + perm += list(range(half, half + ah)) + perm += list(range(ah, 2 * ah)) + perm += list(range(half + ah, half + 2 * ah)) + perm += list(range(2 * ah, 3 * ah)) + perm += list(range(half + 2 * ah, half + 3 * ah)) + perm += list(range(2 * half, d_head)) + + assert axis_dim % 2 == 0 + assert 3 * axis_dim <= d_head + assert len(perm) == d_head + assert sorted(perm) == list(range(d_head)), "perm is not a bijection of d_head" + assert t.shape[0] == n_head * d_head, f"{new_name}: {t.shape[0]} != {n_head}*{d_head}" + assert d_head == 80 + + idx = torch.tensor(perm, dtype=torch.long) + if t.ndim == 2: + return t.reshape(n_head, d_head, t.shape[1])[:, idx, :].reshape(t.shape) + return t.reshape(n_head, d_head)[:, idx].reshape(t.shape) diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index 66d50cca2684..2071e3eaa8a4 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -857,6 +857,8 @@ class MODEL_TENSOR(IntEnum): V_MM_UP = auto() # cogvlm V_MM_DOWN = auto() # cogvlm V_MM_GATE = auto() # cogvlm + V_MM_MERGER_FC1 = auto() # minimax-m3 (patch-merge MLP) + V_MM_MERGER_FC2 = auto() # minimax-m3 (patch-merge MLP) V_TOK_BOI = auto() # cogvlm V_TOK_EOI = auto() # cogvlm V_TOK_IMG_BEGIN = auto() # hunyuanvl @@ -1441,6 +1443,8 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.V_MM_UP: "mm.up", MODEL_TENSOR.V_MM_DOWN: "mm.down", MODEL_TENSOR.V_MM_GATE: "mm.gate", + MODEL_TENSOR.V_MM_MERGER_FC1: "mm.merger.fc1", + MODEL_TENSOR.V_MM_MERGER_FC2: "mm.merger.fc2", MODEL_TENSOR.V_TOK_BOI: "v.boi", MODEL_TENSOR.V_TOK_EOI: "v.eoi", MODEL_TENSOR.V_MM_PRE_NORM: "mm.pre_norm", @@ -1637,6 +1641,8 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.V_RESMPL_QUERY, MODEL_TENSOR.V_TOK_EMBD_IMG_BREAK, MODEL_TENSOR.V_MM_PATCH_MERGER, + MODEL_TENSOR.V_MM_MERGER_FC1, + MODEL_TENSOR.V_MM_MERGER_FC2, MODEL_TENSOR.V_DS_NORM, MODEL_TENSOR.V_DS_FC1, MODEL_TENSOR.V_DS_FC2, @@ -4771,6 +4777,7 @@ class VisionProjectorType: YOUTUVL = "youtuvl" NEMOTRON_V2_VL = "nemotron_v2_vl" HUNYUANVL = "hunyuanvl" + MINIMAXM3 = "minimax_m3" MINICPMV4_6 = "minicpmv4_6" GRANITE_SPEECH = "granite_speech" # audio MIMOVL = "mimovl" diff --git a/gguf-py/gguf/tensor_mapping.py b/gguf-py/gguf/tensor_mapping.py index 59623accfdb6..62d7a827e35c 100644 --- a/gguf-py/gguf/tensor_mapping.py +++ b/gguf-py/gguf/tensor_mapping.py @@ -1838,6 +1838,14 @@ class TensorNameMap: "visual.downsample", # glm4v ), + MODEL_TENSOR.V_MM_MERGER_FC1: ( + "patch_merge_mlp.linear_1", # minimax-m3 + ), + + MODEL_TENSOR.V_MM_MERGER_FC2: ( + "patch_merge_mlp.linear_2", # minimax-m3 + ), + MODEL_TENSOR.V_DS_NORM: ( "model.visual.deepstack_merger_list.{bid}.norm", # deepstack in qwen3vl ), diff --git a/skills/code-review/SKILL.md b/skills/code-review/SKILL.md index 84075fea3252..ba76c481150e 100644 --- a/skills/code-review/SKILL.md +++ b/skills/code-review/SKILL.md @@ -110,6 +110,15 @@ Public API changes carry a higher bar than internal ones (`CONTRIBUTING.md`). Re - Security: don't trust client-supplied headers (e.g. `X-Forwarded-For`) or add footguns; things like IP allowlisting belong at a reverse proxy unless there's a trusted-proxy design. - Wire new behavior into the existing request/response and checkpoint paths correctly; watch for resource leaks across requests. +## Multimodal (`tools/mtmd/`) + +- Tensor names must be prefixed by `v.`, `a.`, `mm.` or `a.mm.` (legacy naming doesn't follow this convention - this is expected, but new code should follow it). +- Do not use explicit sin/cos for RoPE; use `ggml_rope_ext` instead, see `HOWTO-add-model.md`. If it can't express the needed behavior, that's a design discussion, not a PR. +- New GGML ops must not be introduced in the same PR, you must push it as a separate PR. +- In most cases, `build_vit` should be enough to build the transformer graph for vision models. Do not add a loop to build the transformer graph manually, unless you have a very good reason to do so. If you do, please explain why in the PR description. +- If you need a dedicated preprocessor, there is a high chance that it can be a derived class from one of the existing preprocessors. Check carefully before adding a new preprocessor class. +- If the model need a new public API in `mtmd.h`, open a discussion first. + ## General (always) Enforce the `AGENTS.md` / `CONTRIBUTING.md` coding and naming guidelines on every changed line - this is a distinct pass from checking that the code works, and matters just as much for review speed: diff --git a/tools/mtmd/CMakeLists.txt b/tools/mtmd/CMakeLists.txt index d0329ca56743..fd7ddceb0bf0 100644 --- a/tools/mtmd/CMakeLists.txt +++ b/tools/mtmd/CMakeLists.txt @@ -47,6 +47,7 @@ add_library(mtmd models/paddleocr.cpp models/pixtral.cpp models/qwen2vl.cpp + models/minimax-m3.cpp models/qwen3vl.cpp models/mimovl.cpp models/qwen3a.cpp diff --git a/tools/mtmd/clip-impl.h b/tools/mtmd/clip-impl.h index 5b413681f040..42374311ce7b 100644 --- a/tools/mtmd/clip-impl.h +++ b/tools/mtmd/clip-impl.h @@ -131,6 +131,8 @@ #define TN_MM_SOFT_EMB_N "mm.soft_emb_norm.weight" // gemma3 #define TN_MM_PROJECTOR "mm.model.fc.%s" // idefics3, deepseekocr #define TN_MM_PATCH_MERGER "mm.patch_merger.%s" // mistral small 3.1, glm4v +#define TN_MM_MERGER_FC1 "mm.merger.fc1.%s" // minimax-m3 patch-merge MLP +#define TN_MM_MERGER_FC2 "mm.merger.fc2.%s" #define TN_TOK_IMG_BREAK "v.token_embd.img_break" // pixtral #define TN_TOK_GLM_BOI "adapter.boi" // glm-edge (these embeddings are not in text model) #define TN_TOK_GLM_EOI "adapter.eoi" // glm-edge (these embeddings are not in text model) @@ -370,6 +372,7 @@ enum projector_type { PROJECTOR_TYPE_MINICPMV4_6, PROJECTOR_TYPE_GRANITE_SPEECH, PROJECTOR_TYPE_MIMOVL, + PROJECTOR_TYPE_MINIMAX_M3, PROJECTOR_TYPE_GRANITE4_VISION, PROJECTOR_TYPE_UNKNOWN, }; @@ -424,6 +427,7 @@ static std::map PROJECTOR_TYPE_NAMES = { { PROJECTOR_TYPE_MINICPMV4_6, "minicpmv4_6"}, { PROJECTOR_TYPE_GRANITE_SPEECH, "granite_speech"}, { PROJECTOR_TYPE_MIMOVL, "mimovl"}, + { PROJECTOR_TYPE_MINIMAX_M3, "minimax_m3"}, { PROJECTOR_TYPE_GRANITE4_VISION, "granite4_vision"}, }; diff --git a/tools/mtmd/clip-model.h b/tools/mtmd/clip-model.h index 6d4336c4010b..850957d7de1c 100644 --- a/tools/mtmd/clip-model.h +++ b/tools/mtmd/clip-model.h @@ -397,6 +397,10 @@ struct clip_model { ggml_tensor * mm_0_b = nullptr; ggml_tensor * mm_2_w = nullptr; ggml_tensor * mm_2_b = nullptr; + ggml_tensor * mm_merger_fc1_w = nullptr; // minimax-m3 + ggml_tensor * mm_merger_fc1_b = nullptr; + ggml_tensor * mm_merger_fc2_w = nullptr; + ggml_tensor * mm_merger_fc2_b = nullptr; ggml_tensor * image_newline = nullptr; ggml_tensor * view_seperator = nullptr; diff --git a/tools/mtmd/clip.cpp b/tools/mtmd/clip.cpp index b8866506493e..e0e2107a0be3 100644 --- a/tools/mtmd/clip.cpp +++ b/tools/mtmd/clip.cpp @@ -915,6 +915,10 @@ static std::unique_ptr clip_get_graph_builder(clip_ctx * ctx, const { builder = std::make_unique(ctx, img); } break; + case PROJECTOR_TYPE_MINIMAX_M3: + { + builder = std::make_unique(ctx, img); + } break; case PROJECTOR_TYPE_STEP3VL: { builder = std::make_unique(ctx, img); @@ -1469,6 +1473,17 @@ struct clip_model_loader { LOG_WRN("%s: more info: https://github.com/ggml-org/llama.cpp/issues/16842\n\n", __func__); } } break; + case PROJECTOR_TYPE_MINIMAX_M3: + { + hparams.n_merge = 2; // spatial_merge_size + hparams.image_resize_algo = RESIZE_ALGO_BICUBIC_PILLOW; + hparams.image_resize_pad = PAD_NONE; + get_u32(KEY_SPATIAL_MERGE_SIZE, hparams.n_merge, false); + hparams.rope_theta = 10000.0f; // vision_config.rope_theta + // MiniMax-M3: max_pixels 451584 (=672^2) -> 576 merged tokens (image_seq_length) + hparams.set_limit_image_tokens(8, 576); + hparams.set_warmup_n_tokens(16*16); + } break; case PROJECTOR_TYPE_MIMOVL: { hparams.n_merge = 2; // spatial_merge_size @@ -2089,6 +2104,19 @@ struct clip_model_loader { model.mm_1_w = get_tensor(string_format(TN_LLAVA_PROJ, 2, "weight")); model.mm_1_b = get_tensor(string_format(TN_LLAVA_PROJ, 2, "bias"), false); } break; + case PROJECTOR_TYPE_MINIMAX_M3: + { + // per-patch MLP: mm.1 -> gelu -> mm.2 + model.mm_1_w = get_tensor(string_format(TN_LLAVA_PROJ, 1, "weight")); + model.mm_1_b = get_tensor(string_format(TN_LLAVA_PROJ, 1, "bias")); + model.mm_2_w = get_tensor(string_format(TN_LLAVA_PROJ, 2, "weight")); + model.mm_2_b = get_tensor(string_format(TN_LLAVA_PROJ, 2, "bias")); + // 2x2 merge MLP: mm.merge.fc1 -> gelu -> mm.merge.fc2 + model.mm_merger_fc1_w = get_tensor(string_format(TN_MM_MERGER_FC1, "weight")); + model.mm_merger_fc1_b = get_tensor(string_format(TN_MM_MERGER_FC1, "bias")); + model.mm_merger_fc2_w = get_tensor(string_format(TN_MM_MERGER_FC2, "weight")); + model.mm_merger_fc2_b = get_tensor(string_format(TN_MM_MERGER_FC2, "bias")); + } break; case PROJECTOR_TYPE_STEP3VL: { model.mm_0_w = get_tensor(string_format(TN_LLAVA_PROJ, 0, "weight")); @@ -3360,6 +3388,7 @@ int clip_n_output_tokens(const clip_ctx * ctx, const clip_image_f32 * img) { case PROJECTOR_TYPE_QWEN3VL: case PROJECTOR_TYPE_EXAONE4_5: case PROJECTOR_TYPE_MIMOVL: + case PROJECTOR_TYPE_MINIMAX_M3: case PROJECTOR_TYPE_GLM4V: case PROJECTOR_TYPE_YOUTUVL: { @@ -3866,6 +3895,24 @@ bool clip_image_batch_encode(clip_ctx * ctx, int n_threads, const clip_image_f32 set_input_i32("positions", positions); } break; + case PROJECTOR_TYPE_MINIMAX_M3: + { + const int n_merge = hparams.n_merge; + const int gh = image_size_height / patch_size; + const int gw = image_size_width / patch_size; + std::vector pos_h, pos_w; + pos_h.reserve(gh * gw); + pos_w.reserve(gh * gw); + for (int bh = 0; bh < gh / n_merge; bh++) + for (int bw = 0; bw < gw / n_merge; bw++) + for (int mh = 0; mh < n_merge; mh++) + for (int mw = 0; mw < n_merge; mw++) { + pos_h.push_back(bh * n_merge + mh); + pos_w.push_back(bw * n_merge + mw); + } + set_input_i32("minimax_pos_h", pos_h); + set_input_i32("minimax_pos_w", pos_w); + } break; case PROJECTOR_TYPE_DOTS_OCR: { const int pw = image_size_width / patch_size; @@ -4569,6 +4616,8 @@ int clip_n_mmproj_embd(const struct clip_ctx * ctx) { return ctx->model.mm_ffn_down_w->ne[1]; case PROJECTOR_TYPE_GLM_EDGE: return ctx->model.mm_model_mlp_3_w->ne[1]; + case PROJECTOR_TYPE_MINIMAX_M3: + return ctx->model.mm_merger_fc2_b->ne[0]; case PROJECTOR_TYPE_QWEN2VL: case PROJECTOR_TYPE_QWEN25VL: case PROJECTOR_TYPE_EXAONE4_5: diff --git a/tools/mtmd/models/minimax-m3.cpp b/tools/mtmd/models/minimax-m3.cpp new file mode 100644 index 000000000000..447621754e69 --- /dev/null +++ b/tools/mtmd/models/minimax-m3.cpp @@ -0,0 +1,84 @@ +#include "models.h" + +ggml_tensor * clip_graph_minimax_m3::apply_rope( + ggml_tensor * x, ggml_tensor * pos_h, ggml_tensor * pos_w) { + const int64_t Hn = x->ne[1]; + const int64_t P = x->ne[2]; + const size_t es = ggml_element_size(x); + const int dh = (int) x->ne[0]; + const int axd = 2 * ((2 * (dh / 2) / 3) / 2); + + GGML_ASSERT(x->nb[0] == es); + GGML_ASSERT(3 * axd <= dh); + + const float th = hparams.rope_theta; + + // layout of x is [t, h, w, pad] + // t is unrotated, h and w are rotated, pad is unrotated + // note: everything from n_dims onward untouched, so w and pad are rotated in one call. + auto sl = [&](int off, int n) { + return ggml_cont(ctx0, ggml_view_3d(ctx0, x, n, Hn, P, x->nb[1], x->nb[2], (size_t) off * es)); + }; + ggml_tensor * t = sl(0, axd); + ggml_tensor * h = sl(axd, axd); + ggml_tensor * w = sl(2 * axd, dh - 2 * axd); // w + pad + + h = ggml_rope_ext(ctx0, h, pos_h, nullptr, axd, GGML_ROPE_TYPE_NEOX, 0, th, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f); + w = ggml_rope_ext(ctx0, w, pos_w, nullptr, axd, GGML_ROPE_TYPE_NEOX, 0, th, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f); + return ggml_concat(ctx0, ggml_concat(ctx0, t, h, 0), w, 0); +} + +ggml_cgraph * clip_graph_minimax_m3::build() { + GGML_ASSERT(model.patch_bias == nullptr); + GGML_ASSERT(model.class_embedding == nullptr); + GGML_ASSERT(model.patch_embeddings_0 && model.patch_embeddings_1); + GGML_ASSERT(model.mm_1_w && model.mm_2_w); + GGML_ASSERT(model.mm_merger_fc1_w && model.mm_merger_fc2_w); + + const int batch_size = 1; + const int n_pos = n_patches; + const int merge = hparams.n_merge; + + // patch embedding + ggml_tensor * inp_raw = build_inp_raw(); + ggml_tensor * inp = ggml_add(ctx0, + ggml_conv_2d(ctx0, model.patch_embeddings_0, inp_raw, patch_size, patch_size, 0, 0, 1, 1), + ggml_conv_2d(ctx0, model.patch_embeddings_1, inp_raw, patch_size, patch_size, 0, 0, 1, 1)); + + // spatial merge + { + inp = ggml_permute(ctx0, inp, 1, 2, 0, 3); + inp = ggml_cont_4d(ctx0, inp, n_embd * merge, n_patches_x / merge, n_patches_y, batch_size); + inp = ggml_reshape_4d(ctx0, inp, n_embd * merge, n_patches_x / merge, merge, batch_size * (n_patches_y / merge)); + inp = ggml_permute(ctx0, inp, 0, 2, 1, 3); + inp = ggml_cont_3d(ctx0, inp, n_embd, n_patches_x * n_patches_y, batch_size); + } + + // t (time axis) is always 0 for now, so we leave it unrotated + ggml_tensor * pos_h = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_pos); + ggml_set_name(pos_h, "minimax_pos_h"); ggml_set_input(pos_h); + ggml_tensor * pos_w = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_pos); + ggml_set_name(pos_w, "minimax_pos_w"); ggml_set_input(pos_w); + + ggml_tensor * inpL = build_vit( + inp, n_pos, NORM_TYPE_NORMAL, FFN_GELU_ERF, nullptr, + [&](ggml_tensor * c, const clip_layer &) { + return apply_rope(c, pos_h, pos_w); + }); + + // projector + ggml_tensor * emb = inpL; + emb = build_ffn(emb, model.mm_1_w, model.mm_1_b, + nullptr, nullptr, + model.mm_2_w, model.mm_2_b, FFN_GELU_ERF, -1); + + const int64_t proj = emb->ne[0]; + emb = ggml_reshape_2d(ctx0, emb, proj * merge * merge, n_pos / (merge * merge)); + + emb = build_ffn(emb, model.mm_merger_fc1_w, model.mm_merger_fc1_b, + nullptr, nullptr, + model.mm_merger_fc2_w, model.mm_merger_fc2_b, FFN_GELU_ERF, -1); + + ggml_build_forward_expand(gf, emb); + return gf; +} diff --git a/tools/mtmd/models/models.h b/tools/mtmd/models/models.h index 5f1493fa603e..2d7555da41d2 100644 --- a/tools/mtmd/models/models.h +++ b/tools/mtmd/models/models.h @@ -40,6 +40,12 @@ struct clip_graph_qwen3vl : clip_graph_qwen2vl { ggml_cgraph * build() override; }; +struct clip_graph_minimax_m3 : clip_graph { + clip_graph_minimax_m3(clip_ctx * ctx, const clip_image_f32 & img) : clip_graph(ctx, img) {} + ggml_cgraph * build() override; + ggml_tensor * apply_rope(ggml_tensor * x, ggml_tensor * pos_h, ggml_tensor * pos_w); +}; + struct clip_graph_mimovl : clip_graph { clip_graph_mimovl(clip_ctx * ctx, const clip_image_f32 & img) : clip_graph(ctx, img) {} ggml_cgraph * build() override; diff --git a/tools/mtmd/mtmd.cpp b/tools/mtmd/mtmd.cpp index 5915b4cba967..bb49b211efb3 100644 --- a/tools/mtmd/mtmd.cpp +++ b/tools/mtmd/mtmd.cpp @@ -463,6 +463,13 @@ struct mtmd_context { img_end = "<|vision_end|>"; image_preproc = std::make_unique(ctx_v); } break; + case PROJECTOR_TYPE_MINIMAX_M3: + { + // ]<]start of image[>[ ... (image embeddings) ... ]<]end of image[>[ + img_beg = "]<]start of image[>["; + img_end = "]<]end of image[>["; + image_preproc = std::make_unique(ctx_v); + } break; case PROJECTOR_TYPE_YOUTUVL: { // <|vision_start|> ... (image embeddings) ... <|vision_end|> From 88b47a755c72fed4b22fba0fd262e2d7b7d01583 Mon Sep 17 00:00:00 2001 From: rankaiyx Date: Mon, 27 Jul 2026 08:30:22 +0800 Subject: [PATCH 015/190] ui: Fix symbolic math tool JS sandbox prompt (#26131) * Update sandbox.ts * Update sandbox.ts * Update sandbox.ts * Update sandbox.ts * Revise nerdamer description in sandbox constants Updated NERDAMER_DESCRIPTION to clarify usage and warnings. --- tools/ui/src/lib/constants/sandbox.ts | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/tools/ui/src/lib/constants/sandbox.ts b/tools/ui/src/lib/constants/sandbox.ts index 58242678d98b..381621de647e 100644 --- a/tools/ui/src/lib/constants/sandbox.ts +++ b/tools/ui/src/lib/constants/sandbox.ts @@ -14,12 +14,15 @@ export const SANDBOX_EMPTY_OUTPUT = '(no output)'; export const SANDBOX_TRUNCATION_NOTICE = '[output truncated]'; const NERDAMER_DESCRIPTION = ` -Symbolic/numeric math via \`nerdamer\` (pre-loaded, do not require, use it directly). -nerdamer('diff(sin(x)/x,x)') or nerdamer.diff('sin(x)/x','x') → Expression; convert with .toString()/.text()/.toTeX(), or .evaluate() (→ still Expression, then .toString()). -nerdamer(expr,{x:2}) substitutes only; chain .evaluate() or pass 'numer' for numeric result. -solve(expr,var)→Symbol[]; solveEquations([eq1,..])→[[var,val],..] pairs. -Functions: simplify/expand/factor(expr), diff(expr,var[,n]), integrate(expr,var), defint(expr,from,to,var), limit(expr,var,to), laplace(expr,t,s), ilt(expr,s,t), gcd/lcm(a,b), roots/coeffs/partfrac(expr,var), pfactor(n), numer/decimals/erf(expr), product/sum(expr,var,from,to), mean/median/stdev/variance(...vals). -Object.keys(nerdamer).filter(k=>typeof nerdamer[k]==='function') lists all available functions. If you need a function not documented above, list them first — do not guess function names.`; +Symbolic/numeric math via \`nerdamer\` +nerdamer(expr,subs?,opts?)/nerdamer.func(...)→Expression Format via .text(fmt?) (fmt: 'decimals'|'fractions'|'scientific') eval via .evaluate(subs?) +nerdamer(expr,{x:2}) substitutes numeric via opts 'numer' or .evaluate() +simplify/expand/factor(expr) div/gcd/lcm(...) coeffs/partfrac(expr,var) +diff/integrate(expr,var) defint(expr,lo,hi,var?) sum/product(expr,var,lo,hi) limit(expr,var,pt) +solve(expr,var) solveEquations([eq1,eq2],[var1,var2]) +polarform/rectform/arg/realpart/imagpart(z) +set/get Var/Constant(name,val?) setFunction(name,[params],body) +IMPORTANT:Identifier 'nerdamer' has already been declared, use it directly`; /** * Build the sandbox tool definition. When `includeSymbolicMath` is true, From d73c1d6b22a2d3ecc74c2c9cde354015ee72e862 Mon Sep 17 00:00:00 2001 From: Pascal Date: Mon, 27 Jul 2026 07:34:47 +0200 Subject: [PATCH 016/190] server + ui: fix stream routes for model names containing a slash (#26137) * server + ui: refactor resumable stream routes to query string conv_id The conversation id can embed a model name containing slashes (ggml-org/...) in router mode, which the decoded path splits before the :conv_id param is captured, so stop and resume never matched the session. Move the id to the conv_id query string on the public routes and on the internal router -> child hop, where slashes survive encoding. Handlers are unchanged since query and path params land in the same map. Add a regression test with a slashed model name. * server: move stream route docs to server-stream.h Address review: ngxson wants the main server.cpp registration code kept clean and simple, with route-level explanations living in the header. Move the query string rationale and the lookup ownership note next to the handler declarations in server-stream.h, and shorten the wiring comment to a pointer. * server: cancel a pending request when its stream is stopped during model load The conversation was registered in the conv map only after the blocking autoload wait, so a stop issued while the model loaded found nothing to cancel and the request went on to generate an orphan once the load ended. Register the conversation before the wait and give the entry a ticket: a stop erases the entry, and the parked request checks its ticket after the wait and aborts with 400 instead of starting. A newer request on the same conversation replaces the entry, so only the stopped request is cancelled. Add a regression test that stops during the load window. * server + ui: resume a stream after a page reload during model load A pending request died with the client socket when the page was reloaded while its model was loading, so no session ever existed and the conversation had nothing to recover. A session request that waited for a load now detaches from the client socket and reaches the child regardless, the session buffer receives the generation, and the resume route answers 503 while the owner is loading so the client retries instead of dropping its state. The WebUI persists the pending stream at send time, quietly polls on 503, and attaches once the session exists. Add a regression test that drops the client during the load window. * ui: show the model load progress again after a page refresh The resume wait was invisible, so a conversation refreshed while its model was loading showed nothing until the first byte. On a 503 from the resume probe, mark the conversation as loading again so the assistant row persisted at send time renders the processing info, and target the model frozen in the persisted stream state for the progress, since the row has no model yet and the dropdown may not be restored. * fix CI * fix CI bis --- tools/server/README-dev.md | 10 +- tools/server/server-models.cpp | 67 +++++--- tools/server/server-models.h | 31 +++- tools/server/server-stream.cpp | 8 +- tools/server/server-stream.h | 6 + tools/server/server.cpp | 13 +- tools/server/tests/unit/test_stream.py | 153 ++++++++++++++++++ .../ChatMessageAssistant.svelte | 9 +- tools/ui/src/lib/constants/api-endpoints.ts | 6 +- tools/ui/src/lib/services/chat.service.ts | 32 +++- tools/ui/src/lib/stores/chat.svelte.ts | 65 +++++++- 11 files changed, 345 insertions(+), 55 deletions(-) create mode 100644 tools/server/tests/unit/test_stream.py diff --git a/tools/server/README-dev.md b/tools/server/README-dev.md index b4ec9f17d3c6..b41d70c63ac8 100644 --- a/tools/server/README-dev.md +++ b/tools/server/README-dev.md @@ -136,13 +136,13 @@ Producer side: `server_res_generator` extends `server_res_spipe`, which keeps al Lifetime safety: the session holds no back reference to the response, so `spipe` is a plain `unique_ptr` touched only by the http worker. `cancel` raises an atomic the producer polls; the producer finalizes the session from its destructor, which also runs `~server_response_reader::stop()` to cancel the generation at the queue level. A `DELETE` stops work by raising the flag and letting the worker unwind. -Consumer side: `GET /v1/stream/?from=N` opens a `text/event-stream` that replays buffered bytes from offset `N` and blocks for live bytes, so the browser reattaches like a fresh EventSource. An offset below the dropped prefix returns 400. +Consumer side: `GET /v1/stream?conv_id=&from=N` opens a `text/event-stream` that replays buffered bytes from offset `N` and blocks for live bytes, so the browser reattaches like a fresh EventSource. An offset below the dropped prefix returns 400. Routes: -- `GET /v1/stream/:conv_id?from=N`: replay or live reattach. +- `GET /v1/stream?conv_id=&from=N`: replay or live reattach. The id travels in the query string because it can embed a model name containing slashes. - `POST /v1/streams/lookup` with `{"conversation_ids": [...]}`: returns session status only for ids the caller already owns. There is no listing route, so live sessions cannot be enumerated (an earlier `GET /v1/streams` was removed for exactly this reason). -- `DELETE /v1/stream/:conv_id`: explicit Stop, idempotent (`evict_and_cancel`). +- `DELETE /v1/stream?conv_id=`: explicit Stop, idempotent (`evict_and_cancel`). Router mode binds the same paths to proxy handlers. A `conv_id -> child` map (`conv_models`), populated when a POST is routed, resolves the owning child in one lookup with no polling. The lookup groups ids per child; GET and DELETE proxy straight to the owner. This loopback REST hop is expected to move to a websocket IPC later, swapping only the transport. @@ -166,8 +166,8 @@ graph TD GC[GC thread] -- drop after TTL --> Sess end Sess -- read_from offset --> Cons[stream_pipe_consumer] - Cons -- "GET /v1/stream/:id?from=N" --> Client - DEL[DELETE /v1/stream/:id] -- evict_and_cancel --> Sess + Cons -- "GET /v1/stream?conv_id=id&from=N" --> Client + DEL[DELETE /v1/stream?conv_id=id] -- evict_and_cancel --> Sess ``` The diagram shows the buffer touch points. The live wire (chunks streamed to the original client during a normal generation) is the producer's default output, described under "Producer side" above. diff --git a/tools/server/server-models.cpp b/tools/server/server-models.cpp index 923b3533e9cb..188a72a3741c 100644 --- a/tools/server/server-models.cpp +++ b/tools/server/server-models.cpp @@ -1172,7 +1172,7 @@ bool server_models::ensure_model_ready(const std::string & name) { return true; } -server_http_res_ptr server_models::proxy_request(const server_http_req & req, const std::string & method, const std::string & name, bool update_last_used) { +server_http_res_ptr server_models::proxy_request(const server_http_req & req, const std::string & method, const std::string & name, bool update_last_used, bool detached) { auto meta = get_meta(name); if (!meta.has_value()) { throw std::runtime_error("model name=" + name + " is not found"); @@ -1198,7 +1198,10 @@ server_http_res_ptr server_models::proxy_request(const server_http_req & req, co req.headers, req.body, req.files, - req.should_stop, + // a detached request belongs to a replay session that outlives the client socket: + // it reaches the child even when the downstream died during the load wait, the + // session buffer is the recipient and DELETE remains the stop + detached ? std::function([]() { return false; }) : req.should_stop, base_params.timeout_read, base_params.timeout_write ); @@ -1469,13 +1472,9 @@ static bool router_validate_model(std::string & name, server_models & models, bo } // resolve alias to canonical model name name = meta->name; - if (models_autoload) { - models.ensure_model_ready(name); - } else { - if (!meta->is_running()) { - res_err(res, format_error_response("model is not loaded", ERROR_TYPE_INVALID_REQUEST)); - return false; - } + if (!models_autoload && !meta->is_running()) { + res_err(res, format_error_response("model is not loaded", ERROR_TYPE_INVALID_REQUEST)); + return false; } return true; } @@ -1568,6 +1567,9 @@ void server_models_routes::init_routes() { if (!router_validate_model(name, models, autoload, error_res)) { return error_res; } + if (autoload) { + models.ensure_model_ready(name); + } return models.proxy_request(req, method, name, false); }; @@ -1581,12 +1583,23 @@ void server_models_routes::init_routes() { return error_res; } // remember which child serves this conversation so the stream routes can route straight - // to it without polling, keyed on the exact conv id from the header + // to it without polling, keyed on the exact conv id from the header. registered before + // the load wait so a stop issued while the model loads can erase the entry and cancel + // this request instead of leaving an orphan generation std::string conv_id = server_stream_conv_id_from_headers(req.headers); - if (!conv_id.empty()) { - models.conv_models.remember(conv_id, name); + uint64_t ticket = models.conv_models.remember(conv_id, name); + bool waited = autoload && models.ensure_model_ready(name); + if (ticket != 0 && !models.conv_models.alive(conv_id, ticket)) { + SRV_INF("request for conv_id=%s cancelled while model name=%s was loading\n", + conv_id.c_str(), name.c_str()); + res_err(error_res, format_error_response( + "request cancelled by a stop while the model was loading", ERROR_TYPE_INVALID_REQUEST)); + return error_res; } - return models.proxy_request(req, method, name, true); // update last usage for POST request only + // a session request that waited for a load detaches from the client socket: the + // client may have dropped during the wait (page reload) and the session buffer must + // still receive the generation for a later resume + return models.proxy_request(req, method, name, true, waited && ticket != 0); // update last usage for POST request only }; this->post_router_models_load = [this](const server_http_req & req) { @@ -1779,7 +1792,7 @@ void server_models_routes::init_routes() { }; this->router_stream_get = [this](const server_http_req & req) { - // GET /v1/stream/?from=N. resolve the owning child from the conv_id -> model + // GET /v1/stream?conv_id=&from=N. resolve the owning child from the conv_id -> model // map, 404 when nothing maps auto res = std::make_unique(); std::string conv_id = req.get_param("conv_id"); @@ -1789,13 +1802,24 @@ void server_models_routes::init_routes() { } std::optional owner = resolve_child_for_conv(models, conv_id); if (!owner.has_value()) { - res_err(res, format_error_response("Stream not found or expired", ERROR_TYPE_NOT_FOUND)); + // a registered conv whose model is still loading earns a retry: the session appears + // once the load ends and the pending request reaches the child + auto tracked = models.conv_models.lookup(conv_id); + auto meta = tracked.has_value() ? models.get_meta(*tracked) : std::nullopt; + bool transient = meta.has_value() && (meta->status == SERVER_MODEL_STATUS_LOADING || + meta->status == SERVER_MODEL_STATUS_DOWNLOADING || + meta->status == SERVER_MODEL_STATUS_DOWNLOADED); + if (transient) { + res_err(res, format_error_response("Stream owner model is loading, retry later", ERROR_TYPE_UNAVAILABLE)); + } else { + res_err(res, format_error_response("Stream not found or expired", ERROR_TYPE_NOT_FOUND)); + } return res; } std::string from = req.get_param("from"); - std::string child_path = "/v1/stream/" + encode_qs(conv_id); + std::string child_path = "/v1/stream?conv_id=" + encode_qs(conv_id); if (!from.empty()) { - child_path += "?from=" + from; + child_path += "&from=" + from; } SRV_TRC("proxying stream resume to model %s on port %d, path=%s\n", owner->name.c_str(), owner->port, child_path.c_str()); @@ -1875,7 +1899,7 @@ void server_models_routes::init_routes() { }; this->router_stream_delete = [this](const server_http_req & req) { - // DELETE /v1/stream/. resolve the owning child via the map and forward only to + // DELETE /v1/stream?conv_id=. resolve the owning child via the map and forward only to // it, evict_and_cancel is idempotent on the child auto res = std::make_unique(); std::string conv_id = req.get_param("conv_id"); @@ -1883,7 +1907,7 @@ void server_models_routes::init_routes() { res_err(res, format_error_response("Missing conversation id in path", ERROR_TYPE_INVALID_REQUEST)); return res; } - std::string child_path = "/v1/stream/" + encode_qs(conv_id); + std::string child_path = "/v1/stream?conv_id=" + encode_qs(conv_id); auto owner = resolve_child_for_conv(models, conv_id); if (owner.has_value()) { httplib::Client cli(CHILD_ADDR, owner->port); @@ -1892,6 +1916,11 @@ void server_models_routes::init_routes() { cli.set_write_timeout(0, STREAM_LOOKUP_TIMEOUT_MS * 1000); auto resp = cli.Delete(child_path.c_str()); (void) resp; // the child logs its own miss when the session is unknown there + } else if (auto tracked = models.conv_models.lookup(conv_id); tracked.has_value()) { + // the entry exists but its model is still loading: the forget below erases it, + // which cancels the request parked in proxy_post before the generation starts + SRV_INF("router stop for conv_id=%s while model name=%s is loading, cancelling the pending request\n", + conv_id.c_str(), tracked->c_str()); } else { SRV_WRN("router stop for unknown conv_id=%s, no owning child in the conv map\n", conv_id.c_str()); diff --git a/tools/server/server-models.h b/tools/server/server-models.h index 62bed8725b5b..614798186cfc 100644 --- a/tools/server/server-models.h +++ b/tools/server/server-models.h @@ -134,12 +134,24 @@ struct server_models { // proxy_request forwards a POST carrying an X-Conversation-Id. best effort: a stale entry just // makes the child answer not found and the client recovers. owns its lock, one mutex per struct struct conv_model_tracker { - void remember(const std::string & conv_id, const std::string & model) { + // returns the ticket of this registration, 0 when nothing was registered. erasing or + // replacing the entry invalidates the ticket, which is how a stop cancels a request + // parked in the model load wait + uint64_t remember(const std::string & conv_id, const std::string & model) { if (conv_id.empty() || model.empty()) { - return; + return 0; } std::lock_guard lock(mu); - map[conv_id] = model; + uint64_t ticket = next_ticket++; + map[conv_id] = { model, ticket }; + return ticket; + } + + // false means a stop erased the entry or a newer request replaced it + bool alive(const std::string & conv_id, uint64_t ticket) { + std::lock_guard lock(mu); + auto it = map.find(conv_id); + return it != map.end() && it->second.ticket == ticket; } std::optional lookup(const std::string & conv_id) { @@ -151,7 +163,7 @@ struct server_models { if (it == map.end()) { return std::nullopt; } - return it->second; + return it->second.model; } void forget(const std::string & conv_id) { @@ -163,8 +175,13 @@ struct server_models { } private: - std::mutex mu; - std::unordered_map map; + struct entry_t { + std::string model; + uint64_t ticket; + }; + std::mutex mu; + uint64_t next_ticket = 1; + std::unordered_map map; }; common_preset_context ctx_preset; @@ -249,7 +266,7 @@ struct server_models { bool ensure_model_ready(const std::string & name); // proxy an HTTP request to the model instance - server_http_res_ptr proxy_request(const server_http_req & req, const std::string & method, const std::string & name, bool update_last_used); + server_http_res_ptr proxy_request(const server_http_req & req, const std::string & method, const std::string & name, bool update_last_used, bool detached = false); // handle message sent from server_child::notify_to_router() // raw input must starts with CMD_CHILD_TO_ROUTER_STATE, followed by a JSON string diff --git a/tools/server/server-stream.cpp b/tools/server/server-stream.cpp index f0a35b18e525..f6b9b8a9f4cc 100644 --- a/tools/server/server-stream.cpp +++ b/tools/server/server-stream.cpp @@ -453,7 +453,7 @@ static server_http_res_ptr make_error_response(int status, const std::string & m server_http_context::handler_t server_stream_make_get_handler() { return [](const server_http_req & req) -> server_http_res_ptr { - // GET /v1/stream/?from=N replays buffered SSE bytes then blocks for live + // GET /v1/stream?conv_id=&from=N replays buffered SSE bytes then blocks for live // bytes until the session finalizes, streamed as text/event-stream for EventSource std::string conv_id = req.get_param("conv_id"); if (conv_id.empty()) { @@ -560,13 +560,13 @@ server_http_context::handler_t server_stream_make_lookup_handler() { server_http_context::handler_t server_stream_make_delete_handler() { return [](const server_http_req & req) -> server_http_res_ptr { - // DELETE /v1/stream/ is the explicit user Stop, cancels the producer and evicts + // DELETE /v1/stream?conv_id= is the explicit user Stop, cancels the producer and evicts // the buffer. idempotent, returns 204 even if the session was already gone std::string conv_id = req.get_param("conv_id"); if (conv_id.empty()) { return make_error_response(400, "Missing conversation id in path", ERROR_TYPE_INVALID_REQUEST); } - SRV_TRC("DELETE /v1/stream/%s -> evict_and_cancel\n", conv_id.c_str()); + SRV_TRC("DELETE /v1/stream conv_id=%s -> evict_and_cancel\n", conv_id.c_str()); g_stream_sessions.evict_and_cancel(conv_id); auto res = std::make_unique(); res->status = 204; @@ -621,7 +621,7 @@ bool server_res_spipe::conn_alive() { bool server_res_spipe::should_stop() { if (spipe) { - // note: if DELETE /v1/stream/ is called, is_cancelled() will be true + // note: if DELETE /v1/stream is called for this conv, is_cancelled() will be true return spipe->is_cancelled(); } else { return !conn_alive(); diff --git a/tools/server/server-stream.h b/tools/server/server-stream.h index 9753140dd601..1e7461285f43 100644 --- a/tools/server/server-stream.h +++ b/tools/server/server-stream.h @@ -45,7 +45,13 @@ void server_stream_session_manager_start(); void server_stream_session_manager_stop(); // route handler factories wired under /v1/stream/* by server.cpp +// child-side handlers for the resumable stream routes. the conv id travels in the conv_id +// query string because it can embed a model name containing slashes (org/repo), which the +// decoded path would split before the param is captured server_http_context::handler_t server_stream_make_get_handler(); +// POST /v1/streams/lookup with body {"conversation_ids": [...]}: only answers for ids the +// caller already owns (the WebUI passes the convs visible in its sidebar), the server never +// lists ids it has not been asked about, so a random caller cannot enumerate live sessions server_http_context::handler_t server_stream_make_lookup_handler(); server_http_context::handler_t server_stream_make_delete_handler(); diff --git a/tools/server/server.cpp b/tools/server/server.cpp index b6fef99e8747..a3b2a8b0fe1b 100644 --- a/tools/server/server.cpp +++ b/tools/server/server.cpp @@ -272,10 +272,8 @@ int llama_server(common_params & params, int argc, char ** argv) { ctx_http.get ("/slots", ex_wrapper(routes.get_slots)); ctx_http.post("/slots/:id_slot", ex_wrapper(routes.post_slots)); - // resumable streaming, the conversation_id is the session identity end to end. router and - // child wire different handlers under the same paths: a child binds the local session - // factories, the router binds proxies that resolve the owning child through the - // conv_id -> model map + // resumable streaming: a child binds the local session factories, the router binds + // proxies that resolve the owning child, see server-stream.h server_http_context::handler_t stream_get_h; server_http_context::handler_t streams_lookup_h; server_http_context::handler_t stream_delete_h; @@ -288,12 +286,9 @@ int llama_server(common_params & params, int argc, char ** argv) { streams_lookup_h = server_stream_make_lookup_handler(); stream_delete_h = server_stream_make_delete_handler(); } - ctx_http.get ("/v1/stream/:conv_id", ex_wrapper(stream_get_h)); - // POST /v1/streams/lookup with body {"conversation_ids": [...]}. you can only ask for ids - // you already own (the WebUI passes the convs visible in its sidebar). the server never - // lists ids it has not been asked about, so a random caller cannot enumerate live sessions + ctx_http.get ("/v1/stream", ex_wrapper(stream_get_h)); ctx_http.post("/v1/streams/lookup", ex_wrapper(streams_lookup_h)); - ctx_http.del ("/v1/stream/:conv_id", ex_wrapper(stream_delete_h)); + ctx_http.del ("/v1/stream", ex_wrapper(stream_delete_h)); // Google Cloud Platform (Vertex AI) compat ctx_http.register_gcp_compat(); diff --git a/tools/server/tests/unit/test_stream.py b/tools/server/tests/unit/test_stream.py new file mode 100644 index 000000000000..a1ef55567bc7 --- /dev/null +++ b/tools/server/tests/unit/test_stream.py @@ -0,0 +1,153 @@ +import json +import socket +import threading +import time +from urllib.parse import quote +import pytest +from utils import * + +server: ServerProcess + +# a model name with slashes exercises the query string routing of the stream routes: the id +# cannot travel as a path param because the decoded slash would split it before capture +MODEL = "ggml-org/tinygemma3-GGUF:Q8_0" +STREAM_ID = f"conv-stream-test::{MODEL}" +QS = "conv_id=" + quote(STREAM_ID, safe="") + + +@pytest.fixture(autouse=True) +def create_server(): + global server + server = ServerPreset.router() + + +def test_stream_resume_and_stop_with_slashed_model_name(): + global server + server.start() + + content = "" + for data in server.make_stream_request("POST", "/chat/completions", data={ + "model": MODEL, + "stream": True, + "max_tokens": 16, + "messages": [{"role": "user", "content": "hello"}], + }, headers={"X-Conversation-Id": STREAM_ID}): + if data["choices"]: + content += data["choices"][0]["delta"].get("content") or "" + assert len(content) > 0 + + # the finished session replays from the beginning through the router + res = server.make_request("GET", f"/v1/stream?{QS}&from=0") + assert res.status_code == 200 + assert "data: " in str(res.body) + + # the explicit stop reaches the owning child and evicts the session + res = server.make_request("DELETE", f"/v1/stream?{QS}") + assert res.status_code == 204 + res = server.make_request("GET", f"/v1/stream?{QS}&from=0") + assert res.status_code == 404 + + +def test_stream_stop_during_model_load(): + global server + server.start() + + thread_error: list[ServerError] = [] + thread_done = threading.Event() + + def fire_post(): + try: + for _ in server.make_stream_request("POST", "/chat/completions", data={ + "model": MODEL, + "stream": True, + "max_tokens": 512, + "messages": [{"role": "user", "content": "Count from 1 to 1000."}], + }, headers={"X-Conversation-Id": STREAM_ID}): + pass + except ServerError as e: + thread_error.append(e) + finally: + thread_done.set() + + t = threading.Thread(target=fire_post) + t.start() + + # catch the autoload window, tiny models load fast so poll aggressively + saw_loading = False + deadline = time.time() + 5.0 + while time.time() < deadline and not thread_done.is_set(): + res = server.make_request("GET", "/models") + status = next(m["status"]["value"] for m in res.body["data"] if m["id"] == MODEL) + if status == "loading": + saw_loading = True + break + time.sleep(0.002) + if not saw_loading: + t.join() + pytest.skip("load window too short to be observed on this machine") # ty: ignore[too-many-positional-arguments] + + # a stop during the load cancels the parked request instead of leaving an orphan + res = server.make_request("DELETE", f"/v1/stream?{QS}") + assert res.status_code == 204 + assert thread_done.wait(timeout=60) + t.join() + assert len(thread_error) == 1 + assert thread_error[0].code == 400 + assert "cancelled" in json.dumps(thread_error[0].body) + res = server.make_request("GET", f"/v1/stream?{QS}&from=0") + assert res.status_code == 404 + + +def test_stream_resumes_after_reload_during_model_load(): + global server + server.start() + + # raw socket client so the connection can be dropped mid load like a page reload + body = json.dumps({ + "model": MODEL, + "stream": True, + "max_tokens": 16, + "messages": [{"role": "user", "content": "hello"}], + }) + request = ( + f"POST /v1/chat/completions HTTP/1.1\r\n" + f"Host: {server.server_host}:{server.server_port}\r\n" + f"Content-Type: application/json\r\n" + f"X-Conversation-Id: {STREAM_ID}\r\n" + f"Content-Length: {len(body)}\r\n" + f"Connection: close\r\n\r\n{body}" + ) + sock = socket.create_connection((server.server_host, server.server_port)) + sock.sendall(request.encode()) + + # drop the client while the model loads, poll aggressively to catch the window + saw_loading = False + saw_503 = False + deadline = time.time() + 5.0 + while time.time() < deadline: + res = server.make_request("GET", "/models") + status = next(m["status"]["value"] for m in res.body["data"] if m["id"] == MODEL) + if status == "loading": + saw_loading = True + break + if status == "loaded": + break + time.sleep(0.002) + sock.close() + if not saw_loading: + pytest.skip("load window too short to be observed on this machine") # ty: ignore[too-many-positional-arguments] + + # while the model loads the resume route answers retry later, then the session appears, + # receives the whole generation despite the dead client, and replays from the beginning + deadline = time.time() + 60.0 + replay = None + while time.time() < deadline: + res = server.make_request("GET", f"/v1/stream?{QS}&from=0") + if res.status_code == 503: + saw_503 = True + elif res.status_code == 200 and "data: " in str(res.body): + replay = res + break + time.sleep(0.1) + assert saw_503, "resume during the load did not answer 503" + assert replay is not None, "session never became resumable after the client disconnect" diff --git a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageAssistant/ChatMessageAssistant.svelte b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageAssistant/ChatMessageAssistant.svelte index 00578fcf1a1a..199d75fcec95 100644 --- a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageAssistant/ChatMessageAssistant.svelte +++ b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageAssistant/ChatMessageAssistant.svelte @@ -10,7 +10,7 @@ } from '$lib/components/app'; import { getMessageEditContext } from '$lib/contexts'; import { useProcessingState } from '$lib/hooks/use-processing-state.svelte'; - import { isLoading, isChatStreaming } from '$lib/stores/chat.svelte'; + import { chatStore, isLoading, isChatStreaming } from '$lib/stores/chat.svelte'; import { modelLoadProgressText } from '$lib/utils'; import { MessageRole } from '$lib/enums'; import { config } from '$lib/stores/settings.svelte'; @@ -82,8 +82,11 @@ let hasNoContent = $derived(!message?.content?.trim()); let isActivelyProcessing = $derived(isCurrentlyLoading || isStreaming); - // during a router auto-load the message has no model yet, so target the selected one - let loadTargetModel = $derived(message.model ?? modelsStore.selectedModelName); + // during a router auto-load the message has no model yet: target the model frozen in the + // persisted stream state (survives a reload), then fall back to the dropdown selection + let loadTargetModel = $derived( + message.model ?? chatStore.getResumeModel(message.convId) ?? modelsStore.selectedModelName + ); let modelLoadProgress = $derived( isRouter && loadTargetModel ? modelsStore.getLoadProgress(loadTargetModel) : null ); diff --git a/tools/ui/src/lib/constants/api-endpoints.ts b/tools/ui/src/lib/constants/api-endpoints.ts index 37137c1c7699..ab35708a46da 100644 --- a/tools/ui/src/lib/constants/api-endpoints.ts +++ b/tools/ui/src/lib/constants/api-endpoints.ts @@ -21,7 +21,11 @@ export const API_TOOLS = { EXECUTE: '/tools' }; -// resumable stream routes, the conv::model identity is appended as a path segment +// resumable stream routes, the conv::model identity travels as the conv_id query param +// because model names can contain slashes that a path segment cannot carry +// resume retry cadence while the owning model is still loading (server answers 503) +export const STREAM_RESUME_RETRY_MS = 2000; + export const API_STREAM = { BASE: './v1/stream', LOOKUP: './v1/streams/lookup' diff --git a/tools/ui/src/lib/services/chat.service.ts b/tools/ui/src/lib/services/chat.service.ts index d2455614fbaa..4ce396533de8 100644 --- a/tools/ui/src/lib/services/chat.service.ts +++ b/tools/ui/src/lib/services/chat.service.ts @@ -343,6 +343,9 @@ export class ChatService { // model the ::model suffix keeps the per model session distinct if (stream && conversationId) { headers['X-Conversation-Id'] = streamIdentity(conversationId, options.model); + // persist the pending stream before the fetch: a reload during the model load or + // the prompt processing must still find its way back to the session once it exists + ChatService.saveStreamState(conversationId, 0, options.model ?? null); } const response = await fetch(API_CHAT.COMPLETIONS, { @@ -353,6 +356,11 @@ export class ChatService { }); if (!response.ok) { + // a rejected request (including one cancelled by a stop during the model load) + // leaves nothing to resume + if (conversationId) { + ChatService.clearStreamState(conversationId); + } const error = await ChatService.parseErrorResponse(response); if (onError) { @@ -512,7 +520,7 @@ export class ChatService { if (!conversationId) return; try { const id = streamIdentity(conversationId, model); - await fetch(`${API_STREAM.BASE}/${encodeURIComponent(id)}`, { + await fetch(`${API_STREAM.BASE}?conv_id=${encodeURIComponent(id)}`, { method: 'DELETE', headers: getAuthHeaders() }); @@ -605,6 +613,26 @@ export class ChatService { * existing SSE parser drains it like a fresh stream. The server returns 200 on success, 404 if * no session exists for the conv_id, and 400 if the offset is below the dropped prefix. */ + // probe the resume route status without consuming the stream: the SSE route has no HEAD, + // so issue the GET and abort it right after the status line. 0 on network error + static async probeResumeStatus(streamId: string): Promise { + if (!streamId) return 0; + const ac = new AbortController(); + try { + const resp = await fetch( + `${API_STREAM.BASE}?conv_id=${encodeURIComponent(streamId)}&from=0`, + { + headers: getAuthHeaders(), + signal: ac.signal + } + ); + ac.abort(); + return resp.status; + } catch { + return 0; + } + } + static async resumeStream( conversationId: string, signal?: AbortSignal, @@ -614,7 +642,7 @@ export class ChatService { const state = ChatService.getStreamState(conversationId); const from = state?.bytesReceived ?? 0; const id = streamIdentity(conversationId, model); - const url = `${API_STREAM.BASE}/${encodeURIComponent(id)}?from=${from}`; + const url = `${API_STREAM.BASE}?conv_id=${encodeURIComponent(id)}&from=${from}`; return await fetch(url, { method: 'GET', signal, headers: getAuthHeaders() }); } diff --git a/tools/ui/src/lib/stores/chat.svelte.ts b/tools/ui/src/lib/stores/chat.svelte.ts index 5cbfe213b104..222723ab108d 100644 --- a/tools/ui/src/lib/stores/chat.svelte.ts +++ b/tools/ui/src/lib/stores/chat.svelte.ts @@ -14,6 +14,7 @@ import { SvelteMap, SvelteSet } from 'svelte/reactivity'; import { DatabaseService } from '$lib/services/database.service'; import { ChatService } from '$lib/services/chat.service'; +import { STREAM_RESUME_RETRY_MS } from '$lib/constants/api-endpoints'; import { streamIdentity } from '$lib/utils/stream-identity'; import { getAuthHeaders } from '$lib/utils/api-headers'; import { CONTENT_TYPE_HEADER } from '$lib/constants'; @@ -78,7 +79,7 @@ class ChatStore { // true while the active conversation streams reasoning content but no visible content yet isReasoning = $state(false); // resumable stream connection state for the active conversation - // streaming -> bytes flowing normally, resuming -> waiting on /v1/stream/:id reconnect, lost -> unrecoverable + // streaming -> bytes flowing normally, resuming -> waiting on /v1/stream reconnect, lost -> unrecoverable streamConnectionState = $state(StreamConnectionState.STREAMING); chatLoadingStates = new SvelteMap(); chatReasoningStates = new SvelteMap(); @@ -94,6 +95,11 @@ class ChatStore { // off when one conv finishes while another is still streaming. mirrors chatLoadingStates // in scope but tracks the attach + tee replay path specifically private attachingConvs = new SvelteSet(); + // pending resume retry timers while an owning model loads, one per conv + private resumeRetryTimers = new SvelteMap>(); + // convs whose resume waits on a model load: their loading state belongs to the retry loop, + // so discoverActiveStream must not treat it as a live send and bail + private resumePendingConvs = new SvelteSet(); // in-flight discoverActiveStream guard, keyed by conv id private discoveringConvs = new SvelteSet(); private abortControllers = new SvelteMap(); @@ -263,7 +269,7 @@ class ChatStore { const id = streamId || streamIdentity(convId, selectedModelName()); let response: Response; try { - response = await fetch(`./v1/stream/${encodeURIComponent(id)}?from=0`, { + response = await fetch(`./v1/stream?conv_id=${encodeURIComponent(id)}&from=0`, { headers: getAuthHeaders() }); } catch (e) { @@ -438,13 +444,22 @@ class ChatStore { } } + /** + * Model frozen at send time for a stream awaiting resume, from the persisted stream state. + * The load progress indicator targets it after a reload, when the message row has no model + * yet and the dropdown selection may not be restored. + */ + getResumeModel(convId: string): string | null { + return ChatService.getStreamState(convId)?.model ?? null; + } + async discoverActiveStream(convId: string): Promise { if (!convId) return; if (this.chatStreamingStates.has(convId)) return; - if (this.chatLoadingStates.get(convId)) return; + if (this.chatLoadingStates.get(convId) && !this.resumePendingConvs.has(convId)) return; // concurrency guard: another discover may already be running for this conv (typical race // between mount and visibilitychange on tab switch). a second concurrent fetch on the same - // /v1/stream/ would duplicate every byte into the DB message, this guard bounces it + // /v1/stream would duplicate every byte into the DB message, this guard bounces it if (this.discoveringConvs.has(convId)) return; this.discoveringConvs.add(convId); @@ -470,6 +485,38 @@ class ChatStore { if (!localState) { return; } + // quiet status probe first: a full attach flips the loading UI on every try, probing + // keeps the retry loop invisible while the owning model is still loading (503) + const status = await ChatService.probeResumeStatus(streamId); + if (status === 503) { + // make the wait visible: the empty assistant row persisted at send time renders + // the processing info, whose model load percentage flows from the models feed + this.resumePendingConvs.add(convId); + this.setChatLoading(convId, true); + if (!this.resumeRetryTimers.has(convId)) { + this.resumeRetryTimers.set( + convId, + setTimeout(() => { + this.resumeRetryTimers.delete(convId); + void this.discoverActiveStream(convId); + }, STREAM_RESUME_RETRY_MS) + ); + } + return; + } + if (this.resumePendingConvs.delete(convId) && status !== 200) { + // the wait is over without a session to attach, drop the visible loading state + this.setChatLoading(convId, false); + } + if (status === 0) { + // transient network failure, the next mount or visibility change retries + return; + } + if (status !== 200) { + // the session is gone (stopped, TTL expired), nothing to resume anymore + ChatService.clearStreamState(convId); + return; + } await this.attachServerStream(convId, streamId); // if attachServerStream failed (session gone, TTL expired), clear the local state to avoid retrying forever if (!this.chatStreamingStates.has(convId) && !this.chatLoadingStates.get(convId)) { @@ -1469,8 +1516,16 @@ class ChatStore { // detached drain keeps producing tokens until eos or max_tokens. use the frozen identity // captured when the session started, not the live dropdown const streamStateForStop = this.chatStreamingStates.get(convId); - const modelForStop = streamStateForStop?.model; + const modelForStop = streamStateForStop?.model ?? ChatService.getStreamState(convId)?.model; void ChatService.cancelServerStream(convId, modelForStop); + // an explicit stop leaves nothing to resume and kills a pending resume retry + ChatService.clearStreamState(convId); + const retryTimer = this.resumeRetryTimers.get(convId); + if (retryTimer !== undefined) { + clearTimeout(retryTimer); + this.resumeRetryTimers.delete(convId); + } + this.resumePendingConvs.delete(convId); this.abortRequest(convId); this.setChatLoading(convId, false); this.clearChatStreaming(convId); From ad256ded30b5e9dbf43c146b452673a1471b62cd Mon Sep 17 00:00:00 2001 From: Aaron Teo Date: Mon, 27 Jul 2026 16:44:08 +0800 Subject: [PATCH 017/190] args: add `-lm mlock` where it mlocks but doesnt mmap (#26135) * arg: add `-lm mlock` where it mlocks but doesnt mmap Signed-off-by: Aaron Teo * docs: rm unwanted docs changes Signed-off-by: Aaron Teo * docs: revert auto-formatting Signed-off-by: Aaron Teo * bench: fix automated review point 3 Signed-off-by: Aaron Teo * arg: revert the meaning of --mlock to non-mmap'ed mlock Signed-off-by: Aaron Teo * docs: update docs Signed-off-by: Aaron Teo * docs: remove extra changes from `llama-gen-docs` Signed-off-by: Aaron Teo --------- Signed-off-by: Aaron Teo --- common/arg.cpp | 14 +++--- include/llama.h | 9 ++-- src/llama-model-loader.cpp | 2 +- src/llama-model.cpp | 2 +- src/llama.cpp | 11 +++-- tests/test-arg-parser.cpp | 9 ++++ tools/cli/README.md | 4 +- tools/completion/README.md | 4 +- tools/llama-bench/llama-bench.cpp | 78 ++++++++++++++++--------------- tools/server/README.md | 4 +- 10 files changed, 77 insertions(+), 60 deletions(-) diff --git a/common/arg.cpp b/common/arg.cpp index 9753441313a7..84c19e06329a 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -2508,7 +2508,7 @@ common_params_context common_params_parser_init(common_params & params, llama_ex } add_opt(common_arg( {"--mlock"}, - "DEPRECATED in favor of `--load-mode`: mmap + force system to keep model in RAM rather than swapping or compressing", + "DEPRECATED in favor of `--load-mode`: force system to keep model in RAM rather than swapping or compressing", [](common_params & params) { LOG_WRN("DEPRECATED: --mlock is deprecated. use --load-mode mlock instead\n"); params.load_mode = LLAMA_LOAD_MODE_MLOCK; @@ -2537,13 +2537,15 @@ common_params_context common_params_parser_init(common_params & params, llama_ex "model loading mode (default: mmap)\n" "- none: no special loading mode\n" "- mmap: memory-map model (if mmap disabled, slower load but may reduce pageouts if not using mlock)\n" - "- mlock: mmap + force system to keep model in RAM rather than swapping or compressing\n" + "- mlock: force system to keep model in RAM rather than swapping or compressing\n" + "- mmap+mlock: mmap + force system to keep model in RAM rather than swapping or compressing\n" "- dio: use DirectIO if available\n", [](common_params & params, const std::string & value) { - /**/ if (value == "none") { params.load_mode = LLAMA_LOAD_MODE_NONE; } - else if (value == "mmap") { params.load_mode = LLAMA_LOAD_MODE_MMAP; } - else if (value == "mlock") { params.load_mode = LLAMA_LOAD_MODE_MLOCK; } - else if (value == "dio") { params.load_mode = LLAMA_LOAD_MODE_DIRECT_IO; } + /**/ if (value == "none") { params.load_mode = LLAMA_LOAD_MODE_NONE; } + else if (value == "mmap") { params.load_mode = LLAMA_LOAD_MODE_MMAP; } + else if (value == "mlock") { params.load_mode = LLAMA_LOAD_MODE_MLOCK; } + else if (value == "mmap+mlock") { params.load_mode = LLAMA_LOAD_MODE_MMAP_MLOCK; } + else if (value == "dio") { params.load_mode = LLAMA_LOAD_MODE_DIRECT_IO; } else { throw std::invalid_argument("invalid value"); } } ).set_env("LLAMA_ARG_LOAD_MODE")); diff --git a/include/llama.h b/include/llama.h index 9fab69317006..3c6d22be8999 100644 --- a/include/llama.h +++ b/include/llama.h @@ -203,10 +203,11 @@ extern "C" { }; enum llama_load_mode { - LLAMA_LOAD_MODE_NONE = 0, // no special loading mode - LLAMA_LOAD_MODE_MMAP = 1, // memory map the model - LLAMA_LOAD_MODE_MLOCK = 2, // mmap + force system to keep model in RAM rather than swapping or compressing - LLAMA_LOAD_MODE_DIRECT_IO = 3, // use direct I/O if available + LLAMA_LOAD_MODE_NONE = 0, // no special loading mode + LLAMA_LOAD_MODE_MMAP = 1, // memory map the model + LLAMA_LOAD_MODE_MLOCK = 2, // force system to keep model in RAM rather than swapping or compressing + LLAMA_LOAD_MODE_MMAP_MLOCK = 3, // mmap + force system to keep model in RAM rather than swapping or compressing + LLAMA_LOAD_MODE_DIRECT_IO = 4, // use direct I/O if available }; LLAMA_API const char * llama_load_mode_name(enum llama_load_mode load_mode); diff --git a/src/llama-model-loader.cpp b/src/llama-model-loader.cpp index 43447f57d30b..510586e96c20 100644 --- a/src/llama-model-loader.cpp +++ b/src/llama-model-loader.cpp @@ -542,7 +542,7 @@ llama_model_loader::llama_model_loader( tensor_buft_overrides = param_tensor_buft_overrides_p; - this->use_mmap = load_mode == LLAMA_LOAD_MODE_MMAP || load_mode == LLAMA_LOAD_MODE_MLOCK; + this->use_mmap = load_mode == LLAMA_LOAD_MODE_MMAP || load_mode == LLAMA_LOAD_MODE_MMAP_MLOCK; this->use_direct_io = load_mode == LLAMA_LOAD_MODE_DIRECT_IO; if (!fname.empty()) { diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 51796921081f..074acbe1fa27 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -1249,7 +1249,7 @@ void llama_model_base::load_vocab(llama_model_loader & ml) { bool llama_model_base::load_tensors(llama_model_loader & ml) { const auto & split_mode = params.split_mode; - const bool use_mlock = params.load_mode == LLAMA_LOAD_MODE_MLOCK; + const bool use_mlock = params.load_mode == LLAMA_LOAD_MODE_MLOCK || params.load_mode == LLAMA_LOAD_MODE_MMAP_MLOCK; const auto & tensor_split = params.tensor_split; const int n_layer_all = hparams.n_layer_all; diff --git a/src/llama.cpp b/src/llama.cpp index 11ac9656d9f9..d22e4c81a9e3 100644 --- a/src/llama.cpp +++ b/src/llama.cpp @@ -54,6 +54,8 @@ const char * llama_load_mode_name(enum llama_load_mode load_mode) { return "mmap"; case LLAMA_LOAD_MODE_MLOCK: return "mlock"; + case LLAMA_LOAD_MODE_MMAP_MLOCK: + return "mmap+mlock"; case LLAMA_LOAD_MODE_DIRECT_IO: return "dio"; } @@ -61,10 +63,11 @@ const char * llama_load_mode_name(enum llama_load_mode load_mode) { } enum llama_load_mode llama_load_mode_from_str(const char * str) { - if (std::strcmp(str, "none") == 0) { return LLAMA_LOAD_MODE_NONE; } - if (std::strcmp(str, "mmap") == 0) { return LLAMA_LOAD_MODE_MMAP; } - if (std::strcmp(str, "mlock") == 0) { return LLAMA_LOAD_MODE_MLOCK; } - if (std::strcmp(str, "dio") == 0) { return LLAMA_LOAD_MODE_DIRECT_IO; } + if (std::strcmp(str, "none") == 0) { return LLAMA_LOAD_MODE_NONE; } + if (std::strcmp(str, "mmap") == 0) { return LLAMA_LOAD_MODE_MMAP; } + if (std::strcmp(str, "mlock") == 0) { return LLAMA_LOAD_MODE_MLOCK; } + if (std::strcmp(str, "mmap+mlock") == 0) { return LLAMA_LOAD_MODE_MMAP_MLOCK; } + if (std::strcmp(str, "dio") == 0) { return LLAMA_LOAD_MODE_DIRECT_IO; } throw std::invalid_argument(std::string("unknown load mode: ") + str); } diff --git a/tests/test-arg-parser.cpp b/tests/test-arg-parser.cpp index 000ecd9aaa76..1d3584f903c4 100644 --- a/tests/test-arg-parser.cpp +++ b/tests/test-arg-parser.cpp @@ -143,6 +143,10 @@ static void test(void) { assert(true == common_params_parse(argv.size(), list_str_to_char(argv).data(), params, LLAMA_EXAMPLE_COMMON)); assert(params.load_mode == LLAMA_LOAD_MODE_MLOCK); + argv = {"binary_name", "-lm", "mmap+mlock"}; + assert(true == common_params_parse(argv.size(), list_str_to_char(argv).data(), params, LLAMA_EXAMPLE_COMMON)); + assert(params.load_mode == LLAMA_LOAD_MODE_MMAP_MLOCK); + argv = {"binary_name", "-lm", "dio"}; assert(true == common_params_parse(argv.size(), list_str_to_char(argv).data(), params, LLAMA_EXAMPLE_COMMON)); assert(params.load_mode == LLAMA_LOAD_MODE_DIRECT_IO); @@ -187,6 +191,11 @@ static void test(void) { assert(true == common_params_parse(argv.size(), list_str_to_char(argv).data(), params, LLAMA_EXAMPLE_COMMON)); assert(params.load_mode == LLAMA_LOAD_MODE_MLOCK); + setenv("LLAMA_ARG_LOAD_MODE", "mmap+mlock", true); + argv = {"binary_name"}; + assert(true == common_params_parse(argv.size(), list_str_to_char(argv).data(), params, LLAMA_EXAMPLE_COMMON)); + assert(params.load_mode == LLAMA_LOAD_MODE_MMAP_MLOCK); + setenv("LLAMA_ARG_LOAD_MODE", "dio", true); argv = {"binary_name"}; assert(true == common_params_parse(argv.size(), list_str_to_char(argv).data(), params, LLAMA_EXAMPLE_COMMON)); diff --git a/tools/cli/README.md b/tools/cli/README.md index 6ee447b07301..972ea04dc7d7 100644 --- a/tools/cli/README.md +++ b/tools/cli/README.md @@ -55,10 +55,10 @@ | `-dt, --defrag-thold N` | KV cache defragmentation threshold (DEPRECATED)
(env: LLAMA_ARG_DEFRAG_THOLD) | | `-np, --parallel N` | number of parallel sequences to decode (default: 1)
(env: LLAMA_ARG_N_PARALLEL) | | `--rpc SERVERS` | comma-separated list of RPC servers (host:port)
(env: LLAMA_ARG_RPC) | -| `--mlock` | DEPRECATED in favor of `--load-mode`: mmap + force system to keep model in RAM rather than swapping or compressing
(env: LLAMA_ARG_MLOCK) | +| `--mlock` | DEPRECATED in favor of `--load-mode`: force system to keep model in RAM rather than swapping or compressing
(env: LLAMA_ARG_MLOCK) | | `--mmap, --no-mmap` | DEPRECATED in favor of `--load-mode`: whether to memory-map model. (if mmap disabled, slower load but may reduce pageouts if not using mlock)
(env: LLAMA_ARG_MMAP) | | `-dio, --direct-io, -ndio, --no-direct-io` | DEPRECATED in favor of `--load-mode`: use DirectIO if available
(env: LLAMA_ARG_DIO) | -| `-lm, --load-mode MODE` | model loading mode (default: mmap)
- none: no special loading mode
- mmap: memory-map model (if mmap disabled, slower load but may reduce pageouts if not using mlock)
- mlock: mmap + force system to keep model in RAM rather than swapping or compressing
- dio: use DirectIO if available

(env: LLAMA_ARG_LOAD_MODE) | +| `-lm, --load-mode MODE` | model loading mode (default: mmap)
- none: no special loading mode
- mmap: memory-map model (if mmap disabled, slower load but may reduce pageouts if not using mlock)
- mlock: force system to keep model in RAM rather than swapping or compressing
- mmap+mlock: mmap + force system to keep model in RAM rather than swapping or compressing
- dio: use DirectIO if available

(env: LLAMA_ARG_LOAD_MODE) | | `--numa TYPE` | attempt optimizations that help on some NUMA systems
- distribute: spread execution evenly over all nodes
- isolate: only spawn threads on CPUs on the node that execution started on
- numactl: use the CPU map provided by numactl
if run without this previously, it is recommended to drop the system page cache before using this
see https://github.com/ggml-org/llama.cpp/issues/1437
(env: LLAMA_ARG_NUMA) | | `-dev, --device ` | comma-separated list of devices to use for offloading (none = don't offload)
use --list-devices to see a list of available devices
(env: LLAMA_ARG_DEVICE) | | `--list-devices` | print list of available devices and exit | diff --git a/tools/completion/README.md b/tools/completion/README.md index 17f7cd765900..bce71d68d949 100644 --- a/tools/completion/README.md +++ b/tools/completion/README.md @@ -138,10 +138,10 @@ llama-completion.exe -m models\gemma-1.1-7b-it.Q4_K_M.gguf --ignore-eos -n -1 | `-dt, --defrag-thold N` | KV cache defragmentation threshold (DEPRECATED)
(env: LLAMA_ARG_DEFRAG_THOLD) | | `-np, --parallel N` | number of parallel sequences to decode (default: 1)
(env: LLAMA_ARG_N_PARALLEL) | | `--rpc SERVERS` | comma-separated list of RPC servers (host:port)
(env: LLAMA_ARG_RPC) | -| `--mlock` | DEPRECATED in favor of `--load-mode`: mmap + force system to keep model in RAM rather than swapping or compressing
(env: LLAMA_ARG_MLOCK) | +| `--mlock` | DEPRECATED in favor of `--load-mode`: force system to keep model in RAM rather than swapping or compressing
(env: LLAMA_ARG_MLOCK) | | `--mmap, --no-mmap` | DEPRECATED in favor of `--load-mode`: whether to memory-map model. (if mmap disabled, slower load but may reduce pageouts if not using mlock)
(env: LLAMA_ARG_MMAP) | | `-dio, --direct-io, -ndio, --no-direct-io` | DEPRECATED in favor of `--load-mode`: use DirectIO if available
(env: LLAMA_ARG_DIO) | -| `-lm, --load-mode MODE` | model loading mode (default: mmap)
- none: no special loading mode
- mmap: memory-map model (if mmap disabled, slower load but may reduce pageouts if not using mlock)
- mlock: mmap + force system to keep model in RAM rather than swapping or compressing
- dio: use DirectIO if available

(env: LLAMA_ARG_LOAD_MODE) | +| `-lm, --load-mode MODE` | model loading mode (default: mmap)
- none: no special loading mode
- mmap: memory-map model (if mmap disabled, slower load but may reduce pageouts if not using mlock)
- mlock: force system to keep model in RAM rather than swapping or compressing
- mmap+mlock: mmap + force system to keep model in RAM rather than swapping or compressing
- dio: use DirectIO if available

(env: LLAMA_ARG_LOAD_MODE) | | `--numa TYPE` | attempt optimizations that help on some NUMA systems
- distribute: spread execution evenly over all nodes
- isolate: only spawn threads on CPUs on the node that execution started on
- numactl: use the CPU map provided by numactl
if run without this previously, it is recommended to drop the system page cache before using this
see https://github.com/ggml-org/llama.cpp/issues/1437
(env: LLAMA_ARG_NUMA) | | `-dev, --device ` | comma-separated list of devices to use for offloading (none = don't offload)
use --list-devices to see a list of available devices
(env: LLAMA_ARG_DEVICE) | | `--list-devices` | print list of available devices and exit | diff --git a/tools/llama-bench/llama-bench.cpp b/tools/llama-bench/llama-bench.cpp index 29ad352d0cf3..dc1c7caf4a1c 100644 --- a/tools/llama-bench/llama-bench.cpp +++ b/tools/llama-bench/llama-bench.cpp @@ -429,45 +429,45 @@ static void print_usage(int /* argc */, char ** argv) { } printf("\n"); printf("test parameters:\n"); - printf(" -m, --model (default: %s)\n", join(cmd_params_defaults.model, ",").c_str()); - printf(" -hf, -hfr, --hf-repo /[:quant] Hugging Face model repository; quant is optional, case-insensitive\n"); - printf(" default to Q4_K_M, or falls back to the first file in the repo if Q4_K_M doesn't exist.\n"); - printf(" example: ggml-org/GLM-4.7-Flash-GGUF:Q4_K_M\n"); - printf(" (default: unused)\n"); - printf(" -hff, --hf-file Hugging Face model file. If specified, it will override the quant in --hf-repo\n"); - printf(" (default: unused)\n"); - printf(" -hft, --hf-token Hugging Face access token\n"); - printf(" (default: value from HF_TOKEN environment variable)\n"); - printf(" --offline Offline mode: forces use of cache, prevents network access\n"); - printf(" (default: disabled)\n"); - printf(" -p, --n-prompt (default: %s)\n", join(cmd_params_defaults.n_prompt, ",").c_str()); - printf(" -n, --n-gen (default: %s)\n", join(cmd_params_defaults.n_gen, ",").c_str()); - printf(" -pg (default: %s)\n", join(transform_to_str(cmd_params_defaults.n_pg, pair_str), ",").c_str()); - printf(" -d, --n-depth (default: %s)\n", join(cmd_params_defaults.n_depth, ",").c_str()); - printf(" -b, --batch-size (default: %s)\n", join(cmd_params_defaults.n_batch, ",").c_str()); - printf(" -ub, --ubatch-size (default: %s)\n", join(cmd_params_defaults.n_ubatch, ",").c_str()); - printf(" -ctk, --cache-type-k (default: %s)\n", join(transform_to_str(cmd_params_defaults.type_k, ggml_type_name), ",").c_str()); - printf(" -ctv, --cache-type-v (default: %s)\n", join(transform_to_str(cmd_params_defaults.type_v, ggml_type_name), ",").c_str()); - printf(" -t, --threads (default: %s)\n", join(cmd_params_defaults.n_threads, ",").c_str()); - printf(" -C, --cpu-mask (default: %s)\n", join(cmd_params_defaults.cpu_mask, ",").c_str()); - printf(" --cpu-strict <0|1> (default: %s)\n", join(cmd_params_defaults.cpu_strict, ",").c_str()); - printf(" --poll <0...100> (default: %s)\n", join(cmd_params_defaults.poll, ",").c_str()); - printf(" -ngl, --n-gpu-layers (default: %s)\n", join(cmd_params_defaults.n_gpu_layers, ",").c_str()); - printf(" -ncmoe, --n-cpu-moe (default: %s)\n", join(cmd_params_defaults.n_cpu_moe, ",").c_str()); - printf(" -sm, --split-mode (default: %s)\n", join(transform_to_str(cmd_params_defaults.split_mode, split_mode_str), ",").c_str()); - printf(" -mg, --main-gpu (default: %s)\n", join(cmd_params_defaults.main_gpu, ",").c_str()); - printf(" -nkvo, --no-kv-offload <0|1> (default: %s)\n", join(cmd_params_defaults.no_kv_offload, ",").c_str()); - printf(" -fa, --flash-attn (default: %s)\n", join(transform_to_str(cmd_params_defaults.flash_attn, llama_flash_attn_type_name), ",").c_str()); - printf(" -dev, --device (default: auto)\n"); - printf(" -lm, --load-mode (default: %s)\n", join(transform_to_str(cmd_params_defaults.load_mode, llama_load_mode_name), ",").c_str()); - printf(" -mmp, --mmap <0|1> (DEPRECATED IN FAVOUR OF --load-mode)\n"); - printf(" -dio, --direct-io <0|1> (DEPRECATED IN FAVOUR OF --load-mode)\n"); - printf(" -embd, --embeddings <0|1> (default: %s)\n", join(cmd_params_defaults.embeddings, ",").c_str()); - printf(" -ts, --tensor-split (default: 0)\n"); + printf(" -m, --model (default: %s)\n", join(cmd_params_defaults.model, ",").c_str()); + printf(" -hf, -hfr, --hf-repo /[:quant] Hugging Face model repository; quant is optional, case-insensitive\n"); + printf(" default to Q4_K_M, or falls back to the first file in the repo if Q4_K_M doesn't exist.\n"); + printf(" example: ggml-org/GLM-4.7-Flash-GGUF:Q4_K_M\n"); + printf(" (default: unused)\n"); + printf(" -hff, --hf-file Hugging Face model file. If specified, it will override the quant in --hf-repo\n"); + printf(" (default: unused)\n"); + printf(" -hft, --hf-token Hugging Face access token\n"); + printf(" (default: value from HF_TOKEN environment variable)\n"); + printf(" --offline Offline mode: forces use of cache, prevents network access\n"); + printf(" (default: disabled)\n"); + printf(" -p, --n-prompt (default: %s)\n", join(cmd_params_defaults.n_prompt, ",").c_str()); + printf(" -n, --n-gen (default: %s)\n", join(cmd_params_defaults.n_gen, ",").c_str()); + printf(" -pg (default: %s)\n", join(transform_to_str(cmd_params_defaults.n_pg, pair_str), ",").c_str()); + printf(" -d, --n-depth (default: %s)\n", join(cmd_params_defaults.n_depth, ",").c_str()); + printf(" -b, --batch-size (default: %s)\n", join(cmd_params_defaults.n_batch, ",").c_str()); + printf(" -ub, --ubatch-size (default: %s)\n", join(cmd_params_defaults.n_ubatch, ",").c_str()); + printf(" -ctk, --cache-type-k (default: %s)\n", join(transform_to_str(cmd_params_defaults.type_k, ggml_type_name), ",").c_str()); + printf(" -ctv, --cache-type-v (default: %s)\n", join(transform_to_str(cmd_params_defaults.type_v, ggml_type_name), ",").c_str()); + printf(" -t, --threads (default: %s)\n", join(cmd_params_defaults.n_threads, ",").c_str()); + printf(" -C, --cpu-mask (default: %s)\n", join(cmd_params_defaults.cpu_mask, ",").c_str()); + printf(" --cpu-strict <0|1> (default: %s)\n", join(cmd_params_defaults.cpu_strict, ",").c_str()); + printf(" --poll <0...100> (default: %s)\n", join(cmd_params_defaults.poll, ",").c_str()); + printf(" -ngl, --n-gpu-layers (default: %s)\n", join(cmd_params_defaults.n_gpu_layers, ",").c_str()); + printf(" -ncmoe, --n-cpu-moe (default: %s)\n", join(cmd_params_defaults.n_cpu_moe, ",").c_str()); + printf(" -sm, --split-mode (default: %s)\n", join(transform_to_str(cmd_params_defaults.split_mode, split_mode_str), ",").c_str()); + printf(" -mg, --main-gpu (default: %s)\n", join(cmd_params_defaults.main_gpu, ",").c_str()); + printf(" -nkvo, --no-kv-offload <0|1> (default: %s)\n", join(cmd_params_defaults.no_kv_offload, ",").c_str()); + printf(" -fa, --flash-attn (default: %s)\n", join(transform_to_str(cmd_params_defaults.flash_attn, llama_flash_attn_type_name), ",").c_str()); + printf(" -dev, --device (default: auto)\n"); + printf(" -lm, --load-mode (default: %s)\n", join(transform_to_str(cmd_params_defaults.load_mode, llama_load_mode_name), ",").c_str()); + printf(" -mmp, --mmap <0|1> (DEPRECATED IN FAVOUR OF --load-mode)\n"); + printf(" -dio, --direct-io <0|1> (DEPRECATED IN FAVOUR OF --load-mode)\n"); + printf(" -embd, --embeddings <0|1> (default: %s)\n", join(cmd_params_defaults.embeddings, ",").c_str()); + printf(" -ts, --tensor-split (default: 0)\n"); printf(" -ot --override-tensor =;...\n"); - printf(" (default: disabled)\n"); - printf(" -nopo, --no-op-offload <0|1> (default: 0)\n"); - printf(" --no-host <0|1> (default: %s)\n", join(cmd_params_defaults.no_host, ",").c_str()); + printf(" (default: disabled)\n"); + printf(" -nopo, --no-op-offload <0|1> (default: 0)\n"); + printf(" --no-host <0|1> (default: %s)\n", join(cmd_params_defaults.no_host, ",").c_str()); printf("\n"); printf( "Multiple values can be given for each parameter by separating them with ','\n" @@ -785,6 +785,8 @@ static cmd_params parse_cmd_params(int argc, char ** argv) { mode = LLAMA_LOAD_MODE_MMAP; } else if (m == "mlock") { mode = LLAMA_LOAD_MODE_MLOCK; + } else if (m == "mmap+mlock") { + mode = LLAMA_LOAD_MODE_MMAP_MLOCK; } else if (m == "dio") { mode = LLAMA_LOAD_MODE_DIRECT_IO; } else { diff --git a/tools/server/README.md b/tools/server/README.md index d34565455455..25aacf9f516f 100644 --- a/tools/server/README.md +++ b/tools/server/README.md @@ -72,10 +72,10 @@ For the full list of features, please refer to [server's changelog](https://gith | `-ctv, --cache-type-v TYPE` | KV cache data type for V
allowed values: f32, f16, bf16, q8_0, q4_0, q4_1, iq4_nl, q5_0, q5_1
(default: f16)
(env: LLAMA_ARG_CACHE_TYPE_V) | | `-dt, --defrag-thold N` | KV cache defragmentation threshold (DEPRECATED)
(env: LLAMA_ARG_DEFRAG_THOLD) | | `--rpc SERVERS` | comma-separated list of RPC servers (host:port)
(env: LLAMA_ARG_RPC) | -| `--mlock` | DEPRECATED in favor of `--load-mode`: mmap + force system to keep model in RAM rather than swapping or compressing
(env: LLAMA_ARG_MLOCK) | +| `--mlock` | DEPRECATED in favor of `--load-mode`: force system to keep model in RAM rather than swapping or compressing
(env: LLAMA_ARG_MLOCK) | | `--mmap, --no-mmap` | DEPRECATED in favor of `--load-mode`: whether to memory-map model. (if mmap disabled, slower load but may reduce pageouts if not using mlock)
(env: LLAMA_ARG_MMAP) | | `-dio, --direct-io, -ndio, --no-direct-io` | DEPRECATED in favor of `--load-mode`: use DirectIO if available
(env: LLAMA_ARG_DIO) | -| `-lm, --load-mode MODE` | model loading mode (default: mmap)
- none: no special loading mode
- mmap: memory-map model (if mmap disabled, slower load but may reduce pageouts if not using mlock)
- mlock: mmap + force system to keep model in RAM rather than swapping or compressing
- dio: use DirectIO if available

(env: LLAMA_ARG_LOAD_MODE) | +| `-lm, --load-mode MODE` | model loading mode (default: mmap)
- none: no special loading mode
- mmap: memory-map model (if mmap disabled, slower load but may reduce pageouts if not using mlock)
- mlock: force system to keep model in RAM rather than swapping or compressing
- mmap+mlock: mmap + force system to keep model in RAM rather than swapping or compressing
- dio: use DirectIO if available

(env: LLAMA_ARG_LOAD_MODE) | | `--numa TYPE` | attempt optimizations that help on some NUMA systems
- distribute: spread execution evenly over all nodes
- isolate: only spawn threads on CPUs on the node that execution started on
- numactl: use the CPU map provided by numactl
if run without this previously, it is recommended to drop the system page cache before using this
see https://github.com/ggml-org/llama.cpp/issues/1437
(env: LLAMA_ARG_NUMA) | | `-dev, --device ` | comma-separated list of devices to use for offloading (none = don't offload)
use --list-devices to see a list of available devices
(env: LLAMA_ARG_DEVICE) | | `--list-devices` | print list of available devices and exit | From b910200897f1d3193a711438d8482a59e517ed08 Mon Sep 17 00:00:00 2001 From: shalinib-ibm Date: Mon, 27 Jul 2026 14:22:03 +0530 Subject: [PATCH 018/190] ggml-cpu: Enable BF16 tiled gemm optimization on PowerPC (#26068) --- ggml/src/ggml-cpu/llamafile/sgemm.cpp | 146 ++++++++++++++++++++++---- 1 file changed, 126 insertions(+), 20 deletions(-) diff --git a/ggml/src/ggml-cpu/llamafile/sgemm.cpp b/ggml/src/ggml-cpu/llamafile/sgemm.cpp index 23bcd54c122a..99b7d5afa2f9 100644 --- a/ggml/src/ggml-cpu/llamafile/sgemm.cpp +++ b/ggml/src/ggml-cpu/llamafile/sgemm.cpp @@ -1797,14 +1797,6 @@ class tinyBLAS_Q0_AVX { //PPC Implementation #if defined(__MMA__) -#define SAVE_ACC(ACC, ii, jj) \ - __builtin_mma_disassemble_acc(vec_C, ACC); \ - for (int I = 0; I < 4; I++) { \ - for (int J = 0; J < 4; J++) { \ - *((float*)(C+ii+((jj+J)*ldc)+I)) = *((float*)&vec_C[I]+J); \ - } \ - } \ - template struct mma_instr; @@ -1834,10 +1826,49 @@ class tinyBLAS_HP16_PPC { } void matmul(int64_t m, int64_t n) { - mnpack(0, m, 0, n); + int64_t mc = 256; + int64_t nc = 256; + int64_t kc = 256; + #if defined(_AIX) || defined(__BIG_ENDIAN__) + mc = 128; + nc = 128; + kc = 128; + #endif + if (k < kc) { + kc = k; + } + bool can_use_tiled = (m % mc == 0) && (n % nc == 0) && (k % kc == 0); + if (can_use_tiled) { + matmul_tiled(m, n, mc, nc, kc); + } else { + mnpack(0, m, 0, n); + } } private: + __attribute__((always_inline)) + inline void save_acc(acc_t * ACC, int64_t ii, int64_t jj) { + vec_t vec_C[4]; + __builtin_mma_disassemble_acc(vec_C, ACC); + for (int I = 0; I < 4; I++) { + for (int J = 0; J < 4; J++) { + *((float *)(C+ii+((jj+J)*ldc)+I)) = *((float *)&vec_C[I]+J); + } + } + } + + __attribute__((always_inline)) + inline void add_save_acc(acc_t * ACC, int64_t ii, int64_t jj) { + vec_t vec_C[4]; + __builtin_mma_disassemble_acc(vec_C, ACC); + for (int I = 0; I < 4; I++) { + for (int J = 0; J < 4; J++) { + float * c_ptr = (float *)(C+ii+((jj+J)*ldc)+I); + *c_ptr += *((float *)&vec_C[I]+J); + } + } + } + void vector_permute_store(vec_t *c, int numVec, unsigned char *vecOffset) { vec_t t[8], s[8]; vec_t swiz1 = {0, 1, 2, 3, 16, 17, 18, 19, 4, 5, 6, 7, 20, 21, 22, 23}; @@ -1896,6 +1927,7 @@ class tinyBLAS_HP16_PPC { j = (rows >> 3); if (j > 0) { do { + aoffsets[0] = aoffset; if (cols == 4) { aoffsets[0] = aoffset; for (int it = 1; it < 4; ++it) @@ -1910,17 +1942,17 @@ class tinyBLAS_HP16_PPC { } i = (cols >> 3); if (i > 0) { - aoffsets[0] = aoffset; for (int it = 1; it < 8; ++it) { aoffsets[it] = aoffsets[it-1] + lda; } aoffset += 8 * lda; + do { for (int it = 0; it < 8; ++it) c_arr[it] = vec_xl(0, (vector unsigned char*)aoffsets[it]); vector_permute_store(c_arr, 8, vecOffset); for (int it = 0; it < 8; ++it) - aoffsets[it] = aoffsets[it] + 8*lda; + aoffsets[it] = aoffsets[it] + 8; vecOffset += 128; i--; } while(i > 0); @@ -2147,8 +2179,8 @@ class tinyBLAS_HP16_PPC { mma_instr::outer_product(&acc_1, vec_A[x], vec_B[x+4]); } } - SAVE_ACC(&acc_0, ii, jj); - SAVE_ACC(&acc_1, ii, jj+4); + save_acc(&acc_0, ii, jj); + save_acc(&acc_1, ii, jj+4); } void KERNEL_8x4(int64_t ii, int64_t jj) { @@ -2164,8 +2196,8 @@ class tinyBLAS_HP16_PPC { mma_instr::outer_product(&acc_1, vec_A[x+4], vec_B[x]); } } - SAVE_ACC(&acc_0, ii, jj); - SAVE_ACC(&acc_1, ii+4, jj); + save_acc(&acc_0, ii, jj); + save_acc(&acc_1, ii+4, jj); } @@ -2186,13 +2218,64 @@ class tinyBLAS_HP16_PPC { mma_instr::outer_product(&acc_3, vec_A[x+4], vec_B[x+4]); } } - - SAVE_ACC(&acc_0, ii, jj); - SAVE_ACC(&acc_1, ii, jj+4); - SAVE_ACC(&acc_2, ii+4, jj); - SAVE_ACC(&acc_3, ii+4, jj+4); + save_acc(&acc_0, ii, jj); + save_acc(&acc_1, ii, jj+4); + save_acc(&acc_2, ii+4, jj); + save_acc(&acc_3, ii+4, jj+4); } + inline void MMA_16x8(vec_t * vec_A0, vec_t * vec_A1, vec_t * vec_B, acc_t * acc) { + for (int x = 0; x < 4; x ++) { + mma_instr::outer_product(&acc[0], vec_A0[x], vec_B[x]); + mma_instr::outer_product(&acc[1], vec_A0[x], vec_B[x+4]); + mma_instr::outer_product(&acc[2], vec_A0[x+4], vec_B[x]); + mma_instr::outer_product(&acc[3], vec_A0[x+4], vec_B[x+4]); + mma_instr::outer_product(&acc[4], vec_A1[x], vec_B[x]); + mma_instr::outer_product(&acc[5], vec_A1[x], vec_B[x+4]); + mma_instr::outer_product(&acc[6], vec_A1[x+4], vec_B[x]); + mma_instr::outer_product(&acc[7], vec_A1[x+4], vec_B[x+4]); + } + } + void KERNEL(int64_t ii, int64_t jj, int64_t mc, int64_t nc, int64_t kc, vec_t * vec_A, vec_t * vec_B, int64_t kk) { + for (int64_t i = 0; i < mc; i += 16) { + int A_base_addr = (mc / 8) * (i / 8) * 8; + for (int64_t j = 0; j < nc; j += 8) { + int B_base_addr = (nc / 8) * (j / 8) * 8; + acc_t acc[8]; + vec_t A0_block[8]; vec_t A1_block[8]; + for (int x = 0; x < 8; x++) + __builtin_mma_xxsetaccz(&acc[x]); + for (int64_t l = 0; l < kc; l += 8) { + int A0_block_idx = A_base_addr + (l / 8) * 8; + int A1_block_idx = A0_block_idx + (mc / 8) * 8; + int B_block_idx = B_base_addr + (l / 8) * 8; + vec_t* A0_block = &vec_A[A0_block_idx]; + vec_t* A1_block = &vec_A[A1_block_idx]; + vec_t* B_block = &vec_B[B_block_idx]; + MMA_16x8(A0_block, A1_block, B_block, acc); + } + if (kk == 0) { + save_acc(&acc[0], ii + i, jj + j); + save_acc(&acc[1], ii + i, jj + j + 4); + save_acc(&acc[2], ii + i + 4, jj + j); + save_acc(&acc[3], ii + i + 4, jj + j + 4); + save_acc(&acc[4], ii + i + 8, jj + j); + save_acc(&acc[5], ii + i + 8, jj + j + 4); + save_acc(&acc[6], ii + i + 12, jj + j); + save_acc(&acc[7], ii + i + 12, jj + j + 4); + } else { + add_save_acc(&acc[0], ii + i, jj + j); + add_save_acc(&acc[1], ii + i, jj + j + 4); + add_save_acc(&acc[2], ii + i + 4, jj + j); + add_save_acc(&acc[3], ii + i + 4, jj + j + 4); + add_save_acc(&acc[4], ii + i + 8, jj + j); + add_save_acc(&acc[5], ii + i + 8, jj + j + 4); + add_save_acc(&acc[6], ii + i + 12, jj + j); + add_save_acc(&acc[7], ii + i + 12, jj + j + 4); + } + } + } + } template void gemm_small(int64_t m0, int64_t m, int64_t n0, int64_t n) { int64_t ytiles = (m - m0) / RM; @@ -2281,6 +2364,29 @@ class tinyBLAS_HP16_PPC { } } + void matmul_tiled(int64_t m, int64_t n, int64_t mc, int64_t nc, int64_t kc) { + int64_t ytiles = m / mc; + int64_t xtiles = n / nc; + int64_t tiles = xtiles * ytiles; + int64_t duty = (tiles + nth - 1) / nth; + int64_t start = duty * ith; + int64_t end = start + duty; + if (end > tiles) { + end = tiles; + } + for (int64_t job = start; job < end; ++job) { + int64_t ii = (job / xtiles) * mc; + int64_t jj = (job % xtiles) * nc; + for (int64_t kk = 0; kk < k; kk += kc) { + vec_t A_pack[kc * mc / 8]; + vec_t B_pack[kc * nc / 8]; + packNormal(A + (ii * lda) + kk, lda, kc, mc, (uint8_t *)A_pack); + packNormal(B + (jj * ldb) + kk, ldb, kc, nc, (uint8_t *)B_pack); + KERNEL(ii, jj, mc, nc, kc, A_pack, B_pack, kk); + } + } + } + template NOINLINE void gemm(int64_t m0, int64_t m, int64_t n0, int64_t n) { int64_t ytiles = (m - m0) / RM; From 419b881c02dcbdd7ecc47d2b5c935e269642269d Mon Sep 17 00:00:00 2001 From: Xuan-Son Nguyen Date: Mon, 27 Jul 2026 12:00:56 +0200 Subject: [PATCH 019/190] docs: add exception about weight folding (#26168) * docs: add exception about weight folding * add example --- docs/development/HOWTO-add-model.md | 2 ++ skills/add-new-model/SKILL.md | 1 + 2 files changed, 3 insertions(+) diff --git a/docs/development/HOWTO-add-model.md b/docs/development/HOWTO-add-model.md index 632e79881a43..102f479eb02c 100644 --- a/docs/development/HOWTO-add-model.md +++ b/docs/development/HOWTO-add-model.md @@ -144,6 +144,8 @@ Examples: - Gemma 3 folds the `1 +` of its `norm(1 + weight)` normalization into the weights at conversion time, so the graph just does a plain RMS norm. - Qwen3-Next applies its tensor permutation during conversion (in `modify_tensors`), so the graph can consume the already-permuted weights directly. +Exception: a plain `weight * scale` with a constant scale is usually better left to inference time rather than folded into the weight at conversion. The scale conceptually applies to the activation, not the weight, so folding it into the weight can hurt numerical stability, and it shifts the weight's value range in a way that can make quantization worse. In this case, write the scale to GGUF as its own metadata key (e.g. `%s.attention.output_scale`, `%s.attention.value_scale`, `%s.embedding_scale`) and apply it in the graph, instead of pre-multiplying the weight tensor during conversion. + ### Working with ggml_rope_ext PyTorch implementations usually prefer explicitly calculating `freq_cis`/`sin`/`cos` components. However, in llama.cpp, most RoPE operations can be handled via `ggml_rope_ext`, which does not require a sin/cos matrix. This saves memory while allowing the GGML RoPE kernel to be fused with other ops. diff --git a/skills/add-new-model/SKILL.md b/skills/add-new-model/SKILL.md index 68be866c7b8d..f76d1abfd768 100644 --- a/skills/add-new-model/SKILL.md +++ b/skills/add-new-model/SKILL.md @@ -76,6 +76,7 @@ These recur often enough in review comments on past add-model PRs that they're w - Don't ship unfinished or unverified speculative-decoding (e.g. MTP) scaffolding in the base model PR - if it hasn't actually been confirmed to work, pull it out and land it as its own follow-up. - Conversion code should call into the base class's existing hparam logic (e.g. `super().set_gguf_parameters()`) rather than re-deriving it - large blocks of code that duplicate what `TextModel`/`MmprojModel` already provide will get flagged as redundant. - Do constant tensor modifications (e.g. `norm(1 + weight)`) and permutations/chunking at conversion time, not in the graph - see HOWTO-add-model.md's "Prefer conversion-time tensor modifications" tip (Gemma 3 folds its `1 +` into the weights, Qwen3-Next permutes in `modify_tensors`). Doing these at runtime in the graph is very likely to be rejected as over-complicated; if you genuinely can't do it at conversion time, open a discussion first explaining why rather than implementing it in the graph. + - Exception: a plain `weight * scale` with a constant scale is usually better applied at inference time instead of being folded into the weight at conversion. The scale conceptually applies to the activation, not the weight, so folding it in can hurt numerical stability, and it shifts the weight's value range in a way that can make quantization worse. ## Validation checklist From ddfc2288e426b49f5e008c3c2b7c184c7b3ac520 Mon Sep 17 00:00:00 2001 From: Pascal Date: Mon, 27 Jul 2026 12:10:59 +0200 Subject: [PATCH 020/190] common: fix explicit -md precedence over draft sidecar resolution (#26165) * common: fix explicit -md precedence over draft sidecar resolution Follow-up of #25955, an explicit --model-draft file given with -hfd was silently overridden by the sidecar resolution of the draft repo, and its path was never resolved to a local file. An explicit draft file selection now disables the sidecar resolution, so the manual CLI configuration wins over the automatic one. * common: apply the -hfd tag to the sidecar resolution The sidecar selection was anchored on the primary of the draft plan, so a tag without a matching full model aborted the whole plan, and the sidecar quant silently followed the default model pick. The tag now anchors the sidecar directly: exact tag match first, then closest quant to the tag, and a requested sidecar resolves even when no full model matches the tag. A wired draft sidecar also counts as an explicit draft, so the main plan no longer downloads a second one. * common: promote speculative load logs from trace to info Show the loaded draft model and the MTP draft context at the default verbosity, for consistency with the mmproj and primary logs. Co-authored-by: Georgi Gerganov --------- Co-authored-by: Georgi Gerganov --- common/arg.cpp | 12 ++++++++ common/download.cpp | 69 +++++++++++++++++++++++++++++++----------- common/speculative.cpp | 4 +-- 3 files changed, 66 insertions(+), 19 deletions(-) diff --git a/common/arg.cpp b/common/arg.cpp index 84c19e06329a..3bc9574d8175 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -539,6 +539,13 @@ void common_models_handler_apply(common_models_handler & handler, common_params } }; + // an explicit draft file selection (e.g. -md with -hfd) disables the sidecar resolution of the draft repo + if (!params.speculative.draft.mparams.hf_file.empty()) { + plan_spec.mtp = {}; + plan_spec.dflash = {}; + plan_spec.eagle3 = {}; + } + // infer the speculative type from the sidecar shipped by the draft repo when none is requested if (spec_types_is_default(params)) { if (!plan_spec.mtp.local_path.empty()) { @@ -588,6 +595,11 @@ void common_models_handler_apply(common_models_handler & handler, common_params }); } + // a wired draft sidecar counts as an explicit draft for the main plan fallback below + if (spec_sidecar_found) { + had_spec_url = true; + } + // handle plan_spec (e.g. --spec-draft-hf) if (!plan_spec.model_files.empty() && !had_spec_url && !spec_sidecar_found) { add_tasks(plan_spec.model_files, plan_spec.primary, params.speculative.draft.mparams); diff --git a/common/download.cpp b/common/download.cpp index e8e938426f2a..3776c6c7eb68 100644 --- a/common/download.cpp +++ b/common/download.cpp @@ -568,16 +568,30 @@ static hf_cache::hf_files get_split_files(const hf_cache::hf_files & files, } // pick the best sibling GGUF whose filename contains `keyword` (e.g. "mmproj" / "mtp"), -// preferring deeper shared directory prefix with the model, then closest quantization +// preferring deeper shared directory prefix with the model, then exact `tag` match, +// then closest quantization to the tag when given, or to the model otherwise static hf_cache::hf_file find_best_sibling(const hf_cache::hf_files & files, const std::string & model, - const std::string & keyword) { + const std::string & keyword, + const std::string & tag = "") { hf_cache::hf_file best; size_t best_depth = 0; int best_diff = 0; + bool best_exact = false; bool found = false; - auto model_bits = extract_quant_bits(model); + std::string tag_upper = tag; + for (char & c : tag_upper) { + c = (char) std::toupper((unsigned char) c); + } + + int model_bits = 0; + if (!tag_upper.empty()) { + auto pos = tag_upper.find_first_of("0123456789"); + model_bits = pos == std::string::npos ? 0 : std::stoi(tag_upper.substr(pos)); + } else { + model_bits = extract_quant_bits(model); + } auto model_parts = string_split(model, '/'); auto model_dir = model_parts.end() - 1; @@ -600,10 +614,19 @@ static hf_cache::hf_file find_best_sibling(const hf_cache::hf_files & files, auto bits = extract_quant_bits(f.path); auto diff = std::abs(bits - model_bits); - if (!found || depth > best_depth || (depth == best_depth && diff < best_diff)) { + std::string path_upper = f.path; + for (char & c : path_upper) { + c = (char) std::toupper((unsigned char) c); + } + bool exact = !tag_upper.empty() && path_upper.find("-" + tag_upper + ".") != std::string::npos; + + if (!found || depth > best_depth || + (depth == best_depth && exact && !best_exact) || + (depth == best_depth && exact == best_exact && diff < best_diff)) { best = f; best_depth = depth; best_diff = diff; + best_exact = exact; found = true; } } @@ -616,18 +639,21 @@ static hf_cache::hf_file find_best_mmproj(const hf_cache::hf_files & files, } static hf_cache::hf_file find_best_mtp(const hf_cache::hf_files & files, - const std::string & model) { - return find_best_sibling(files, model, "mtp-"); + const std::string & model, + const std::string & tag = "") { + return find_best_sibling(files, model, "mtp-", tag); } static hf_cache::hf_file find_best_eagle3(const hf_cache::hf_files & files, - const std::string & model) { - return find_best_sibling(files, model, "eagle3-"); + const std::string & model, + const std::string & tag = "") { + return find_best_sibling(files, model, "eagle3-", tag); } static hf_cache::hf_file find_best_dflash(const hf_cache::hf_files & files, - const std::string & model) { - return find_best_sibling(files, model, "dflash-"); + const std::string & model, + const std::string & tag = "") { + return find_best_sibling(files, model, "dflash-", tag); } static bool gguf_filename_is_model(const std::string & filepath) { @@ -736,27 +762,36 @@ common_download_hf_plan common_download_get_hf_plan(const common_params_model & } } else { primary = find_best_model(all, tag); - if (primary.path.empty()) { + // a requested sidecar can resolve on its own, without a full model of the same tag + if (primary.path.empty() && !opts.download_mtp && !opts.download_dflash && !opts.download_eagle3) { LOG_ERR("%s: no GGUF files found in repository %s\n", __func__, repo.c_str()); list_available_gguf_files(all); return plan; } } - plan.primary = primary; - plan.model_files = get_split_files(all, primary); + if (!primary.path.empty()) { + plan.primary = primary; + plan.model_files = get_split_files(all, primary); + } - if (opts.download_mmproj) { + if (opts.download_mmproj && !primary.path.empty()) { plan.mmproj = find_best_mmproj(all, primary.path); } if (opts.download_mtp) { - plan.mtp = find_best_mtp(all, primary.path); + plan.mtp = find_best_mtp(all, primary.path, tag); } if (opts.download_dflash) { - plan.dflash = find_best_dflash(all, primary.path); + plan.dflash = find_best_dflash(all, primary.path, tag); } if (opts.download_eagle3) { - plan.eagle3 = find_best_eagle3(all, primary.path); + plan.eagle3 = find_best_eagle3(all, primary.path, tag); + } + + if (primary.path.empty() && + plan.mtp.local_path.empty() && plan.dflash.local_path.empty() && plan.eagle3.local_path.empty()) { + LOG_ERR("%s: no GGUF files found in repository %s\n", __func__, repo.c_str()); + list_available_gguf_files(all); } return plan; diff --git a/common/speculative.cpp b/common/speculative.cpp index 3cb08767bd46..ee94d7c37662 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -2284,7 +2284,7 @@ common_speculative_init_result::common_speculative_init_result( std::string model_path; if (has_draft) { model_path = params.speculative.draft.mparams.path; - LOG_TRC("%s: loading draft model '%s'\n", __func__, model_path.c_str()); + LOG_INF("%s: loading draft model '%s'\n", __func__, model_path.c_str()); llama_model * model_dft = llama_model_load_from_file(params.model.path.c_str(), mparams); if (model_dft == NULL) { @@ -2304,7 +2304,7 @@ common_speculative_init_result::common_speculative_init_result( } else if (spec_mtp) { model_path = params.model.path; - LOG_TRC("%s: creating MTP draft context against the target model '%s'\n", __func__, model_path.c_str()); + LOG_INF("%s: creating MTP draft context against the target model '%s'\n", __func__, model_path.c_str()); llama_context * ctx_dft = llama_init_from_model(model_tgt, cparams); if (ctx_dft == nullptr) { From 7ef790f90a4772eadd6966815f6d3ae7d93c7e03 Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Mon, 27 Jul 2026 13:11:20 +0300 Subject: [PATCH 021/190] tests : remove unnecessary sync in test-save-load-state (#26166) --- tests/test-save-load-state.cpp | 2 -- 1 file changed, 2 deletions(-) diff --git a/tests/test-save-load-state.cpp b/tests/test-save-load-state.cpp index bbb025617f6f..6e93ce6fb8da 100644 --- a/tests/test-save-load-state.cpp +++ b/tests/test-save-load-state.cpp @@ -44,8 +44,6 @@ static llama_tokens generate_tokens(llama_context * ctx, llama_sampler * smpl, i n_past++; } - llama_synchronize(ctx); - return result; } From dee2a846b82f15d27f84a48fa387cb53e0d99c25 Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Mon, 27 Jul 2026 14:54:46 +0300 Subject: [PATCH 022/190] ggml : adjust logic for offloading ops to weight's backend (#25832) * ggml : adjust logic for offloading ops to weight's backend * llama : dsv4 graph fixes --- ggml/src/ggml-backend.cpp | 43 +++++++++++++++++++++++---------------- src/llama-context.cpp | 5 +++-- src/models/deepseek4.cpp | 6 +++++- 3 files changed, 34 insertions(+), 20 deletions(-) diff --git a/ggml/src/ggml-backend.cpp b/ggml/src/ggml-backend.cpp index 87615921c09b..7f4e252dca39 100644 --- a/ggml/src/ggml-backend.cpp +++ b/ggml/src/ggml-backend.cpp @@ -906,26 +906,35 @@ static int ggml_backend_sched_backend_id_from_cur(ggml_backend_sched_t sched, st } // operations with weights are preferably run on the same backend as the weights - for (int i = 0; i < GGML_MAX_SRC; i++) { - const struct ggml_tensor * src = tensor->src[i]; - if (src == NULL) { - continue; - } - // skip ROPE since the rope freqs tensor is too small to choose a backend based on it - // not an ideal solution - if (tensor->op != GGML_OP_ROPE && src->buffer != NULL && src->buffer->usage == GGML_BACKEND_BUFFER_USAGE_WEIGHTS) { - int src_backend_id = ggml_backend_sched_backend_from_buffer(sched, src, tensor); - // check if a backend with higher prio wants to offload the op - if (sched->op_offload && src_backend_id == sched->n_backends - 1 && ggml_backend_buffer_is_host(src->buffer)) { - for (int b = 0; b < src_backend_id; b++) { - if (ggml_backend_supports_op(sched->backends[b], tensor) && ggml_backend_offload_op(sched->backends[b], tensor)) { - SET_CAUSE(tensor, "1.off"); - return b; + // TODO: there are exceptions (see below) - not an ideal solution + bool allow = true; + + // skip ROPE since the rope freqs tensor is too small to choose a backend based on it + allow = allow && tensor->op != GGML_OP_ROPE; + + // skip FLASH_ATTN_EXT since the sinks tensor is too small to choose a based based on it + allow = allow && tensor->op != GGML_OP_FLASH_ATTN_EXT; + + if (allow) { + for (int i = 0; i < GGML_MAX_SRC; i++) { + const struct ggml_tensor * src = tensor->src[i]; + if (src == NULL) { + continue; + } + if (src->buffer != NULL && src->buffer->usage == GGML_BACKEND_BUFFER_USAGE_WEIGHTS) { + int src_backend_id = ggml_backend_sched_backend_from_buffer(sched, src, tensor); + // check if a backend with higher prio wants to offload the op + if (sched->op_offload && src_backend_id == sched->n_backends - 1 && ggml_backend_buffer_is_host(src->buffer)) { + for (int b = 0; b < src_backend_id; b++) { + if (ggml_backend_supports_op(sched->backends[b], tensor) && ggml_backend_offload_op(sched->backends[b], tensor)) { + SET_CAUSE(tensor, "1.off"); + return b; + } } } + SET_CAUSE(tensor, "1.wgt%d", i); + return src_backend_id; } - SET_CAUSE(tensor, "1.wgt%d", i); - return src_backend_id; } } diff --git a/src/llama-context.cpp b/src/llama-context.cpp index c512477c0eab..012894e13f68 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -2473,11 +2473,12 @@ llm_graph_cb llama_context::graph_get_cb() const { ggml_set_name(cur, name); } - // norm may be automatically assigned to the backend of the previous layer, increasing data transfer between backends + // - norm may be automatically assigned to the backend of the previous layer, increasing data transfer between backends + // - force the last op of the layer on the specified backend to avoid running it on the backend of the next layer due to scheduling // FIXME: fix in ggml_backend_sched const bool full_offload = model.n_gpu_layers() > model.hparams.n_layer_all; if (ubatch.n_tokens < 32 || full_offload) { - if (il != -1 && strcmp(name, "norm") == 0) { + if (il != -1 && (strcmp(name, "norm") == 0 || strcmp(name, "l_last") == 0)) { const auto & dev_layer = model.dev_layer(il); for (const auto & backend : backends) { if (ggml_backend_get_device(backend.get()) == dev_layer) { diff --git a/src/models/deepseek4.cpp b/src/models/deepseek4.cpp index 5ad6473ce203..2d41dace0b28 100644 --- a/src/models/deepseek4.cpp +++ b/src/models/deepseek4.cpp @@ -1133,6 +1133,10 @@ llama_model_deepseek4::graph::graph(const llama_model & model, const llm_graph_p &post, &comb, il); cb(cur, "hc_ffn_pre", il); + ggml_build_forward_expand(gf, residual); + ggml_build_forward_expand(gf, post); + ggml_build_forward_expand(gf, comb); + cur = build_norm(cur, model.layers[il].ffn_norm, nullptr, LLM_NORM_RMS, il); cb(cur, "ffn_norm", il); @@ -1175,7 +1179,7 @@ llama_model_deepseek4::graph::graph(const llama_model & model, const llm_graph_p inpL = build_hc_post(cur, residual, post, comb, il); inpL = build_cvec(inpL, il); - cb(inpL, "l_out", il); + cb(inpL, "l_last", il); } if (inp_out_ids) { From 8e8681e0e20820a7736960381d71dec06a830163 Mon Sep 17 00:00:00 2001 From: Titaniumtown Date: Mon, 27 Jul 2026 05:33:11 -0700 Subject: [PATCH 023/190] sycl(build): parallelize ocloc invocations (#25903) --- ggml/src/ggml-sycl/CMakeLists.txt | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/ggml/src/ggml-sycl/CMakeLists.txt b/ggml/src/ggml-sycl/CMakeLists.txt index 1c17d20df12b..a8d9c0d804bf 100644 --- a/ggml/src/ggml-sycl/CMakeLists.txt +++ b/ggml/src/ggml-sycl/CMakeLists.txt @@ -199,9 +199,20 @@ if (GGML_SYCL_DEVICE_ARCH) -fsycl-targets=spir64_gen "SHELL:-Xsycl-target-backend=spir64_gen \"-device ${GGML_SYCL_DEVICE_ARCH}\"" ) + + # Pass through parallel job (process) count for parallelising the + # `llvm-foreach -- ocloc` invocation for compiling AOT device images. + include(ProcessorCount) + ProcessorCount(_ggml_sycl_nproc) + if (_ggml_sycl_nproc LESS 1) + set(_ggml_sycl_nproc 1) + endif() + set(GGML_SYCL_MAX_PARALLEL_LINK_JOBS ${_ggml_sycl_nproc} CACHE STRING + "Parallel ocloc jobs for spir64_gen AOT device-image lowering") target_link_options( ggml-sycl PRIVATE -fsycl-targets=spir64_gen "SHELL:-Xsycl-target-backend=spir64_gen \"-device ${GGML_SYCL_DEVICE_ARCH}\"" + -fsycl-max-parallel-link-jobs=${GGML_SYCL_MAX_PARALLEL_LINK_JOBS} ) endif() From 0324696b8e5fe340dc94b64714e4c9aab03084a2 Mon Sep 17 00:00:00 2001 From: Jonas Jankaitis <111707981+John-194@users.noreply.github.com> Date: Mon, 27 Jul 2026 16:21:37 +0300 Subject: [PATCH 024/190] fit : count nextn (MTP) blocks in n_gpu_layers so front layers stay on GPU (#26177) --- common/fit.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/common/fit.cpp b/common/fit.cpp index c79221cb00fa..c82d066ad444 100644 --- a/common/fit.cpp +++ b/common/fit.cpp @@ -136,7 +136,7 @@ static std::vector common_get_device_memory_data_impl( devs.push_back(llama_model_get_device(model, i)); } - hp_ngl = llama_model_n_layer(model); + hp_ngl = llama_model_n_layer(model) + llama_model_n_layer_nextn(model); hp_n_ctx_train = llama_model_n_ctx_train(model); hp_n_expert = llama_model_n_expert(model); From b77d646751d01c0962bc203b6809e9d94f7d50b7 Mon Sep 17 00:00:00 2001 From: zql <37731799+zqlcode@users.noreply.github.com> Date: Mon, 27 Jul 2026 23:04:18 +0800 Subject: [PATCH 025/190] model: Add support for Nanbeige4.2 (#25994) * support nanbeige4.2 model * fix * fix flake8 Lint check * fix loop bound check and drop redundant head_dim --------- Co-authored-by: root --- conversion/__init__.py | 1 + conversion/nanbeige.py | 24 +++++ gguf-py/gguf/constants.py | 25 ++++- gguf-py/gguf/gguf_writer.py | 6 ++ src/llama-arch.cpp | 3 + src/llama-arch.h | 3 + src/llama-context.cpp | 1 + src/llama-model.cpp | 3 + src/models/models.h | 16 ++++ src/models/nanbeige.cpp | 184 ++++++++++++++++++++++++++++++++++++ 10 files changed, 265 insertions(+), 1 deletion(-) create mode 100644 conversion/nanbeige.py create mode 100644 src/models/nanbeige.cpp diff --git a/conversion/__init__.py b/conversion/__init__.py index b2bb7e5161eb..45c001b78f96 100644 --- a/conversion/__init__.py +++ b/conversion/__init__.py @@ -167,6 +167,7 @@ "ModernBertForMaskedLM": "bert", "ModernBertForSequenceClassification": "bert", "ModernBertModel": "bert", + "NanbeigeForCausalLM": "nanbeige", "NemotronForCausalLM": "nemotron", "NemotronHForCausalLM": "nemotron", "NeoBERT": "bert", diff --git a/conversion/nanbeige.py b/conversion/nanbeige.py new file mode 100644 index 000000000000..f1fc425b3a09 --- /dev/null +++ b/conversion/nanbeige.py @@ -0,0 +1,24 @@ +from __future__ import annotations + +from .base import ModelBase, gguf, logger +from .llama import LlamaModel + + +@ModelBase.register("NanbeigeForCausalLM") +class NanbeigeModel(LlamaModel): + model_arch = gguf.MODEL_ARCH.NANBEIGE + undo_permute = True + + def set_gguf_parameters(self): + super().set_gguf_parameters() + hparams = self.hparams + + n_loops = int(hparams.get("num_loops", 1) or 1) + if n_loops < 1: + n_loops = 1 + self.gguf_writer.add_num_loops(n_loops) + logger.info(f"gguf: num_loops = {n_loops}") + + skip_loop_final_norm = bool(hparams.get("skip_loop_final_norm", False)) + self.gguf_writer.add_skip_loop_final_norm(skip_loop_final_norm) + logger.info(f"gguf: skip_loop_final_norm = {skip_loop_final_norm}") diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index 2071e3eaa8a4..78e3c29a0d38 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -145,6 +145,8 @@ class LLM: TOKEN_SHIFT_COUNT = "{arch}.token_shift_count" INTERLEAVE_MOE_LAYER_STEP = "{arch}.interleave_moe_layer_step" FULL_ATTENTION_INTERVAL = "{arch}.full_attention_interval" + NUM_LOOPS = "{arch}.num_loops" + SKIP_LOOP_FINAL_NORM = "{arch}.skip_loop_final_norm" HASH_LAYER_COUNT = "{arch}.hash_layer_count" ACTIVATION_SPARSITY_SCALE = "{arch}.activation_sparsity_scale" ALTUP_ACTIVE_IDX = "{arch}.altup.active_idx" @@ -545,6 +547,7 @@ class MODEL_ARCH(IntEnum): KIMI_LINEAR = auto() TALKIE = auto() MELLUM = auto() + NANBEIGE = auto() class VISION_PROJECTOR_TYPE(IntEnum): @@ -1134,6 +1137,7 @@ class MODEL_TENSOR(IntEnum): MODEL_ARCH.KIMI_LINEAR: "kimi-linear", MODEL_ARCH.TALKIE: "talkie", MODEL_ARCH.MELLUM: "mellum", + MODEL_ARCH.NANBEIGE: "nanbeige", } VISION_PROJECTOR_TYPE_NAMES: dict[VISION_PROJECTOR_TYPE, str] = { @@ -4505,7 +4509,22 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.FFN_DOWN_EXP, MODEL_TENSOR.FFN_UP_EXP, ], - # TODO + MODEL_ARCH.NANBEIGE: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ROPE_FREQS, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.ATTN_ROT_EMBD, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + ], } # tensors that will not be serialized @@ -4572,6 +4591,10 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.ROPE_FREQS, MODEL_TENSOR.ATTN_ROT_EMBD, ], + MODEL_ARCH.NANBEIGE: [ + MODEL_TENSOR.ROPE_FREQS, + MODEL_TENSOR.ATTN_ROT_EMBD, + ], } # diff --git a/gguf-py/gguf/gguf_writer.py b/gguf-py/gguf/gguf_writer.py index ba08f8d65004..ecf1f17ee8b3 100644 --- a/gguf-py/gguf/gguf_writer.py +++ b/gguf-py/gguf/gguf_writer.py @@ -908,6 +908,12 @@ def add_wkv_head_size(self, size: int) -> None: def add_token_shift_count(self, count: int) -> None: self.add_uint32(Keys.LLM.TOKEN_SHIFT_COUNT.format(arch=self.arch), count) + def add_num_loops(self, count: int) -> None: + self.add_uint32(Keys.LLM.NUM_LOOPS.format(arch=self.arch), count) + + def add_skip_loop_final_norm(self, value: bool) -> None: + self.add_bool(Keys.LLM.SKIP_LOOP_FINAL_NORM.format(arch=self.arch), value) + def add_interleave_moe_layer_step(self, value: int) -> None: self.add_uint32(Keys.LLM.INTERLEAVE_MOE_LAYER_STEP.format(arch=self.arch), value) diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index 39bf2c79590b..c01706785070 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -143,6 +143,7 @@ static const std::map LLM_ARCH_NAMES = { { LLM_ARCH_KIMI_LINEAR, "kimi-linear" }, { LLM_ARCH_TALKIE, "talkie" }, { LLM_ARCH_MELLUM, "mellum" }, + { LLM_ARCH_NANBEIGE, "nanbeige" }, { LLM_ARCH_UNKNOWN, "(unknown)" }, }; @@ -221,6 +222,8 @@ static const std::map LLM_KV_NAMES = { { LLM_KV_TOKEN_SHIFT_COUNT, "%s.token_shift_count" }, { LLM_KV_INTERLEAVE_MOE_LAYER_STEP, "%s.interleave_moe_layer_step" }, { LLM_KV_FULL_ATTENTION_INTERVAL, "%s.full_attention_interval" }, + { LLM_KV_NUM_LOOPS, "%s.num_loops" }, + { LLM_KV_SKIP_LOOP_FINAL_NORM, "%s.skip_loop_final_norm" }, { LLM_KV_ATTENTION_HEAD_COUNT, "%s.attention.head_count" }, { LLM_KV_ATTENTION_HEAD_COUNT_KV, "%s.attention.head_count_kv" }, diff --git a/src/llama-arch.h b/src/llama-arch.h index 2e3916a0beee..1c9aebb0bbdc 100644 --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -148,6 +148,7 @@ enum llm_arch { LLM_ARCH_EAGLE3, LLM_ARCH_MINIMAX_M3, LLM_ARCH_DFLASH, + LLM_ARCH_NANBEIGE, LLM_ARCH_UNKNOWN, }; @@ -226,6 +227,8 @@ enum llm_kv { LLM_KV_TOKEN_SHIFT_COUNT, LLM_KV_INTERLEAVE_MOE_LAYER_STEP, LLM_KV_FULL_ATTENTION_INTERVAL, + LLM_KV_NUM_LOOPS, + LLM_KV_SKIP_LOOP_FINAL_NORM, LLM_KV_ATTENTION_HEAD_COUNT, LLM_KV_ATTENTION_HEAD_COUNT_KV, diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 012894e13f68..9b399d6096b1 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -2339,6 +2339,7 @@ uint32_t llama_context::graph_max_nodes(uint32_t n_tokens) const { model.arch == LLM_ARCH_QWEN35 || model.arch == LLM_ARCH_QWEN35MOE || model.arch == LLM_ARCH_DEEPSEEK4 || + model.arch == LLM_ARCH_NANBEIGE || model.arch == LLM_ARCH_MINIMAX_M3) { return std::max(n_tokens * 40, 32u * model.n_tensors()); } diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 074acbe1fa27..be0a0df55d62 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -85,6 +85,8 @@ static llama_model * llama_model_mapping(llm_arch arch, const llama_model_params return new llama_model_stablelm(params); case LLM_ARCH_MELLUM: return new llama_model_mellum(params); + case LLM_ARCH_NANBEIGE: + return new llama_model_nanbeige(params); case LLM_ARCH_QWEN: return new llama_model_qwen(params); case LLM_ARCH_QWEN2: @@ -2491,6 +2493,7 @@ llama_rope_type llama_model_rope_type(const llama_model * model) { case LLM_ARCH_LLAMA_EMBED: case LLM_ARCH_MAINCODER: case LLM_ARCH_GLM_DSA: + case LLM_ARCH_NANBEIGE: return LLAMA_ROPE_TYPE_NORM; // the pairs of head values are offset by n_rot/2 diff --git a/src/models/models.h b/src/models/models.h index 916459e12782..92ebfafa1e29 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -424,6 +424,22 @@ struct llama_model_mellum : public llama_model_base { std::unique_ptr build_arch_graph(const llm_graph_params & params) const override; }; +struct llama_model_nanbeige : public llama_model_base { + llama_model_nanbeige(const struct llama_model_params & params) : llama_model_base(params) {} + void load_arch_hparams(llama_model_loader & ml) override; + void load_arch_tensors(llama_model_loader & ml) override; + + int n_loops = 1; + int n_layer_phys = 0; + bool skip_loop_final_norm = false; + + struct graph : public llm_graph_context { + graph(const llama_model & model, const llm_graph_params & params); + }; + + std::unique_ptr build_arch_graph(const llm_graph_params & params) const override; +}; + struct llama_model_qwen : public llama_model_base { llama_model_qwen(const struct llama_model_params & params) : llama_model_base(params) {} void load_arch_hparams(llama_model_loader & ml) override; diff --git a/src/models/nanbeige.cpp b/src/models/nanbeige.cpp new file mode 100644 index 000000000000..3a546600fa27 --- /dev/null +++ b/src/models/nanbeige.cpp @@ -0,0 +1,184 @@ +#include "models.h" + +void llama_model_nanbeige::load_arch_hparams(llama_model_loader & ml) { + ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps); + + uint32_t n_loops_u = 1; + ml.get_key(LLM_KV_NUM_LOOPS, n_loops_u, false); + GGML_ASSERT(n_loops_u >= 1); + + skip_loop_final_norm = false; + ml.get_key(LLM_KV_SKIP_LOOP_FINAL_NORM, skip_loop_final_norm, false); + + n_layer_phys = (int) hparams.n_layer(); + + // Bound-check before casting: signed int mul can overflow and bypass the guard. + GGML_ASSERT((size_t) n_layer_phys * (size_t) n_loops_u <= (size_t) LLAMA_MAX_LAYERS); + n_loops = (int) n_loops_u; + + // Expand logical layer count before load_tensors() allocates layers / KV. + if (n_loops > 1) { + for (int j = 1; j < n_loops; ++j) { + for (int i = 0; i < n_layer_phys; ++i) { + const int dst = i + j * n_layer_phys; + hparams.n_head_arr[dst] = hparams.n_head_arr[i]; + hparams.n_head_kv_arr[dst] = hparams.n_head_kv_arr[i]; + hparams.n_ff_arr[dst] = hparams.n_ff_arr[i]; + hparams.is_swa_impl[dst] = hparams.is_swa_impl[i]; + hparams.is_recr_impl[dst] = hparams.is_recr_impl[i]; + } + } + hparams.n_layer_all = (uint32_t) ((size_t) n_layer_phys * (size_t) n_loops); + } + + type = LLM_TYPE_UNKNOWN; +} + +void llama_model_nanbeige::load_arch_tensors(llama_model_loader &) { + LLAMA_LOAD_LOCALS; + + tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0); + + output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0); + output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_vocab}, TENSOR_NOT_REQUIRED); + if (output == NULL) { + output = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, TENSOR_DUPLICATED); + } + + const int n_phys = n_layer_phys > 0 ? n_layer_phys : n_layer; + for (int i = 0; i < n_phys; ++i) { + auto & layer = layers[i]; + + layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0); + + create_tensor_qkv(layer, i, n_embd, n_embd_head_k * n_head, n_embd_k_gqa, n_embd_v_gqa, 0); + layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd_head_k * n_head, n_embd}, 0); + + layer.rope_freqs = create_tensor(tn(LLM_TENSOR_ROPE_FREQS, "weight", i), {n_rot/2}, + TENSOR_NOT_REQUIRED | (i != 0 ? TENSOR_DUPLICATED : 0)); + + layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0); + layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd, n_ff}, 0); + layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), { n_ff, n_embd}, 0); + layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, n_ff}, 0); + } + + // Share physical weights across loops; each slot still has its own KV index. + if (n_loops > 1) { + for (int j = 1; j < n_loops; ++j) { + for (int i = 0; i < n_phys; ++i) { + layers[i + j * n_phys] = layers[i]; + } + } + } +} + +std::unique_ptr llama_model_nanbeige::build_arch_graph(const llm_graph_params & params) const { + return std::make_unique(*this, params); +} + +llama_model_nanbeige::graph::graph(const llama_model & model, const llm_graph_params & params) : + llm_graph_context(params) { + const auto & nb = static_cast(model); + + const int64_t n_embd_head = hparams.n_embd_head_v(); + GGML_ASSERT(n_embd_head == hparams.n_embd_head_k()); + + const int n_phys = nb.n_layer_phys > 0 ? nb.n_layer_phys : (int) n_layer; + const int n_loops = nb.n_loops > 0 ? nb.n_loops : 1; + + ggml_tensor * cur; + ggml_tensor * inpL; + + inpL = build_inp_embd(model.tok_embd); + + ggml_tensor * inp_pos = build_inp_pos(); + + auto * inp_attn = build_attn_inp_kv(); + + const float kq_scale = hparams.f_attention_scale == 0.0f + ? 1.0f / sqrtf(float(n_embd_head)) + : hparams.f_attention_scale; + + ggml_tensor * inp_out_ids = build_inp_out_ids(); + + for (int il = 0; il < n_layer; ++il) { + ggml_tensor * inpSA = inpL; + + cur = build_norm(inpL, model.layers[il].attn_norm, NULL, LLM_NORM_RMS, il); + cb(cur, "attn_norm", il); + + { + ggml_tensor * rope_factors = model.get_rope_factors(cparams, il); + + auto [Qcur, Kcur, Vcur] = build_qkv(model.layers[il], cur, + n_embd_head, n_head, n_head_kv, il); + + Qcur = ggml_rope_ext( + ctx0, Qcur, inp_pos, rope_factors, + n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow); + + Kcur = ggml_rope_ext( + ctx0, Kcur, inp_pos, rope_factors, + n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow); + + cb(Qcur, "Qcur", il); + cb(Kcur, "Kcur", il); + cb(Vcur, "Vcur", il); + + cur = build_attn(inp_attn, + model.layers[il].wo, model.layers[il].wo_b, model.layers[il].wo_s, + Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il); + cb(cur, "attn_out", il); + } + + if (il == n_layer - 1 && inp_out_ids) { + cur = ggml_get_rows(ctx0, cur, inp_out_ids); + inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids); + } + + ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpSA); + cb(ffn_inp, "ffn_inp", il); + + cur = build_norm(ffn_inp, model.layers[il].ffn_norm, NULL, LLM_NORM_RMS, il); + cb(cur, "ffn_norm", il); + + cur = build_ffn(cur, + model.layers[il].ffn_up, model.layers[il].ffn_up_b, model.layers[il].ffn_up_s, + model.layers[il].ffn_gate, model.layers[il].ffn_gate_b, model.layers[il].ffn_gate_s, + model.layers[il].ffn_down, model.layers[il].ffn_down_b, model.layers[il].ffn_down_s, + NULL, LLM_FFN_SILU, LLM_FFN_PAR, il); + cb(cur, "ffn_out", il); + + cur = ggml_add(ctx0, cur, ffn_inp); + cb(cur, "ffn_out", il); + + cur = build_cvec(cur, il); + cb(cur, "l_out", il); + + inpL = cur; + + if (n_loops > 1 && + ((il + 1) % n_phys) == 0 && + (il + 1) < n_layer && + !nb.skip_loop_final_norm) { + cur = build_norm(inpL, model.output_norm, NULL, LLM_NORM_RMS, il); + cb(cur, "loop_norm", il); + inpL = cur; + } + } + + cur = inpL; + + cur = build_norm(cur, model.output_norm, NULL, LLM_NORM_RMS, -1); + cb(cur, "result_norm", -1); + res->t_embd = cur; + + cur = build_lora_mm(model.output, cur, model.output_s); + cb(cur, "result_output", -1); + res->t_logits = cur; + + ggml_build_forward_expand(gf, cur); +} From 0e4a0362239713ea95a6864a17a8de4b0ad90d62 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Adrien=20Gallou=C3=ABt?= Date: Mon, 27 Jul 2026 18:19:59 +0200 Subject: [PATCH 026/190] common : add common_print_available_devices() (#26170) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Adrien Gallouët --- common/arg.cpp | 40 ++++++++++++++++++++----------- common/arg.h | 3 +++ tools/llama-bench/llama-bench.cpp | 17 +------------ 3 files changed, 30 insertions(+), 30 deletions(-) diff --git a/common/arg.cpp b/common/arg.cpp index 3bc9574d8175..79480e06f9d2 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -1061,6 +1061,31 @@ static std::vector parse_device_list(const std::string & val return devices; } +void common_print_available_devices() { + constexpr size_t MiB = 1024 * 1024; + std::vector devices; + + ggml_backend_load_all(); + + for (size_t i = 0; i < ggml_backend_dev_count(); ++i) { + auto * dev = ggml_backend_dev_get(i); + if (ggml_backend_dev_type(dev) != GGML_BACKEND_DEVICE_TYPE_CPU) { + devices.push_back(dev); + } + } + printf("Available devices:\n"); + + if (devices.empty()) { + printf(" (none)\n"); + return; + } + for (auto * dev : devices) { + size_t free, total; + ggml_backend_dev_memory(dev, &free, &total); + printf(" %s: %s (%zu MiB, %zu MiB free)\n", ggml_backend_dev_name(dev), ggml_backend_dev_description(dev), total / MiB, free / MiB); + } +} + static void add_rpc_devices(const std::string & servers) { auto rpc_servers = string_split(servers, ','); if (rpc_servers.empty()) { @@ -2588,20 +2613,7 @@ common_params_context common_params_parser_init(common_params & params, llama_ex {"--list-devices"}, "print list of available devices and exit", [](common_params &) { - ggml_backend_load_all(); - std::vector devices; - for (size_t i = 0; i < ggml_backend_dev_count(); ++i) { - auto * dev = ggml_backend_dev_get(i); - if (ggml_backend_dev_type(dev) != GGML_BACKEND_DEVICE_TYPE_CPU) { - devices.push_back(dev); - } - } - printf("Available devices:\n"); - for (auto * dev : devices) { - size_t free, total; - ggml_backend_dev_memory(dev, &free, &total); - printf(" %s: %s (%zu MiB, %zu MiB free)\n", ggml_backend_dev_name(dev), ggml_backend_dev_description(dev), total / 1024 / 1024, free / 1024 / 1024); - } + common_print_available_devices(); exit(0); } )); diff --git a/common/arg.h b/common/arg.h index 54a38b9cce4a..8f609e356fe2 100644 --- a/common/arg.h +++ b/common/arg.h @@ -123,6 +123,9 @@ struct common_params_context { // if one argument has invalid value, it will automatically display usage of the specific argument (and not the full usage message) bool common_params_parse(int argc, char ** argv, common_params & params, llama_example ex, void(*print_usage)(int, char **) = nullptr); +// load all backends and print the list of available (non-CPU) devices to stdout +void common_print_available_devices(); + // parse input arguments from CLI into a map bool common_params_to_map(int argc, char ** argv, llama_example ex, std::map & out_map); diff --git a/tools/llama-bench/llama-bench.cpp b/tools/llama-bench/llama-bench.cpp index dc1c7caf4a1c..c17a27b54019 100644 --- a/tools/llama-bench/llama-bench.cpp +++ b/tools/llama-bench/llama-bench.cpp @@ -670,22 +670,7 @@ static cmd_params parse_cmd_params(int argc, char ** argv) { break; } } else if (arg == "--list-devices") { - std::vector devices; - for (size_t i = 0; i < ggml_backend_dev_count(); ++i) { - auto * dev = ggml_backend_dev_get(i); - if (ggml_backend_dev_type(dev) != GGML_BACKEND_DEVICE_TYPE_CPU) { - devices.push_back(dev); - } - } - printf("Available devices:\n"); - if (devices.empty()) { - printf(" (none)\n"); - } - for (auto * dev : devices) { - size_t free, total; - ggml_backend_dev_memory(dev, &free, &total); - printf(" %s: %s (%zu MiB, %zu MiB free)\n", ggml_backend_dev_name(dev), ggml_backend_dev_description(dev), total / 1024 / 1024, free / 1024 / 1024); - } + common_print_available_devices(); exit(0); } else if (arg == "-t" || arg == "--threads") { if (++i >= argc) { From 1cbfd1988311775425d36c0ce066590f7d3049cf Mon Sep 17 00:00:00 2001 From: Xuan-Son Nguyen Date: Mon, 27 Jul 2026 23:17:09 +0200 Subject: [PATCH 027/190] mtmd: support MiMo-V2.5 audio input (RVQ-based model) (#26190) * gguf converter for mimo audio * fix conv * cpp impl * nits * nits 2 --- conversion/mimo.py | 123 +++++++++++++++-- gguf-py/gguf/constants.py | 49 +++++++ gguf-py/gguf/gguf_writer.py | 18 +++ gguf-py/gguf/tensor_mapping.py | 59 +++++++++ src/llama-quant.cpp | 4 + tools/mtmd/CMakeLists.txt | 1 + tools/mtmd/clip-graph.h | 8 ++ tools/mtmd/clip-impl.h | 26 ++++ tools/mtmd/clip-model.h | 22 ++++ tools/mtmd/clip.cpp | 167 ++++++++++++++++++++++- tools/mtmd/models/mimo-audio.cpp | 218 +++++++++++++++++++++++++++++++ tools/mtmd/models/models.h | 5 + tools/mtmd/mtmd-audio.cpp | 66 ++++++++++ tools/mtmd/mtmd-audio.h | 9 ++ tools/mtmd/mtmd.cpp | 6 + 15 files changed, 770 insertions(+), 11 deletions(-) create mode 100644 tools/mtmd/models/mimo-audio.cpp diff --git a/conversion/mimo.py b/conversion/mimo.py index 11ec2867940a..ca2ed28ad391 100644 --- a/conversion/mimo.py +++ b/conversion/mimo.py @@ -1,8 +1,9 @@ from __future__ import annotations +import json import re -from typing import Callable, TYPE_CHECKING +from typing import Any, Callable, Iterable, TYPE_CHECKING import torch @@ -229,7 +230,13 @@ def prepare_tensors(self): @ModelBase.register("MiMoV2ForCausalLM") -class MiMoV2VisionModel(MmprojModel): +class MiMoV2VisionAudioModel(MmprojModel): + has_audio_encoder = True + + _audio_tok_hparams: dict[str, Any] | None = None + _rvq_codebook_sizes: list[int] | None = None + _code_embd: dict[int, Tensor] | None = None + def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) assert self.hparams_vision is not None @@ -253,10 +260,22 @@ def __init__(self, *args, **kwargs): self.visual_token_window_size = int(hp.get("visual_token_window_size", -1)) self.use_sink = bool(hp.get("use_sink", False)) + def get_audio_config(self) -> dict[str, Any] | None: + if self._audio_tok_hparams is None: + path = self.dir_model / "audio_tokenizer" / "config.json" + with open(path, "r", encoding="utf-8") as f: + cfg = json.load(f) + # aliases so MmprojModel.find_aparam() / n_block_keys can resolve them + cfg["hidden_size"] = cfg["d_model"] + cfg["intermediate_size"] = cfg["encoder_ffn_dim"] + cfg["num_attention_heads"] = cfg["encoder_attention_heads"] + self._audio_tok_hparams = cfg + return self._audio_tok_hparams + def set_gguf_parameters(self): super().set_gguf_parameters() - self.gguf_writer.add_clip_projector_type(gguf.VisionProjectorType.MIMOVL) + self.gguf_writer.add_clip_vision_projector_type(gguf.VisionProjectorType.MIMOVL) self.gguf_writer.add_vision_use_silu(True) self.gguf_writer.add_vision_head_count_kv(self.num_kv_heads) self.gguf_writer.add_vision_spatial_merge_size(self.spatial_merge_size) @@ -266,19 +285,45 @@ def set_gguf_parameters(self): self.gguf_writer.add_vision_min_pixels(int(self.preprocessor_config["min_pixels"])) self.gguf_writer.add_vision_max_pixels(int(self.preprocessor_config["max_pixels"])) + assert self.hparams_audio is not None + self.gguf_writer.add_clip_audio_projector_type(gguf.VisionProjectorType.MIMO_AUDIO) + self.gguf_writer.add_audio_num_mel_bins(self.hparams_audio["n_mels"]) + self.gguf_writer.add_audio_attention_layernorm_eps(self.hparams_audio.get("layer_norm_eps", 1e-5)) + + assert self._rvq_codebook_sizes is not None + self.gguf_writer.add_audio_rvq_num_quantizers(len(self._rvq_codebook_sizes)) + self.gguf_writer.add_audio_rvq_codebook_size(self._rvq_codebook_sizes) + + n_layer = self.hparams_audio["encoder_layers"] + swa_per_block = self.hparams_audio.get("swa_per_block", 1) + if self.hparams_audio.get("hybrid_attention") and swa_per_block > 1: + wa_pattern = [0 if i % swa_per_block < swa_per_block - 1 else -1 for i in range(n_layer)] + else: + wa_pattern = [-1] * n_layer + self.gguf_writer.add_audio_wa_pattern_mode(wa_pattern) + self.gguf_writer.add_audio_window_size(int(self.hparams_audio["encoder_attn_window_size"][0])) + + audio_cfg = self.global_config["audio_config"] + self.gguf_writer.add_audio_local_block_count(int(audio_cfg["input_local_layers"])) + self.gguf_writer.add_audio_local_group_size(int(audio_cfg["group_size"])) + def tensor_force_quant(self, name, new_name, bid, n_dims): - # Sinks must be F32: any sink-style softmax/mask add in ggml requires - # F32, and we fold sinks into a host-built F32 mask at encode time. - if new_name.endswith(".attn_sinks"): + # for audio encoder: keep codebook in F32 + if new_name in ( + gguf.TENSOR_NAMES[gguf.MODEL_TENSOR.A_ENC_RVQ_CODEBOOK] + ".weight", + gguf.TENSOR_NAMES[gguf.MODEL_TENSOR.A_MM_CODE_EMBD] + ".weight", + ): + return gguf.GGMLQuantizationType.F32 + if ("encoder.conv" in name or "encoder.down_sample_layer" in name) and name.endswith(".weight"): return gguf.GGMLQuantizationType.F32 return super().tensor_force_quant(name, new_name, bid, n_dims) @classmethod def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None: name, _ = item - if not name.startswith("visual."): - return None - return super().filter_tensors(item) + if name.startswith("visual.") or name.startswith("speech_embeddings.") or name.startswith("audio_encoder."): + return super().filter_tensors(item) + return None def modify_tensors(self, data_torch, name, bid): # Conv3D patch embed: split along the temporal axis (kt=2) into two Conv2D @@ -292,4 +337,64 @@ def modify_tensors(self, data_torch, name, bid): yield (embd_name + ".weight.1", data_torch[:, :, 1, ...]) return + if m := re.match(r"^speech_embeddings\.(\d+)\.weight$", name): + if self._code_embd is None: + self._code_embd = {} + self._code_embd[int(m.group(1))] = data_torch + + n_channels = int(self.global_config["audio_config"]["audio_channels"]) + if len(self._code_embd) < n_channels: + return + merged = torch.stack([self._code_embd.pop(i) for i in range(n_channels)], dim=0) + yield (self.format_tensor_name(gguf.MODEL_TENSOR.A_MM_CODE_EMBD), merged) + return + + if "conv1.bias" in name or "conv2.bias" in name: + # transpose conv1/conv2 bias so it broadcasts against [n_frames, C_out, 1] + data_torch = data_torch.unsqueeze(-1) + + if name == "audio_encoder.projection.mlp.0.weight": + yield (self.format_tensor_name(gguf.MODEL_TENSOR.A_MMPROJ, 1), data_torch) + return + if name == "audio_encoder.projection.mlp.2.weight": + yield (self.format_tensor_name(gguf.MODEL_TENSOR.A_MMPROJ, 2), data_torch) + return + yield from super().modify_tensors(data_torch, name, bid) + + def generate_extra_tensors(self) -> Iterable[tuple[str, Tensor]]: + # note: audio encoder is in its own subdir "audio_tokenizer" + from safetensors.torch import load_file + + tok_dir = self.dir_model / "audio_tokenizer" + state_dict = load_file(tok_dir / "model.safetensors") + + codebook_re = re.compile(r"^encoder\.quantizer\.vq\.layers\.(\d+)\._codebook\.embed$") + codebooks: dict[int, Tensor] = {} + + # EMA/training-only RVQ buffers - not needed for inference (nearest-codebook + # lookup only reads "_codebook.embed") + skip_suffixes = ( + "_codebook.cluster_size", + "_codebook.embed_avg", + "_codebook.inited", + ) + for name, tensor in state_dict.items(): + if name.endswith(skip_suffixes): + continue + if m := codebook_re.match(name): + codebooks[int(m.group(1))] = tensor + continue + yield name, tensor + + # gather codebooks and merge into 3D tensor, similar to MoE MLP tensors + n_q = len(codebooks) + ordered = [codebooks[i] for i in range(n_q)] + self._rvq_codebook_sizes = [int(cb.shape[0]) for cb in ordered] + max_bins = max(self._rvq_codebook_sizes) + dim = ordered[0].shape[1] + merged = ordered[0].new_zeros(n_q, max_bins, dim) + for i, cb in enumerate(ordered): + merged[i, : cb.shape[0], :] = cb + + yield (self.format_tensor_name(gguf.MODEL_TENSOR.A_ENC_RVQ_CODEBOOK), merged) diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index 78e3c29a0d38..f9264425f6e1 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -376,6 +376,12 @@ class ClipAudio: CONV_KERNEL_SIZE = "clip.audio.conv_kernel_size" MAX_POS_EMB = "clip.audio.max_pos_emb" FEATURE_LAYERS = "clip.audio.feature_layer" # Granite Speech Plus + RVQ_NUM_QUANTIZERS = "clip.audio.rvq.num_quantizers" + RVQ_CODEBOOK_SIZE = "clip.audio.rvq.codebook_size" + WA_PATTERN_MODE = "clip.audio.wa_pattern_mode" # per-layer -1 (full) / 0 (windowed) + WINDOW_SIZE = "clip.audio.window_size" + LOCAL_BLOCK_COUNT = "clip.audio.local_block_count" # mimo-v2.5: input_local_transformer layer count + LOCAL_GROUP_SIZE = "clip.audio.local_group_size" # mimo-v2.5: input_local_transformer grouping size class Attention: HEAD_COUNT = "clip.audio.attention.head_count" @@ -945,6 +951,9 @@ class MODEL_TENSOR(IntEnum): A_ENC_FFN_SCALE_1 = auto() # gemma3n A_ENC_FFN_GATE_1 = auto() # lfm2, gemma3n A_ENC_FFN_DOWN_1 = auto() # lfm2, gemma3n + A_ENC_DOWNSAMPLE_CONV = auto() # mimo-audio-tokenizer: post-transformer downsample conv + A_ENC_DOWNSAMPLE_NORM = auto() # mimo-audio-tokenizer: post-transformer downsample norm + A_ENC_RVQ_CODEBOOK = auto() # mimo-audio-tokenizer: residual vector quantizer codebook, per quantizer index A_MMPROJ = auto() A_MMPROJ_FC = auto() A_MM_NORM_PRE = auto() @@ -953,6 +962,17 @@ class MODEL_TENSOR(IntEnum): A_MM_HARD_EMB_NORM = auto() # gemma3n A_MM_SOFT_EMB_NORM = auto() # gemma3n A_MM_INP_PROJ = auto() # gemma3n + A_MM_CODE_EMBD = auto() # mimo: text-side RVQ code embedding table ("text codebook"), merged 3D [n_channels, vocab, dim] + A_MM_LOCAL_ATTN_Q = auto() # mimo: input_local_transformer (LLM-side connector) + A_MM_LOCAL_ATTN_K = auto() + A_MM_LOCAL_ATTN_V = auto() + A_MM_LOCAL_ATTN_OUT = auto() + A_MM_LOCAL_FFN_GATE = auto() + A_MM_LOCAL_FFN_UP = auto() + A_MM_LOCAL_FFN_DOWN = auto() + A_MM_LOCAL_LN1 = auto() + A_MM_LOCAL_LN2 = auto() + A_MM_LOCAL_NORM = auto() # final norm after all input_local_transformer layers A_PER_DIM_K_SCALE = auto() # gemma4 A_PER_DIM_SCALE = auto() # gemma4 # nextn/mtp @@ -1532,6 +1552,9 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.A_ENC_FFN_UP_1: "a.blk.{bid}.ffn_up_1", MODEL_TENSOR.A_ENC_FFN_GATE_1: "a.blk.{bid}.ffn_gate_1", MODEL_TENSOR.A_ENC_FFN_DOWN_1: "a.blk.{bid}.ffn_down_1", + MODEL_TENSOR.A_ENC_DOWNSAMPLE_CONV: "a.downsample.conv", + MODEL_TENSOR.A_ENC_DOWNSAMPLE_NORM: "a.downsample.norm", + MODEL_TENSOR.A_ENC_RVQ_CODEBOOK: "a.rvq.codebook", MODEL_TENSOR.A_MMPROJ: "mm.a.mlp.{bid}", MODEL_TENSOR.A_MMPROJ_FC: "mm.a.fc", MODEL_TENSOR.A_MM_NORM_PRE: "mm.a.norm_pre", @@ -1540,6 +1563,17 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.A_MM_SOFT_EMB_NORM: "mm.a.soft_emb_norm", # gemma3n MODEL_TENSOR.A_MM_EMBEDDING: "mm.a.embedding", # gemma3n MODEL_TENSOR.A_MM_HARD_EMB_NORM: "mm.a.hard_emb_norm", # gemma3n + MODEL_TENSOR.A_MM_CODE_EMBD: "mm.a.code_embd", + MODEL_TENSOR.A_MM_LOCAL_ATTN_Q: "mm.a.local_blk.{bid}.attn_q", + MODEL_TENSOR.A_MM_LOCAL_ATTN_K: "mm.a.local_blk.{bid}.attn_k", + MODEL_TENSOR.A_MM_LOCAL_ATTN_V: "mm.a.local_blk.{bid}.attn_v", + MODEL_TENSOR.A_MM_LOCAL_ATTN_OUT: "mm.a.local_blk.{bid}.attn_out", + MODEL_TENSOR.A_MM_LOCAL_FFN_GATE: "mm.a.local_blk.{bid}.ffn_gate", + MODEL_TENSOR.A_MM_LOCAL_FFN_UP: "mm.a.local_blk.{bid}.ffn_up", + MODEL_TENSOR.A_MM_LOCAL_FFN_DOWN: "mm.a.local_blk.{bid}.ffn_down", + MODEL_TENSOR.A_MM_LOCAL_LN1: "mm.a.local_blk.{bid}.ln1", + MODEL_TENSOR.A_MM_LOCAL_LN2: "mm.a.local_blk.{bid}.ln2", + MODEL_TENSOR.A_MM_LOCAL_NORM: "mm.a.local_norm", MODEL_TENSOR.A_PER_DIM_K_SCALE: "a.blk.{bid}.per_dim_k_scale", # gemma4 MODEL_TENSOR.A_PER_DIM_SCALE: "a.blk.{bid}.per_dim_scale", # gemma4 # lfm2 audio @@ -1741,10 +1775,24 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.A_ENC_FFN_UP_1, MODEL_TENSOR.A_ENC_FFN_GATE_1, MODEL_TENSOR.A_ENC_FFN_DOWN_1, + MODEL_TENSOR.A_ENC_DOWNSAMPLE_CONV, + MODEL_TENSOR.A_ENC_DOWNSAMPLE_NORM, + MODEL_TENSOR.A_ENC_RVQ_CODEBOOK, MODEL_TENSOR.A_MMPROJ, MODEL_TENSOR.A_MMPROJ_FC, MODEL_TENSOR.A_MM_NORM_PRE, MODEL_TENSOR.A_MM_NORM_MID, + MODEL_TENSOR.A_MM_CODE_EMBD, + MODEL_TENSOR.A_MM_LOCAL_ATTN_Q, + MODEL_TENSOR.A_MM_LOCAL_ATTN_K, + MODEL_TENSOR.A_MM_LOCAL_ATTN_V, + MODEL_TENSOR.A_MM_LOCAL_ATTN_OUT, + MODEL_TENSOR.A_MM_LOCAL_FFN_GATE, + MODEL_TENSOR.A_MM_LOCAL_FFN_UP, + MODEL_TENSOR.A_MM_LOCAL_FFN_DOWN, + MODEL_TENSOR.A_MM_LOCAL_LN1, + MODEL_TENSOR.A_MM_LOCAL_LN2, + MODEL_TENSOR.A_MM_LOCAL_NORM, MODEL_TENSOR.A_ENC_NORM_CONV, MODEL_TENSOR.A_ENC_LINEAR_POS, MODEL_TENSOR.A_ENC_POS_BIAS_U, @@ -4804,6 +4852,7 @@ class VisionProjectorType: MINICPMV4_6 = "minicpmv4_6" GRANITE_SPEECH = "granite_speech" # audio MIMOVL = "mimovl" + MIMO_AUDIO = "mimo_audio" GRANITE4_VISION = "granite4_vision" diff --git a/gguf-py/gguf/gguf_writer.py b/gguf-py/gguf/gguf_writer.py index ecf1f17ee8b3..bd8629aa119e 100644 --- a/gguf-py/gguf/gguf_writer.py +++ b/gguf-py/gguf/gguf_writer.py @@ -1350,6 +1350,24 @@ def add_audio_attention_layernorm_eps(self, value: float) -> None: def add_audio_num_mel_bins(self, value: int) -> None: self.add_uint32(Keys.ClipAudio.NUM_MEL_BINS, value) + def add_audio_rvq_num_quantizers(self, value: int) -> None: + self.add_uint32(Keys.ClipAudio.RVQ_NUM_QUANTIZERS, value) + + def add_audio_rvq_codebook_size(self, values: Sequence[int]) -> None: + self.add_array(Keys.ClipAudio.RVQ_CODEBOOK_SIZE, values) + + def add_audio_wa_pattern_mode(self, modes: Sequence[int]) -> None: + self.add_array(Keys.ClipAudio.WA_PATTERN_MODE, modes) + + def add_audio_window_size(self, value: int) -> None: + self.add_uint32(Keys.ClipAudio.WINDOW_SIZE, value) + + def add_audio_local_block_count(self, value: int) -> None: + self.add_uint32(Keys.ClipAudio.LOCAL_BLOCK_COUNT, value) + + def add_audio_local_group_size(self, value: int) -> None: + self.add_uint32(Keys.ClipAudio.LOCAL_GROUP_SIZE, value) + def add_audio_stack_factor(self, value: int) -> None: self.add_uint32(Keys.ClipAudio.Projector.STACK_FACTOR, value) diff --git a/gguf-py/gguf/tensor_mapping.py b/gguf-py/gguf/tensor_mapping.py index 62d7a827e35c..8299ac25b432 100644 --- a/gguf-py/gguf/tensor_mapping.py +++ b/gguf-py/gguf/tensor_mapping.py @@ -2095,6 +2095,7 @@ class TensorNameMap: "conformer.pre_encode.conv.{bid}", # lfm2 "model.audio_tower.subsample_conv_projection.conv_{bid}.conv", # gemma3n "conformer.subsample_conv_projection.layer{bid}.conv", # gemma4 + "encoder.conv{bid}", # mimo-audio-tokenizer ), MODEL_TENSOR.A_ENC_CONV1D_NORM: ( @@ -2119,6 +2120,7 @@ class TensorNameMap: MODEL_TENSOR.A_POST_NORM: ( "audio_tower.layer_norm", # ultravox "audio_tower.ln_post", # qwen2omni + "encoder.layer_norm", # mimo-audio-tokenizer ), MODEL_TENSOR.A_ENC_ATTN_Q: ( @@ -2127,6 +2129,7 @@ class TensorNameMap: "conformer.layers.{bid}.attention.attn.q_proj", # gemma3n "conformer.layers.{bid}.self_attn.q_proj", # gemma4 "encoder.layers.{bid}.attn.to_q", # granite_speech + "encoder.layers.{bid}.self_attn.q_proj", # mimo-audio-tokenizer ), MODEL_TENSOR.A_ENC_ATTN_K: ( @@ -2135,6 +2138,7 @@ class TensorNameMap: "conformer.layers.{bid}.attention.attn.k_proj", # gemma3n "conformer.layers.{bid}.self_attn.k_proj", # gemma4 "encoder.layers.{bid}.attn.to_k", # granite_speech (split from to_kv) + "encoder.layers.{bid}.self_attn.k_proj", # mimo-audio-tokenizer ), MODEL_TENSOR.A_ENC_ATTN_V: ( @@ -2143,6 +2147,7 @@ class TensorNameMap: "conformer.layers.{bid}.attention.attn.v_proj", # gemma3n "conformer.layers.{bid}.self_attn.v_proj", # gemma4 "encoder.layers.{bid}.attn.to_v", # granite_speech (split from to_kv) + "encoder.layers.{bid}.self_attn.v_proj", # mimo-audio-tokenizer ), MODEL_TENSOR.A_ENC_ATTN_K_REL: ( @@ -2171,6 +2176,7 @@ class TensorNameMap: "conformer.layers.{bid}.norm_self_att", # lfm2 "conformer.layers.{bid}.attention.pre_attn_norm", # gemma3n "encoder.layers.{bid}.attn.pre_norm", # granite_speech + "encoder.layers.{bid}.self_attn_layer_norm", # mimo-audio-tokenizer ), MODEL_TENSOR.A_ENC_OUTPUT: ( @@ -2179,6 +2185,7 @@ class TensorNameMap: "conformer.layers.{bid}.attention.post", # gemma3n "conformer.layers.{bid}.self_attn.post", # gemma4 "encoder.layers.{bid}.attn.to_out", # granite_speech + "encoder.layers.{bid}.self_attn.out_proj", # mimo-audio-tokenizer ), MODEL_TENSOR.A_ENC_OUTPUT_NORM: ( @@ -2186,6 +2193,7 @@ class TensorNameMap: "conformer.layers.{bid}.norm_out", # lfm2 "conformer.layers.{bid}.attention.post_norm", # gemma3n "encoder.layers.{bid}.post_norm", # granite_speech + "encoder.layers.{bid}.final_layer_norm", # mimo-audio-tokenizer ), MODEL_TENSOR.A_ENC_FFN_NORM: ( @@ -2210,6 +2218,7 @@ class TensorNameMap: "conformer.layers.{bid}.ffw_layer_start.ffw_layer_1", # gemma3n "conformer.layers.{bid}.feed_forward1.ffw_layer_1", # gemma4 "encoder.layers.{bid}.ff1.up_proj", # granite_speech + "encoder.layers.{bid}.fc1", # mimo-audio-tokenizer ), MODEL_TENSOR.A_ENC_FFN_GATE: (), @@ -2220,6 +2229,7 @@ class TensorNameMap: "conformer.layers.{bid}.ffw_layer_start.ffw_layer_2", # gemma3n "conformer.layers.{bid}.feed_forward1.ffw_layer_2", # gemma4 "encoder.layers.{bid}.ff1.down_proj", # granite_speech + "encoder.layers.{bid}.fc2", # mimo-audio-tokenizer ), MODEL_TENSOR.A_ENC_FFN_UP_1: ( @@ -2243,6 +2253,19 @@ class TensorNameMap: "encoder.layers.{bid}.ff2.pre_norm", # granite_speech ), + MODEL_TENSOR.A_ENC_DOWNSAMPLE_CONV: ( + "encoder.down_sample_layer.0", # mimo-audio-tokenizer + ), + + MODEL_TENSOR.A_ENC_DOWNSAMPLE_NORM: ( + "encoder.down_sample_norm", # mimo-audio-tokenizer + ), + + # note: the raw per-quantizer "encoder.quantizer.vq.layers.{i}._codebook.embed" + # tensors are merged (padded + stacked, like MoE experts) into this single 3D + # tensor in conversion code, so no raw-name mapping is registered here. + MODEL_TENSOR.A_ENC_RVQ_CODEBOOK: (), + MODEL_TENSOR.A_ENC_FFN_POST_NORM_1: ( "conformer.layers.{bid}.ffw_layer_end.post_layer_norm", # gemma3n "conformer.layers.{bid}.feed_forward2.post_layer_norm", # gemma4 @@ -2294,6 +2317,42 @@ class TensorNameMap: "audio.multi_modal_projector.ln_mid", # ultravox ), + # note: the raw per-channel "speech_embeddings.{i}" tensors are merged + # (stacked, like MoE experts) into this single 3D tensor in conversion + # code, so no raw-name mapping is registered here. + MODEL_TENSOR.A_MM_CODE_EMBD: (), + + MODEL_TENSOR.A_MM_LOCAL_ATTN_Q: ( + "audio_encoder.input_local_transformer.layers.{bid}.self_attn.q_proj", # mimo-v2.5 + ), + MODEL_TENSOR.A_MM_LOCAL_ATTN_K: ( + "audio_encoder.input_local_transformer.layers.{bid}.self_attn.k_proj", # mimo-v2.5 + ), + MODEL_TENSOR.A_MM_LOCAL_ATTN_V: ( + "audio_encoder.input_local_transformer.layers.{bid}.self_attn.v_proj", # mimo-v2.5 + ), + MODEL_TENSOR.A_MM_LOCAL_ATTN_OUT: ( + "audio_encoder.input_local_transformer.layers.{bid}.self_attn.o_proj", # mimo-v2.5 + ), + MODEL_TENSOR.A_MM_LOCAL_FFN_GATE: ( + "audio_encoder.input_local_transformer.layers.{bid}.mlp.gate_proj", # mimo-v2.5 + ), + MODEL_TENSOR.A_MM_LOCAL_FFN_UP: ( + "audio_encoder.input_local_transformer.layers.{bid}.mlp.up_proj", # mimo-v2.5 + ), + MODEL_TENSOR.A_MM_LOCAL_FFN_DOWN: ( + "audio_encoder.input_local_transformer.layers.{bid}.mlp.down_proj", # mimo-v2.5 + ), + MODEL_TENSOR.A_MM_LOCAL_LN1: ( + "audio_encoder.input_local_transformer.layers.{bid}.input_layernorm", # mimo-v2.5 + ), + MODEL_TENSOR.A_MM_LOCAL_LN2: ( + "audio_encoder.input_local_transformer.layers.{bid}.post_attention_layernorm", # mimo-v2.5 + ), + MODEL_TENSOR.A_MM_LOCAL_NORM: ( + "audio_encoder.input_local_transformer.norm", # mimo-v2.5 + ), + MODEL_TENSOR.A_ENC_CONV_DW: ( "conformer.layers.{bid}.conv.depthwise_conv", # lfm2 "conformer.layers.{bid}.lconv1d.depthwise_conv1d", # gemma3n diff --git a/src/llama-quant.cpp b/src/llama-quant.cpp index 7c0bac07d096..92ebc11b99f3 100644 --- a/src/llama-quant.cpp +++ b/src/llama-quant.cpp @@ -359,6 +359,10 @@ static bool tensor_allows_quantization(const llama_model_quantize_params * param quantize &= name.find(".patch_embd") == std::string::npos; quantize &= name.find(".patch_merger") == std::string::npos; + // audio codebook + quantize &= name.find("a.rvq.codebook") == std::string::npos; + quantize &= name.find("mm.a.code_embd") == std::string::npos; + return quantize; } diff --git a/tools/mtmd/CMakeLists.txt b/tools/mtmd/CMakeLists.txt index fd7ddceb0bf0..18a8288ba048 100644 --- a/tools/mtmd/CMakeLists.txt +++ b/tools/mtmd/CMakeLists.txt @@ -51,6 +51,7 @@ add_library(mtmd models/qwen3vl.cpp models/mimovl.cpp models/qwen3a.cpp + models/mimo-audio.cpp models/step3vl.cpp models/siglip.cpp models/whisper-enc.cpp diff --git a/tools/mtmd/clip-graph.h b/tools/mtmd/clip-graph.h index a95de20a3122..29352abb4c0b 100644 --- a/tools/mtmd/clip-graph.h +++ b/tools/mtmd/clip-graph.h @@ -13,6 +13,14 @@ struct build_vit_opts { ggml_tensor * attn_mask = nullptr; + // TODO @ngxson : merge attn_mask and attn_mask_layers into one call + std::vector attn_mask_layers; // one per layer + + // hook at layer output embeddings + std::function callback_layer_out = nullptr; + + // whether to skip the automatic post-layernorm (model.post_ln_w) applied at the end + bool skip_post_ln = false; }; struct clip_graph { diff --git a/tools/mtmd/clip-impl.h b/tools/mtmd/clip-impl.h index 42374311ce7b..09204113801f 100644 --- a/tools/mtmd/clip-impl.h +++ b/tools/mtmd/clip-impl.h @@ -82,6 +82,12 @@ #define KEY_A_PROJ_WINDOW_SIZE "clip.audio.projector.window_size" #define KEY_A_PROJ_DOWNSAMPLE_RATE "clip.audio.projector.downsample_rate" #define KEY_A_PROJ_HEAD_COUNT "clip.audio.projector.head_count" +#define KEY_A_RVQ_NUM_QUANTIZERS "clip.audio.rvq.num_quantizers" // mimo-audio-tokenizer +#define KEY_A_RVQ_CODEBOOK_SIZE "clip.audio.rvq.codebook_size" // mimo-audio-tokenizer: per-quantizer bin count +#define KEY_A_WA_PATTERN_MODE "clip.audio.wa_pattern_mode" // mimo-audio-tokenizer, per-layer -1 (full) / 0 (windowed) +#define KEY_A_ATTN_WINDOW_SIZE "clip.audio.window_size" // mimo-audio-tokenizer: sliding-window radius +#define KEY_A_LOCAL_BLOCK_COUNT "clip.audio.local_block_count" // mimo-v2.5: input_local_transformer layer count +#define KEY_A_LOCAL_GROUP_SIZE "clip.audio.local_group_size" // mimo-v2.5: input_local_transformer grouping size // // tensor name constants @@ -175,6 +181,24 @@ #define TN_MM_NORM_PRE "mm.a.norm_pre.%s" #define TN_MM_NORM_MID "mm.a.norm_mid.%s" +// mimo-audio-tokenizer +#define TN_A_DOWNSAMPLE_CONV "a.downsample.conv.%s" +#define TN_A_DOWNSAMPLE_NORM "a.downsample.norm.%s" +#define TN_A_RVQ_CODEBOOK "a.rvq.codebook.%s" +// mimo-v2.5: text-side RVQ code embedding ("text codebook") +#define TN_MM_A_CODE_EMBD "mm.a.code_embd.%s" +// mimo-v2.5: LLM-side connector (input_local_transformer) +#define TN_MM_A_LOCAL_ATTN_Q "mm.a.local_blk.%d.attn_q.%s" +#define TN_MM_A_LOCAL_ATTN_K "mm.a.local_blk.%d.attn_k.%s" +#define TN_MM_A_LOCAL_ATTN_V "mm.a.local_blk.%d.attn_v.%s" +#define TN_MM_A_LOCAL_ATTN_OUT "mm.a.local_blk.%d.attn_out.%s" +#define TN_MM_A_LOCAL_FFN_GATE "mm.a.local_blk.%d.ffn_gate.%s" +#define TN_MM_A_LOCAL_FFN_UP "mm.a.local_blk.%d.ffn_up.%s" +#define TN_MM_A_LOCAL_FFN_DOWN "mm.a.local_blk.%d.ffn_down.%s" +#define TN_MM_A_LOCAL_LN1 "mm.a.local_blk.%d.ln1.%s" +#define TN_MM_A_LOCAL_LN2 "mm.a.local_blk.%d.ln2.%s" +#define TN_MM_A_LOCAL_NORM "mm.a.local_norm.%s" + // cogvlm #define TN_MM_POST_FC_NORM "mm.post_fc_norm.%s" #define TN_MM_H_TO_4H "mm.up.%s" @@ -374,6 +398,7 @@ enum projector_type { PROJECTOR_TYPE_MIMOVL, PROJECTOR_TYPE_MINIMAX_M3, PROJECTOR_TYPE_GRANITE4_VISION, + PROJECTOR_TYPE_MIMO_AUDIO, PROJECTOR_TYPE_UNKNOWN, }; @@ -429,6 +454,7 @@ static std::map PROJECTOR_TYPE_NAMES = { { PROJECTOR_TYPE_MIMOVL, "mimovl"}, { PROJECTOR_TYPE_MINIMAX_M3, "minimax_m3"}, { PROJECTOR_TYPE_GRANITE4_VISION, "granite4_vision"}, + { PROJECTOR_TYPE_MIMO_AUDIO, "mimo_audio"}, }; static projector_type clip_projector_type_from_string(const std::string & str) { diff --git a/tools/mtmd/clip-model.h b/tools/mtmd/clip-model.h index 850957d7de1c..8dc87549766e 100644 --- a/tools/mtmd/clip-model.h +++ b/tools/mtmd/clip-model.h @@ -124,6 +124,14 @@ struct clip_hparams { int32_t audio_window_len = -1; int32_t audio_hop_len = -1; + // mimo-audio-tokenizer: residual vector quantizer + int32_t rvq_num_quantizers = 0; + std::vector rvq_codebook_size; // per-quantizer bin count (ragged, e.g. 1024/1024/256/128x17) + + // mimo-v2.5: LLM-side connector (input_local_transformer) + int32_t audio_local_n_layer = 0; + int32_t audio_local_group_size = 0; + // legacy bool has_llava_projector = false; int minicpmv_version = 0; @@ -537,6 +545,20 @@ struct clip_model { ggml_tensor * mm_norm_pre_b = nullptr; ggml_tensor * mm_norm_mid_w = nullptr; + // mimo-audio-tokenizer: post-transformer downsample + RVQ codebook + ggml_tensor * downsample_conv_w = nullptr; // no bias + ggml_tensor * downsample_norm_w = nullptr; + ggml_tensor * downsample_norm_b = nullptr; + ggml_tensor * rvq_codebook = nullptr; // merged 3D [n_q, max_bins, dim] + + // mimo-v2.5: text-side RVQ code embedding ("text codebook") + ggml_tensor * mm_a_code_embd = nullptr; // merged 3D [n_channels, vocab, dim] + + // mimo-v2.5: LLM-side connector (input_local_transformer, separate from the + // audio_tokenizer's own encoder `layers`) + std::vector mm_a_local_layers; + ggml_tensor * mm_a_local_norm_w = nullptr; + // qwen3a ggml_tensor * conv2d_1_w = nullptr; ggml_tensor * conv2d_1_b = nullptr; diff --git a/tools/mtmd/clip.cpp b/tools/mtmd/clip.cpp index e0e2107a0be3..04614b93bd27 100644 --- a/tools/mtmd/clip.cpp +++ b/tools/mtmd/clip.cpp @@ -340,6 +340,11 @@ ggml_tensor * clip_graph::build_vit( auto & layer = model.layers[il]; ggml_tensor * cur = inpL; // inpL = residual, cur = hidden_states + ggml_tensor * attn_mask = opts.attn_mask; + if (opts.attn_mask_layers.size() > (size_t) il) { + attn_mask = opts.attn_mask_layers[il]; + } + // layernorm1 cur = build_norm(cur, layer.ln_1_w, layer.ln_1_b, norm_t, eps, il); cb(cur, "layer_inp_normed", il); @@ -452,7 +457,7 @@ ggml_tensor * clip_graph::build_vit( // build_attn returns a flat 2D [n_embd, n_pos*B] cur = build_attn(layer.o_w, layer.o_b, - Qcur, Kcur, Vcur, opts.attn_mask, kq_scale, il); + Qcur, Kcur, Vcur, attn_mask, kq_scale, il); cb(cur, "attn_out", il); } @@ -471,6 +476,10 @@ ggml_tensor * clip_graph::build_vit( inpL = cur; // inpL = residual, cur = hidden_states + if (opts.callback_layer_out) { + opts.callback_layer_out(cur, il); + } + cb(cur, "ffn_inp", il); // layernorm2 (pre-ffn norm) @@ -519,7 +528,7 @@ ggml_tensor * clip_graph::build_vit( } // post-layernorm - if (model.post_ln_w) { + if (model.post_ln_w && !opts.skip_post_ln) { inpL = build_norm(inpL, model.post_ln_w, model.post_ln_b, norm_t, eps, -1); } @@ -1012,6 +1021,10 @@ static std::unique_ptr clip_get_graph_builder(clip_ctx * ctx, const { builder = std::make_unique(ctx, img); } break; + case PROJECTOR_TYPE_MIMO_AUDIO: + { + builder = std::make_unique(ctx, img); + } break; case PROJECTOR_TYPE_YOUTUVL: { builder = std::make_unique(ctx, img); @@ -1575,6 +1588,45 @@ struct clip_model_loader { hparams.audio_window_len = 400; hparams.audio_hop_len = 160; } break; + case PROJECTOR_TYPE_MIMO_AUDIO: + { + get_u32(KEY_A_RVQ_NUM_QUANTIZERS, hparams.rvq_num_quantizers, false); + get_arr_int(KEY_A_RVQ_CODEBOOK_SIZE, hparams.rvq_codebook_size, false); + if (hparams.rvq_num_quantizers <= 0) { + throw std::runtime_error(string_format("%s: mimo_audio: missing %s\n", __func__, KEY_A_RVQ_NUM_QUANTIZERS)); + } + if ((int) hparams.rvq_codebook_size.size() != hparams.rvq_num_quantizers) { + throw std::runtime_error(string_format( + "%s: mimo_audio: %s length (%zu) must equal %s (%d)\n", __func__, + KEY_A_RVQ_CODEBOOK_SIZE, hparams.rvq_codebook_size.size(), + KEY_A_RVQ_NUM_QUANTIZERS, hparams.rvq_num_quantizers)); + } + hparams.ffn_op = FFN_GELU_ERF; // PyTorch F.gelu default (approximate="none") + hparams.rope_theta = 10000.0f; + + // audio preprocessing params (mel spectrogram) + hparams.audio_sample_rate = 24000; + hparams.audio_n_fft = 960; + hparams.audio_window_len = 960; + hparams.audio_hop_len = 240; + + get_u32(KEY_A_ATTN_WINDOW_SIZE, hparams.attn_window_size); + std::vector wa_pattern; + get_arr_int(KEY_A_WA_PATTERN_MODE, wa_pattern, true); + if ((int) wa_pattern.size() != hparams.n_layer) { + throw std::runtime_error(string_format( + "%s: mimo_audio: %s length (%zu) must equal n_layer (%d)\n", __func__, + KEY_A_WA_PATTERN_MODE, wa_pattern.size(), hparams.n_layer)); + } + hparams.wa_pattern_mode.assign(wa_pattern.begin(), wa_pattern.end()); + + get_u32(KEY_A_LOCAL_BLOCK_COUNT, hparams.audio_local_n_layer); + get_u32(KEY_A_LOCAL_GROUP_SIZE, hparams.audio_local_group_size); + if (hparams.audio_local_group_size <= 0) { + throw std::runtime_error(string_format( + "%s: mimo_audio: %s must be > 0\n", __func__, KEY_A_LOCAL_GROUP_SIZE)); + } + } break; case PROJECTOR_TYPE_PADDLEOCR: { hparams.n_merge = 2; @@ -2444,6 +2496,54 @@ struct clip_model_loader { model.mm_2_w = get_tensor(string_format(TN_MM_AUDIO_MLP, 2, "weight")); model.mm_2_b = get_tensor(string_format(TN_MM_AUDIO_MLP, 2, "bias")); } break; + case PROJECTOR_TYPE_MIMO_AUDIO: + { + model.conv1d_1_w = get_tensor(string_format(TN_CONV1D, 1, "weight")); + model.conv1d_1_b = get_tensor(string_format(TN_CONV1D, 1, "bias")); + model.conv1d_2_w = get_tensor(string_format(TN_CONV1D, 2, "weight")); + model.conv1d_2_b = get_tensor(string_format(TN_CONV1D, 2, "bias")); + model.downsample_conv_w = get_tensor(string_format(TN_A_DOWNSAMPLE_CONV, "weight")); + model.downsample_norm_w = get_tensor(string_format(TN_A_DOWNSAMPLE_NORM, "weight")); + model.downsample_norm_b = get_tensor(string_format(TN_A_DOWNSAMPLE_NORM, "bias")); + model.rvq_codebook = get_tensor(string_format(TN_A_RVQ_CODEBOOK, "weight"), false); + model.mm_a_code_embd = get_tensor(string_format(TN_MM_A_CODE_EMBD, "weight"), false); + if (!model.rvq_codebook || !model.mm_a_code_embd) { + throw std::runtime_error(string_format("%s: mimo_audio: missing %s or %s\n", __func__, + TN_A_RVQ_CODEBOOK, TN_MM_A_CODE_EMBD)); + } + // hparams.rvq_codebook_size comes from GGUF metadata and is independent of the + // tensors' actual shapes - bound it so codebook/code_embd views built from it + // (mimo-audio.cpp) can never read past either tensor's allocated bins/vocab. + for (int32_t bins : hparams.rvq_codebook_size) { + if (bins <= 0 || bins > model.rvq_codebook->ne[1] || bins > model.mm_a_code_embd->ne[1]) { + throw std::runtime_error(string_format( + "%s: mimo_audio: %s entry (%d) out of range for codebook/code_embd tensors\n", + __func__, KEY_A_RVQ_CODEBOOK_SIZE, bins)); + } + } + + // LLM-side connector: input_local_transformer + projection + model.mm_a_local_layers.resize(hparams.audio_local_n_layer); + for (int il = 0; il < hparams.audio_local_n_layer; il++) { + auto & layer = model.mm_a_local_layers[il]; + layer.q_w = get_tensor(string_format(TN_MM_A_LOCAL_ATTN_Q, il, "weight")); + layer.q_b = get_tensor(string_format(TN_MM_A_LOCAL_ATTN_Q, il, "bias")); + layer.k_w = get_tensor(string_format(TN_MM_A_LOCAL_ATTN_K, il, "weight")); + layer.k_b = get_tensor(string_format(TN_MM_A_LOCAL_ATTN_K, il, "bias")); + layer.v_w = get_tensor(string_format(TN_MM_A_LOCAL_ATTN_V, il, "weight")); + layer.v_b = get_tensor(string_format(TN_MM_A_LOCAL_ATTN_V, il, "bias")); + layer.o_w = get_tensor(string_format(TN_MM_A_LOCAL_ATTN_OUT, il, "weight")); + layer.ff_gate_w = get_tensor(string_format(TN_MM_A_LOCAL_FFN_GATE, il, "weight")); + layer.ff_up_w = get_tensor(string_format(TN_MM_A_LOCAL_FFN_UP, il, "weight")); + layer.ff_down_w = get_tensor(string_format(TN_MM_A_LOCAL_FFN_DOWN, il, "weight")); + layer.ln_1_w = get_tensor(string_format(TN_MM_A_LOCAL_LN1, il, "weight")); + layer.ln_2_w = get_tensor(string_format(TN_MM_A_LOCAL_LN2, il, "weight")); + } + model.mm_a_local_norm_w = get_tensor(string_format(TN_MM_A_LOCAL_NORM, "weight")); + + model.mm_1_w = get_tensor(string_format(TN_MM_AUDIO_MLP, 1, "weight")); + model.mm_2_w = get_tensor(string_format(TN_MM_AUDIO_MLP, 2, "weight")); + } break; case PROJECTOR_TYPE_VOXTRAL: { model.conv1d_1_w = get_tensor(string_format(TN_CONV1D, 1, "weight")); @@ -3549,6 +3649,15 @@ int clip_n_output_tokens(const clip_ctx * ctx, const clip_image_f32 * img) { { n_patches = img->nx(); // no downsampling: one token per raw waveform frame } break; + case PROJECTOR_TYPE_MIMO_AUDIO: + { + // conv1(s=1) + conv2(s=2) -> RVQ-encoder downsample conv(k=2,s=2) + int n = img->nx(); + n = (n - 1) / 2 + 1; // conv1 + conv2 + n = (n - 2) / 2 + 1; // downsample conv + const int group_size = params.audio_local_group_size; + n_patches = (n + group_size - 1) / group_size; + } break; case PROJECTOR_TYPE_GRANITE_SPEECH: { const int ws = ctx->model.hparams.audio_proj_window_size; @@ -4376,6 +4485,58 @@ bool clip_image_batch_encode(clip_ctx * ctx, int n_threads, const clip_image_f32 set_input_f32("pos_emb", pos_emb); } } break; + case PROJECTOR_TYPE_MIMO_AUDIO: + { + GGML_ASSERT(imgs.entries.size() == 1); + const int n_frames = imgs.entries.front().nx(); + const int n_pos = (n_frames - 1) / 2 + 1; // matches conv1(s=1)+conv2(s=2) output length + + std::vector positions(n_pos); + for (int i = 0; i < n_pos; i++) { + positions[i] = i; + } + set_input_i32("mimo_audio_positions", positions); + + const int window = hparams.attn_window_size; + GGML_ASSERT(window > 0); + + const float neg_inf = std::numeric_limits::lowest(); + std::vector full_mask((size_t) n_pos * n_pos); + std::vector window_mask((size_t) n_pos * n_pos); + for (int q = 0; q < n_pos; q++) { + for (int k = 0; k < n_pos; k++) { + const bool causal_ok = k <= q; + full_mask[(size_t) q * n_pos + k] = causal_ok ? 0.0f : neg_inf; + window_mask[(size_t) q * n_pos + k] = (causal_ok && (q - k) <= window) ? 0.0f : neg_inf; + } + } + set_input_f32("mimo_audio_full_mask", full_mask); + set_input_f32("mimo_audio_window_mask", window_mask); + + // input_local_transformer: block-diagonal mask + in-group positions + { + const int n_pos_ds = (n_pos - 2) / 2 + 1; // matches downsample conv (k=2,s=2,p=0) + const int group_size = hparams.audio_local_group_size; + GGML_ASSERT(group_size > 0); + const int n_groups = (n_pos_ds + group_size - 1) / group_size; + const int n_padded = n_groups * group_size; + + std::vector local_positions(n_padded); + for (int i = 0; i < n_padded; i++) { + local_positions[i] = i % group_size; + } + set_input_i32("mimo_audio_local_positions", local_positions); + + std::vector local_mask((size_t) n_padded * n_padded); + for (int q = 0; q < n_padded; q++) { + for (int k = 0; k < n_padded; k++) { + const bool same_group = (q / group_size) == (k / group_size); + local_mask[(size_t) q * n_padded + k] = same_group ? 0.0f : neg_inf; + } + } + set_input_f32("mimo_audio_local_mask", local_mask); + } + } break; case PROJECTOR_TYPE_LFM2A: { GGML_ASSERT(imgs.entries.size() == 1); @@ -4678,6 +4839,8 @@ int clip_n_mmproj_embd(const struct clip_ctx * ctx) { return ctx->model.qf_proj_blocks.size() * ctx->model.hparams.projection_dim; case PROJECTOR_TYPE_GLM4V: return ctx->model.mm_ffn_down_w->ne[1]; + case PROJECTOR_TYPE_MIMO_AUDIO: + return ctx->model.mm_2_w->ne[1]; default: GGML_ABORT("Unknown projector type"); } diff --git a/tools/mtmd/models/mimo-audio.cpp b/tools/mtmd/models/mimo-audio.cpp new file mode 100644 index 000000000000..481b36cc8d60 --- /dev/null +++ b/tools/mtmd/models/mimo-audio.cpp @@ -0,0 +1,218 @@ +#include "models.h" + +ggml_cgraph * clip_graph_mimo_audio::build() { + ggml_tensor * inp = build_inp_raw(1); // [n_frames, n_mel, 1] + + ggml_tensor * cur = ggml_conv_1d_ph(ctx0, model.conv1d_1_w, inp, 1, 1); + cur = ggml_add(ctx0, cur, model.conv1d_1_b); + cur = ggml_gelu_erf(ctx0, cur); + + cur = ggml_conv_1d_ph(ctx0, model.conv1d_2_w, cur, 2, 1); + cur = ggml_add(ctx0, cur, model.conv1d_2_b); + cur = ggml_gelu_erf(ctx0, cur); + + ggml_tensor * inpL = ggml_cont(ctx0, ggml_transpose(ctx0, cur)); // [n_embd, n_pos] + const int64_t n_pos = inpL->ne[1]; + cb(inpL, "after_conv1d", -1); + + GGML_ASSERT((int) hparams.wa_pattern_mode.size() == n_layer); + + ggml_tensor * inp_pos = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_pos); + ggml_set_name(inp_pos, "mimo_audio_positions"); + ggml_set_input(inp_pos); + + ggml_tensor * full_mask = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_pos, n_pos); + ggml_set_name(full_mask, "mimo_audio_full_mask"); + ggml_set_input(full_mask); + + ggml_tensor * window_mask = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_pos, n_pos); + ggml_set_name(window_mask, "mimo_audio_window_mask"); + ggml_set_input(window_mask); + + build_vit_opts opts; + opts.attn_mask_layers.resize(n_layer); + for (int il = 0; il < n_layer; il++) { + opts.attn_mask_layers[il] = hparams.wa_pattern_mode[il] == -1 ? full_mask : window_mask; + } + // the skip connection below must be added before the post-transformer norm, + // so build_vit must not apply that norm itself + opts.skip_post_ln = true; + + // encoder_skip_layer_id=3 (1-indexed) -> capture output of layer index 2 + const int skip_capture_il = 2; + GGML_ASSERT(n_layer > skip_capture_il); + ggml_tensor * skip_hidden = nullptr; + opts.callback_layer_out = [&](ggml_tensor * layer_cur, int il) { + if (il == skip_capture_il) { + skip_hidden = layer_cur; + } + }; + + auto add_pos = [&](ggml_tensor * x, const clip_layer &) { + return ggml_rope_ext(ctx0, x, inp_pos, nullptr, d_head, + GGML_ROPE_TYPE_NEOX, 0, hparams.rope_theta, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f); + }; + + inpL = build_vit(inpL, n_pos, NORM_TYPE_NORMAL, hparams.ffn_op, nullptr, add_pos, opts); + inpL = ggml_reshape_2d(ctx0, inpL, n_embd, n_pos); // build_vit restores a (size-1) batch dim + + GGML_ASSERT(skip_hidden != nullptr); + inpL = ggml_add(ctx0, inpL, skip_hidden); + + inpL = build_norm(inpL, model.post_ln_w, model.post_ln_b, NORM_TYPE_NORMAL, eps, -1); + cb(inpL, "after_transformer", -1); + + // downsample: strided conv (no bias) + gelu + layernorm + { + ggml_tensor * ds = ggml_cont(ctx0, ggml_transpose(ctx0, inpL)); // [n_pos, n_embd] + ds = ggml_conv_1d(ctx0, model.downsample_conv_w, ds, 2, 0, 1); + ds = ggml_gelu_erf(ctx0, ds); + ds = ggml_cont(ctx0, ggml_transpose(ctx0, ds)); // [n_embd, n_pos/2] + ds = build_norm(ds, model.downsample_norm_w, model.downsample_norm_b, NORM_TYPE_NORMAL, eps, -1); + inpL = ds; + } + cb(inpL, "after_downsample", -1); + + // RVQ quantize: codebook ne=[dim, max_bins, n_q] + // quantize input vector to codes (type=I32) + std::vector codes; + { + GGML_ASSERT(model.rvq_codebook != nullptr); + const int64_t dim = model.rvq_codebook->ne[0]; + GGML_ASSERT(dim == inpL->ne[0]); + GGML_ASSERT((int64_t) hparams.rvq_codebook_size.size() == model.rvq_codebook->ne[2]); + + ggml_tensor * residual = inpL; // [dim, n_pos_ds] + + for (size_t q = 0; q < hparams.rvq_codebook_size.size(); q++) { + const int64_t bins = hparams.rvq_codebook_size[q]; + ggml_tensor * codebook_q = ggml_view_2d(ctx0, model.rvq_codebook, dim, bins, + model.rvq_codebook->nb[1], q * model.rvq_codebook->nb[2]); + codebook_q = ggml_cont(ctx0, codebook_q); + + ggml_tensor * codebook_norm = ggml_sum_rows(ctx0, ggml_sqr(ctx0, codebook_q)); // [1, bins] + codebook_norm = ggml_cont(ctx0, ggml_transpose(ctx0, codebook_norm)); // [bins, 1] + + ggml_tensor * dot = ggml_mul_mat(ctx0, codebook_q, residual); // [bins, n_pos_ds] + ggml_tensor * scores = ggml_sub(ctx0, ggml_scale(ctx0, dot, 2.0f), codebook_norm); + + ggml_tensor * idx = ggml_argmax(ctx0, scores); // [n_pos_ds] + codes.push_back(idx); + + ggml_tensor * quant = ggml_get_rows(ctx0, codebook_q, idx); // [dim, n_pos_ds] + residual = ggml_sub(ctx0, residual, quant); + cb(idx, "rvq_code", (int) q); + } + } + + // convert codes to LLM embeddings + ggml_tensor * code_embd_sum = nullptr; + { + GGML_ASSERT(model.mm_a_code_embd != nullptr); + const int64_t dim = model.mm_a_code_embd->ne[0]; + const int64_t vocab = model.mm_a_code_embd->ne[1]; + GGML_ASSERT((int64_t) codes.size() == model.mm_a_code_embd->ne[2]); + GGML_ASSERT(dim == inpL->ne[0]); + + for (size_t i = 0; i < codes.size(); i++) { + ggml_tensor * table_i = ggml_view_2d(ctx0, model.mm_a_code_embd, dim, vocab, + model.mm_a_code_embd->nb[1], i * model.mm_a_code_embd->nb[2]); + table_i = ggml_cont(ctx0, table_i); + + ggml_tensor * embd_i = ggml_get_rows(ctx0, table_i, codes[i]); // [dim, n_pos_ds] + code_embd_sum = code_embd_sum ? ggml_add(ctx0, code_embd_sum, embd_i) : embd_i; + } + cb(code_embd_sum, "code_embd_sum", -1); + } + + // input_local_transformer + // groups of `group_size` consecutive downsampled frames are processed together, attending only within their own group. + // Implemented as a block-diagonal mask + in-group-repeating positions + // (rather than a real batch dim) - same technique as the encoder's masks above, and as gemma4a's / deepseekocr2's chunked attention. + + // note: hand-rolled here instead of build_vit() because this is a second, independent layer stack + // (own layer array/count, RMSNorm instead of LN, SiLU FFN, own RoPE theta) + + ggml_tensor * projected; + { + const int group_size = hparams.audio_local_group_size; + GGML_ASSERT(group_size > 0); + const int64_t n_pos_ds = code_embd_sum->ne[1]; + const int64_t n_groups = (n_pos_ds + group_size - 1) / group_size; + const int64_t n_padded = n_groups * group_size; + + ggml_tensor * cur_local = code_embd_sum; + if (n_padded != n_pos_ds) { + cur_local = ggml_pad(ctx0, cur_local, 0, (int) (n_padded - n_pos_ds), 0, 0); + } + + ggml_tensor * local_pos = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_padded); + ggml_set_name(local_pos, "mimo_audio_local_positions"); + ggml_set_input(local_pos); + + ggml_tensor * local_mask = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_padded, n_padded); + ggml_set_name(local_mask, "mimo_audio_local_mask"); + ggml_set_input(local_mask); + + const float local_rope_theta = 640000.0f; // audio_config.rope_theta (differs from the encoder's) + auto apply_local_rope = [&](ggml_tensor * x) { + return ggml_rope_ext(ctx0, x, local_pos, nullptr, d_head, + GGML_ROPE_TYPE_NEOX, 0, local_rope_theta, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f); + }; + + for (int il = 0; il < hparams.audio_local_n_layer; il++) { + auto & layer = model.mm_a_local_layers[il]; + + ggml_tensor * attn_in = build_norm(cur_local, layer.ln_1_w, nullptr, NORM_TYPE_RMS, eps, il); + + ggml_tensor * Qcur = build_mm(layer.q_w, attn_in); + if (layer.q_b) { + Qcur = ggml_add(ctx0, Qcur, layer.q_b); + } + ggml_tensor * Kcur = build_mm(layer.k_w, attn_in); + if (layer.k_b) { + Kcur = ggml_add(ctx0, Kcur, layer.k_b); + } + ggml_tensor * Vcur = build_mm(layer.v_w, attn_in); + if (layer.v_b) { + Vcur = ggml_add(ctx0, Vcur, layer.v_b); + } + + Qcur = ggml_reshape_3d(ctx0, Qcur, d_head, n_head, n_padded); + Kcur = ggml_reshape_3d(ctx0, Kcur, d_head, n_head, n_padded); + Vcur = ggml_reshape_3d(ctx0, Vcur, d_head, n_head, n_padded); + + Qcur = apply_local_rope(Qcur); + Kcur = apply_local_rope(Kcur); + + ggml_tensor * attn_out = build_attn(layer.o_w, nullptr, Qcur, Kcur, Vcur, local_mask, kq_scale, il); + cur_local = ggml_add(ctx0, cur_local, attn_out); + + ggml_tensor * ffn_in = build_norm(cur_local, layer.ln_2_w, nullptr, NORM_TYPE_RMS, eps, il); + ggml_tensor * ffn_out = build_ffn(ffn_in, + layer.ff_up_w, nullptr, + layer.ff_gate_w, nullptr, + layer.ff_down_w, nullptr, + FFN_SILU, il); + cur_local = ggml_add(ctx0, cur_local, ffn_out); + } + + cur_local = build_norm(cur_local, model.mm_a_local_norm_w, nullptr, NORM_TYPE_RMS, eps, -1); + cb(cur_local, "after_local_transformer", -1); + + // flatten each group of `group_size` frames into one (group_size*n_embd)-dim vector + // (matching AudioProjection's flattened input) + ggml_tensor * grouped = ggml_reshape_2d(ctx0, cur_local, n_embd * group_size, n_groups); + + // AudioProjection: Linear (no bias) -> GELU -> Linear (no bias) + projected = build_ffn(grouped, + model.mm_1_w, nullptr, + nullptr, nullptr, + model.mm_2_w, nullptr, + FFN_GELU_ERF, -1); + cb(projected, "after_projection", -1); + } + + ggml_build_forward_expand(gf, projected); + return gf; +} diff --git a/tools/mtmd/models/models.h b/tools/mtmd/models/models.h index 2d7555da41d2..caed438ec513 100644 --- a/tools/mtmd/models/models.h +++ b/tools/mtmd/models/models.h @@ -210,6 +210,11 @@ struct clip_graph_qwen3a : clip_graph { ggml_cgraph * build() override; }; +struct clip_graph_mimo_audio : clip_graph { + clip_graph_mimo_audio(clip_ctx * ctx, const clip_image_f32 & img) : clip_graph(ctx, img) {} + ggml_cgraph * build() override; +}; + struct clip_graph_kimik25 : clip_graph { clip_graph_kimik25(clip_ctx * ctx, const clip_image_f32 & img) : clip_graph(ctx, img) {} ggml_cgraph * build() override; diff --git a/tools/mtmd/mtmd-audio.cpp b/tools/mtmd/mtmd-audio.cpp index b72fd067a508..ed68951c0151 100644 --- a/tools/mtmd/mtmd-audio.cpp +++ b/tools/mtmd/mtmd-audio.cpp @@ -725,6 +725,72 @@ bool mtmd_audio_preprocessor_qwen3a::preprocess(const float * sa return true; } +// +// mtmd_audio_preprocessor_mimo_audio +// +// Matches torchaudio.transforms.MelSpectrogram(power=1.0, center=True) followed by +// log(clip(spec, min=1e-7)): HTK mel scale, no Slaney area norm, magnitude (not power) +// spectrogram, natural log, reflect-padded by n_fft/2 on each side. +// + +void mtmd_audio_preprocessor_mimo_audio::initialize() { + cache.fill_sin_cos_table(hparams.audio_n_fft); + cache.fill_hann_window(hparams.audio_window_len, true); + cache.fill_mel_filterbank_matrix( + hparams.n_mel_bins, hparams.audio_n_fft, hparams.audio_sample_rate, + 0.0f, hparams.audio_sample_rate / 2.0f, + /*slaney_area_norm=*/ false, + /*scale=*/ 1.0f, + /*use_htk=*/ true + ); +} + +bool mtmd_audio_preprocessor_mimo_audio::preprocess(const float * samples, + size_t n_samples, + std::vector & output) { + if (n_samples == 0) { + return false; + } + + GGML_ASSERT(!cache.sin_vals.empty()); + GGML_ASSERT(!cache.cos_vals.empty()); + GGML_ASSERT(!cache.filters.data.empty()); + + const int pad = hparams.audio_n_fft / 2; + + std::vector padded(n_samples + 2 * pad, 0.0f); + for (int i = 0; i < pad; i++) { + int src = pad - i; + padded[i] = (src < (int)n_samples) ? samples[src] : 0.0f; + } + std::copy(samples, samples + n_samples, padded.begin() + pad); + for (int i = 0; i < pad; i++) { + int src = (int)n_samples - 2 - i; + padded[n_samples + pad + i] = (src >= 0) ? samples[src] : 0.0f; + } + + filter_params params; + params.n_mel = hparams.n_mel_bins; + params.n_fft_bins = 1 + (hparams.audio_n_fft / 2); + params.hann_window_size = hparams.audio_window_len; + params.hop_length = hparams.audio_hop_len; + params.sample_rate = hparams.audio_sample_rate; + params.no_padding = true; // reflect padding already applied above + params.use_natural_log = true; + params.use_magnitude = true; + params.mel_floor = 1e-7f; + params.norm_per_feature = false; + + mtmd_audio_mel out; + bool ok = log_mel_spectrogram(padded.data(), (int)padded.size(), 4, params, cache, out); + if (!ok) { + return false; + } + + output.push_back(std::move(out)); + return true; +} + // // mtmd_audio_preprocessor_conformer // diff --git a/tools/mtmd/mtmd-audio.h b/tools/mtmd/mtmd-audio.h index ad96bd847cfc..d8ec72b9d54e 100644 --- a/tools/mtmd/mtmd-audio.h +++ b/tools/mtmd/mtmd-audio.h @@ -111,6 +111,15 @@ struct mtmd_audio_preprocessor_qwen3a : mtmd_audio_preprocessor { mtmd_audio_cache cache; }; +struct mtmd_audio_preprocessor_mimo_audio : mtmd_audio_preprocessor { + mtmd_audio_preprocessor_mimo_audio(const clip_ctx * ctx) : mtmd_audio_preprocessor(ctx) {} + void initialize() override; + bool preprocess(const float * samples, size_t n_samples, std::vector & output) override; + + private: + mtmd_audio_cache cache; +}; + // // streaming ISTFT - converts spectrogram frames back to audio one frame at a time // diff --git a/tools/mtmd/mtmd.cpp b/tools/mtmd/mtmd.cpp index bb49b211efb3..6e61cf3e520b 100644 --- a/tools/mtmd/mtmd.cpp +++ b/tools/mtmd/mtmd.cpp @@ -730,6 +730,12 @@ struct mtmd_context { aud_end = ""; audio_preproc = std::make_unique(ctx_a); } break; + case PROJECTOR_TYPE_MIMO_AUDIO: + { + aud_beg = "<|mimo_audio_start|>"; + aud_end = "<|mimo_audio_end|>"; + audio_preproc = std::make_unique(ctx_a); + } break; default: throw std::runtime_error(string_format("%s: unexpected audio projector type %d\n", __func__, proj)); } From 91f8c9c5fb038c086e13e9cd823c29b33b07ba54 Mon Sep 17 00:00:00 2001 From: Beinsezii <39478211+Beinsezii@users.noreply.github.com> Date: Mon, 27 Jul 2026 16:13:48 -0700 Subject: [PATCH 028/190] Disable -ffast-math on HIP (#25495) --- ggml/src/ggml-hip/CMakeLists.txt | 2 -- 1 file changed, 2 deletions(-) diff --git a/ggml/src/ggml-hip/CMakeLists.txt b/ggml/src/ggml-hip/CMakeLists.txt index 5351dcae12db..bbc51797c182 100644 --- a/ggml/src/ggml-hip/CMakeLists.txt +++ b/ggml/src/ggml-hip/CMakeLists.txt @@ -154,5 +154,3 @@ if (GGML_HIP_RCCL) endif() target_link_libraries(ggml-hip PRIVATE ggml-base hip::host roc::rocblas roc::hipblas) - -target_compile_options(ggml-hip PRIVATE "$<$:-ffast-math;-fno-finite-math-only>") From c6292cfb8e964bd9f2146f3c603507627e4a22c3 Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Tue, 28 Jul 2026 08:41:04 +0300 Subject: [PATCH 029/190] contrib : add guideline about the "merge ready" label (#26178) * contrib : add guideline about the "merge ready" label * cont : add ref [no ci] --- CONTRIBUTING.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 91fa381dd019..003133478811 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -73,6 +73,7 @@ For more info, please refer to the [AGENTS.md](AGENTS.md) file. - When merging a PR, make sure you have a good understanding of the changes - If a PR does not warrant a new release, add `[no release]` in the squashed commit to spare CI resources - Be mindful of maintenance: most of the work going into a feature happens after the PR is merged. If the PR author is not committed to contribute long-term, someone else needs to take responsibility (you) +- Add the ["merge ready"](https://github.com/ggml-org/llama.cpp/pulls?q=is%3Apr+is%3Aopen+draft%3Ano+sort%3Aupdated-desc+label%3A%22merge+ready%22+) label to a PR to indicate when a PR can be fast-merged without waiting for 2 independent reviews. [(more info)](https://github.com/ggml-org/llama.cpp/pull/26178) Maintainers reserve the right to decline review or close pull requests for any reason, without any questions, particularly under any of the following conditions: - The proposed change is already mentioned in the roadmap or an existing issue, and it has been assigned to someone. From f87067841bac583bc089a225382248d857791ca8 Mon Sep 17 00:00:00 2001 From: Ruixiang Wang Date: Tue, 28 Jul 2026 08:58:16 +0200 Subject: [PATCH 030/190] spec: add eagle3-v3 support for gpt-oss model (#25794) --- common/speculative.cpp | 16 +++++++++++++--- conversion/llama.py | 18 ++++++++++++++++-- gguf-py/gguf/constants.py | 2 ++ gguf-py/gguf/gguf_writer.py | 3 +++ src/llama-arch.cpp | 1 + src/llama-arch.h | 1 + src/llama-hparams.h | 1 + src/models/eagle3.cpp | 15 +++++++++++++++ src/models/openai-moe.cpp | 8 +++++++- 9 files changed, 59 insertions(+), 6 deletions(-) diff --git a/common/speculative.cpp b/common/speculative.cpp index ee94d7c37662..3a6c07368358 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -437,6 +437,7 @@ struct common_speculative_impl_draft_eagle3 : public common_speculative_impl { int32_t n_embd_dec = 0; // draft hidden size int32_t n_embd_enc = 0; // target_layer_ids_n * target_hidden_size int32_t n_embd_tgt = 0; // target model hidden size + int32_t n_layer_tgt = 0; // target model layer count const int32_t * target_layer_ids = nullptr; // model_dft's extract layer indices uint32_t target_layer_ids_n = 0; @@ -478,6 +479,7 @@ struct common_speculative_impl_draft_eagle3 : public common_speculative_impl { n_embd_tgt = llama_model_n_embd(model_tgt); n_embd_dec = llama_model_n_embd(model_dft); n_embd_enc = (int32_t) target_layer_ids_n * n_embd_tgt; + n_layer_tgt = llama_model_n_layer(model_tgt); const int32_t n_b = (int32_t) llama_n_batch(ctx_dft); batch = llama_batch_init(/*n_tokens=*/ n_b, /*embd=*/ n_embd_dec, /*n_seq_max=*/ 1); @@ -510,9 +512,15 @@ struct common_speculative_impl_draft_eagle3 : public common_speculative_impl { } } - // turn on extraction of the target layers' input embeddings + // turn on extraction of the target layers' hidden states for (uint32_t k = 0; k < target_layer_ids_n; ++k) { - llama_set_embeddings_layer_inp(ctx_tgt, (uint32_t) target_layer_ids[k], true); + if (target_layer_ids[k] < n_layer_tgt) { + llama_set_embeddings_layer_inp(ctx_tgt, (uint32_t) target_layer_ids[k], true); + } else if (target_layer_ids[k] == n_layer_tgt) { + llama_set_embeddings_nextn(ctx_tgt, true, /*masked*/ false); + } else { + GGML_ABORT("EAGLE3: target layer id %d exceeds target n_layer %d", target_layer_ids[k], n_layer_tgt); + } } // turn on extraction of the draft model's pre-norm hidden state @@ -600,7 +608,9 @@ struct common_speculative_impl_draft_eagle3 : public common_speculative_impl { features_buf.resize((size_t) n_tokens * n_embd_enc, 0.0f); for (uint32_t k = 0; k < target_layer_ids_n; ++k) { - const float * layer = llama_get_embeddings_layer_inp(ctx_tgt, (uint32_t) target_layer_ids[k]); + const float * layer = target_layer_ids[k] < n_layer_tgt + ? llama_get_embeddings_layer_inp(ctx_tgt, (uint32_t) target_layer_ids[k]) + : llama_get_embeddings_nextn(ctx_tgt); if (!layer) { GGML_ABORT("EAGLE3: target layer %d input not extracted.", target_layer_ids[k]); } diff --git a/conversion/llama.py b/conversion/llama.py index 315a619c9c26..9b3373f911e9 100644 --- a/conversion/llama.py +++ b/conversion/llama.py @@ -69,9 +69,14 @@ def __init__(self, *args, **kwargs): target_config = {**target_config, **target_config["text_config"]} self.target_vocab_size = target_config["vocab_size"] - # target_layers: derived from target model layer count (low/mid/high) + # target_layers: use the eagle3 config's explicit aux hidden-state layer ids + # if present, else derive from the target layer count. target_num_layers = target_config["num_hidden_layers"] - target_layers = [2, target_num_layers // 2, target_num_layers - 3] + aux_layer_ids = eagle3_raw_config.get("eagle_aux_hidden_state_layer_ids") + if aux_layer_ids: + target_layers = aux_layer_ids + else: + target_layers = [2, target_num_layers // 2, target_num_layers - 3] logger.info(f"EAGLE-3: target_layers = {target_layers} (target model has {target_num_layers} layers)") self.gguf_writer.add_target_layers(target_layers) @@ -90,6 +95,12 @@ def __init__(self, *args, **kwargs): logger.info(f"EAGLE-3: norm_before_residual = {norm_before_residual}") self.gguf_writer.add_norm_before_residual(norm_before_residual) + # norm_before_fc: RMSNorm applied to the fused target features before the + # fc projection (e.g. nvidia/gpt-oss-120b-Eagle3-v3) + norm_before_fc = eagle3_raw_config.get("norm_before_fc", False) + logger.info(f"EAGLE-3: norm_before_fc = {norm_before_fc}") + self.gguf_writer.add_norm_before_fc(norm_before_fc) + def set_vocab(self): # eagle3: use tokenizer from target model if provided original_dir_model = None @@ -222,6 +233,9 @@ def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iter if name == "fc.weight": yield (name, data_torch) return + if name == "input_norm.weight": + yield (self.format_tensor_name(gguf.MODEL_TENSOR.ENC_OUTPUT_NORM), data_torch) + return if name == "d2t": # store for manual int64 handling in prepare_tensors (avoid F32 conversion) if not hasattr(self, '_eagle3_int_tensors'): diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index f9264425f6e1..4df94b9a3add 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -161,6 +161,7 @@ class LLM: TARGET_HIDDEN_SIZE = "{arch}.target_hidden_size" BLOCK_SIZE = "{arch}.block_size" NORM_BEFORE_RESIDUAL = "{arch}.norm_before_residual" + NORM_BEFORE_FC = "{arch}.norm_before_fc" class Attention: HEAD_COUNT = "{arch}.attention.head_count" @@ -4343,6 +4344,7 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.FFN_DOWN, MODEL_TENSOR.FFN_UP, MODEL_TENSOR.FC, + MODEL_TENSOR.ENC_OUTPUT_NORM, MODEL_TENSOR.D2T, ], MODEL_ARCH.DFLASH: [ diff --git a/gguf-py/gguf/gguf_writer.py b/gguf-py/gguf/gguf_writer.py index bd8629aa119e..657ed69b6895 100644 --- a/gguf-py/gguf/gguf_writer.py +++ b/gguf-py/gguf/gguf_writer.py @@ -971,6 +971,9 @@ def add_target_hidden_size(self, value: int) -> None: def add_norm_before_residual(self, value: bool) -> None: self.add_bool(Keys.LLM.NORM_BEFORE_RESIDUAL.format(arch=self.arch), value) + def add_norm_before_fc(self, value: bool) -> None: + self.add_bool(Keys.LLM.NORM_BEFORE_FC.format(arch=self.arch), value) + def add_attention_output_group_count(self, count: int) -> None: self.add_uint32(Keys.Attention.OUTPUT_GROUP_COUNT.format(arch=self.arch), count) diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index c01706785070..a0945e501d53 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -317,6 +317,7 @@ static const std::map LLM_KV_NAMES = { { LLM_KV_TARGET_LAYERS, "%s.target_layers" }, { LLM_KV_TARGET_HIDDEN_SIZE, "%s.target_hidden_size" }, { LLM_KV_NORM_BEFORE_RESIDUAL, "%s.norm_before_residual" }, + { LLM_KV_NORM_BEFORE_FC, "%s.norm_before_fc" }, { LLM_KV_SHORTCONV_L_CACHE, "%s.shortconv.l_cache" }, // sentence-transformers dense modules feature dims diff --git a/src/llama-arch.h b/src/llama-arch.h index 1c9aebb0bbdc..fb5eb3920bcd 100644 --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -363,6 +363,7 @@ enum llm_kv { LLM_KV_TARGET_LAYERS, LLM_KV_TARGET_HIDDEN_SIZE, LLM_KV_NORM_BEFORE_RESIDUAL, + LLM_KV_NORM_BEFORE_FC, LLM_KV_SHORTCONV_L_CACHE, diff --git a/src/llama-hparams.h b/src/llama-hparams.h index 727df6ca21e2..fc770bf003e6 100644 --- a/src/llama-hparams.h +++ b/src/llama-hparams.h @@ -47,6 +47,7 @@ struct llama_hparams { bool use_par_res; bool swin_norm; bool norm_before_residual = false; + bool norm_before_fc = false; uint32_t n_ctx_train; // context size the model was trained on uint32_t n_embd; diff --git a/src/models/eagle3.cpp b/src/models/eagle3.cpp index 9d96fae5944e..be466056df69 100644 --- a/src/models/eagle3.cpp +++ b/src/models/eagle3.cpp @@ -28,6 +28,10 @@ void llama_model_eagle3::load_arch_hparams(llama_model_loader & ml) { LLAMA_LOG_INFO("%s: EAGLE3gnorm_before_residual = true\n", __func__); } + // eagle3 norm_before_fc (optional, default false) + // compatible with eagle3.1 (e.g. nvidia/gpt-oss-120b-Eagle3-v3) + ml.get_key(LLM_KV_NORM_BEFORE_FC, hparams.norm_before_fc, false); + type = LLM_TYPE_UNKNOWN; } @@ -53,6 +57,11 @@ void llama_model_eagle3::load_arch_tensors(llama_model_loader &) { // Feature fusion layer: projects 3 target layers to draft hidden size fc = create_tensor(tn(LLM_TENSOR_FC, "weight"), {n_embd_inp, n_embd}, 0); + // RMSNorm on the fused target features (input to fc), only when norm_before_fc is set. + if (hparams.norm_before_fc) { + output_norm_enc = create_tensor(tn(LLM_TENSOR_ENC_OUTPUT_NORM, "weight"), {n_embd_inp}, 0); + } + // Output layer (uses draft vocab size) output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0); output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_draft_vocab}, TENSOR_NOT_REQUIRED); @@ -130,6 +139,12 @@ llama_model_eagle3::graph::graph(const llama_model & model, const llm_grap cur = build_inp_embd_enc(); + // RMSNorm on the fused target features before fc + if (hparams.norm_before_fc) { + cur = build_norm(cur, model.output_norm_enc, NULL, LLM_NORM_RMS, -1); + cb(cur, "enc_input_norm", -1); + } + // Feature fusion layer cur = build_lora_mm(model.fc, cur); cb(cur, "fc_out", -1); diff --git a/src/models/openai-moe.cpp b/src/models/openai-moe.cpp index 6d74f9c7e6ef..c91bae1c35c6 100644 --- a/src/models/openai-moe.cpp +++ b/src/models/openai-moe.cpp @@ -116,7 +116,7 @@ llama_model_openai_moe::graph::graph(const llama_model & model, const llm_graph_ cb(cur, "attn_out", il); } - if (il == n_layer - 1) { + if (il == n_layer - 1 && inp_out_ids && cparams.embeddings_nextn_masked) { // skip computing output for unused tokens cur = ggml_get_rows(ctx0, cur, inp_out_ids); inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids); @@ -154,6 +154,12 @@ llama_model_openai_moe::graph::graph(const llama_model & model, const llm_graph_ } cur = inpL; + res->t_h_nextn = cur; + + if (!cparams.embeddings_nextn_masked && inp_out_ids) { + cur = ggml_get_rows(ctx0, cur, inp_out_ids); + } + cur = build_norm(cur, model.output_norm, NULL, LLM_NORM_RMS, -1); From f95de9776b5b90dd993f36d2bd66a3eee21c887f Mon Sep 17 00:00:00 2001 From: Nick Lafleur <55208706+nicklafleur@users.noreply.github.com> Date: Tue, 28 Jul 2026 03:44:06 -0400 Subject: [PATCH 031/190] ggml-metal: FWHT kernel for metal backend (#25924) * metal fwht wip * shape guard and formatting * formatting * Formatting and typos Co-authored-by: YiChen Lv <63285796+forforever73@users.noreply.github.com> * fix narrowing issue Co-authored-by: YiChen Lv <63285796+forforever73@users.noreply.github.com> * cont : minor style --------- Co-authored-by: YiChen Lv <63285796+forforever73@users.noreply.github.com> Co-authored-by: Georgi Gerganov --- ggml/src/ggml-metal/ggml-metal-device.cpp | 15 ++++++ ggml/src/ggml-metal/ggml-metal-device.h | 1 + ggml/src/ggml-metal/ggml-metal-impl.h | 4 ++ ggml/src/ggml-metal/ggml-metal-ops.cpp | 52 ++++++++++++++++++ ggml/src/ggml-metal/ggml-metal-ops.h | 1 + ggml/src/ggml-metal/ggml-metal.metal | 64 ++++++++++++++++++++++- tests/test-backend-ops.cpp | 7 +++ 7 files changed, 143 insertions(+), 1 deletion(-) diff --git a/ggml/src/ggml-metal/ggml-metal-device.cpp b/ggml/src/ggml-metal/ggml-metal-device.cpp index 270c1411a059..16e98eb519ea 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.cpp +++ b/ggml/src/ggml-metal/ggml-metal-device.cpp @@ -1252,6 +1252,21 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_argsort_merge(gg return res; } +ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_fwht(ggml_metal_library_t lib, int n) { + char base[256]; + char name[256]; + + snprintf(base, 256, "kernel_fwht_f32_%d", n); + snprintf(name, 256, "%s", base); + + ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); + if (!res.pipeline) { + res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr); + } + + return res; +} + // note: reuse the argsort kernel for top_k ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_top_k(ggml_metal_library_t lib, const ggml_tensor * op) { assert(op->op == GGML_OP_TOP_K); diff --git a/ggml/src/ggml-metal/ggml-metal-device.h b/ggml/src/ggml-metal/ggml-metal-device.h index b36fa8110b57..d0956df50675 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.h +++ b/ggml/src/ggml-metal/ggml-metal-device.h @@ -139,6 +139,7 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_argmax (ggml_metal_library_t lib, const struct ggml_tensor * op); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_argsort (ggml_metal_library_t lib, const struct ggml_tensor * op); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_argsort_merge (ggml_metal_library_t lib, const struct ggml_tensor * op); +struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_fwht (ggml_metal_library_t lib, int n); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_top_k (ggml_metal_library_t lib, const struct ggml_tensor * op); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_top_k_merge (ggml_metal_library_t lib, const struct ggml_tensor * op); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_bin (ggml_metal_library_t lib, const struct ggml_tensor * op, int32_t n_fuse ); diff --git a/ggml/src/ggml-metal/ggml-metal-impl.h b/ggml/src/ggml-metal/ggml-metal-impl.h index 330278d003df..9f350aad5b74 100644 --- a/ggml/src/ggml-metal/ggml-metal-impl.h +++ b/ggml/src/ggml-metal/ggml-metal-impl.h @@ -1157,6 +1157,10 @@ typedef struct { int32_t len; } ggml_metal_kargs_argsort_merge; +typedef struct { + int32_t nrows; +} ggml_metal_kargs_fwht; + typedef struct { int64_t ne0; float start; diff --git a/ggml/src/ggml-metal/ggml-metal-ops.cpp b/ggml/src/ggml-metal/ggml-metal-ops.cpp index c716f118f6d3..76626a451836 100644 --- a/ggml/src/ggml-metal/ggml-metal-ops.cpp +++ b/ggml/src/ggml-metal/ggml-metal-ops.cpp @@ -1979,6 +1979,46 @@ int ggml_metal_op_pool_1d(ggml_metal_op_t ctx, int idx) { return 1; } +// supported FWHT sizes, must stay in sync with the +// kernel_fwht_f32_ templates in ggml-metal.metal +static bool ggml_metal_fwht_supported_size(int64_t n) { + return n == 64 || n == 128 || n == 256 || n == 512; +} + +int ggml_metal_op_fwht(ggml_metal_op_t ctx, int idx) { + ggml_tensor * op = ctx->node(idx); + + ggml_metal_library_t lib = ctx->lib; + ggml_metal_encoder_t enc = ctx->enc; + + ggml_tensor * src1 = op->src[1]; + + const int64_t n = src1->ne[0]; + const int64_t nrows = ggml_nrows(src1); + + ggml_metal_kargs_fwht args = { + /*.nrows = */ (int32_t) nrows, + }; + + auto pipeline = ggml_metal_library_get_pipeline_fwht(lib, n); + + ggml_metal_encoder_set_pipeline(enc, pipeline); + ggml_metal_encoder_set_bytes(enc, &args, sizeof(args), 0); + ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(src1), 1); + ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(op), 2); + + const int th_max = ggml_metal_pipeline_max_theads_per_threadgroup(pipeline); + const int simd_size = 32; + + int sg_per_tg = 2; + sg_per_tg = std::min(sg_per_tg, th_max/simd_size); + sg_per_tg = std::max(sg_per_tg, 1); + + const int64_t n_tg = (nrows + sg_per_tg - 1) / sg_per_tg; + ggml_metal_encoder_dispatch_threadgroups(enc, n_tg, 1, 1, 32*sg_per_tg, 1, 1); + + return 1; +} int ggml_metal_op_pool_2d(ggml_metal_op_t ctx, int idx) { ggml_tensor * op = ctx->node(idx); @@ -2046,6 +2086,18 @@ int ggml_metal_op_mul_mat(ggml_metal_op_t ctx, int idx) { ggml_metal_library_t lib = ctx->lib; ggml_metal_encoder_t enc = ctx->enc; + const int32_t hint = ggml_get_op_params_i32(op, 1); + + if (hint == GGML_HINT_SRC0_IS_HADAMARD) { + if (op->src[1]->type == GGML_TYPE_F32 && + op->type == GGML_TYPE_F32 && + ggml_is_contiguous(op->src[1]) && + ggml_is_contiguous(op) && + ggml_are_same_shape(op->src[1], op) && + ggml_metal_fwht_supported_size(op->src[1]->ne[0])) { + return ggml_metal_op_fwht(ctx, idx); + } + } const ggml_metal_device_props * props_dev = ggml_metal_device_get_props(ctx->dev); GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne); diff --git a/ggml/src/ggml-metal/ggml-metal-ops.h b/ggml/src/ggml-metal/ggml-metal-ops.h index 89a6ad82f1c2..2783ecb8b611 100644 --- a/ggml/src/ggml-metal/ggml-metal-ops.h +++ b/ggml/src/ggml-metal/ggml-metal-ops.h @@ -64,6 +64,7 @@ int ggml_metal_op_set (ggml_metal_op_t ctx, int idx); int ggml_metal_op_cpy (ggml_metal_op_t ctx, int idx); int ggml_metal_op_pool_1d (ggml_metal_op_t ctx, int idx); int ggml_metal_op_pool_2d (ggml_metal_op_t ctx, int idx); +int ggml_metal_op_fwht (ggml_metal_op_t ctx, int idx); int ggml_metal_op_mul_mat (ggml_metal_op_t ctx, int idx); int ggml_metal_op_mul_mat_id (ggml_metal_op_t ctx, int idx); int ggml_metal_op_add_id (ggml_metal_op_t ctx, int idx); diff --git a/ggml/src/ggml-metal/ggml-metal.metal b/ggml/src/ggml-metal/ggml-metal.metal index 969fddfa5b89..f14ee0792ba5 100644 --- a/ggml/src/ggml-metal/ggml-metal.metal +++ b/ggml/src/ggml-metal/ggml-metal.metal @@ -5762,7 +5762,7 @@ kernel void kernel_upscale_bicubic_f32( const float w_y2 = bicubic_weight1(1.0f - fd1); const float w_y3 = bicubic_weight2(2.0f - fd1); - const device const char * src_slice = src0 + i03 * args.nb03 + i02 * args.nb02; + const device char * src_slice = src0 + i03 * args.nb03 + i02 * args.nb02; device float * dst_ptr = (device float *)(dst + i3 * args.nb3 + i2 * args.nb2 + i1 * args.nb1); @@ -6172,6 +6172,68 @@ kernel void kernel_argsort_merge_f32_i32( template [[host_name("kernel_argsort_merge_f32_i32_asc")]] kernel argsort_merge_t kernel_argsort_merge_f32_i32; template [[host_name("kernel_argsort_merge_f32_i32_desc")]] kernel argsort_merge_t kernel_argsort_merge_f32_i32; +template +kernel void kernel_fwht_f32( + constant ggml_metal_kargs_fwht & args, + device const float * src, + device float * dst, + uint3 tgpig[[threadgroup_position_in_grid]], + ushort sgitg[[simdgroup_index_in_threadgroup]], + ushort tiisg[[thread_index_in_simdgroup]], + ushort3 ntg[[threads_per_threadgroup]]) { + + constexpr int NW = N_SIMDWIDTH; + constexpr int NE = N / NW; + + const float scale = 1.0f / sqrt((float) N); + + const int sg_per_tg = ntg.x / NW; + const int64_t r = tgpig.x * sg_per_tg + sgitg; + if (r >= args.nrows) { + return; + } + + src += r * N; + dst += r * N; + + const int lane = tiisg; + + float reg[NE]; + for (int i = 0; i < NE; i++) { + reg[i] = src[i*NW + lane]*scale; + } + for (int i = 1; i < NW; i *= 2) { + for (int j = 0; j < NE; j++) { + const float val = reg[j]; + const float val2 = simd_shuffle_xor(val, i); + reg[j] = (lane & i) == 0 ? val2 + val : val2 - val; + } + } + + for (int i = NW; i < N; i *= 2) { + const int step = i / NW; + for (int j = 0; j < NE; j += (2 * step)) { + for (int k = 0; k < step; k++) { + const float x = reg[j + k ]; + const float y = reg[j + k + step]; + reg[j + k] = x + y; + reg[j + k + step] = x - y; + } + } + } + + for (int i = 0; i < NE; i++) { + dst[i*NW + lane] = reg[i]; + } +} + +typedef decltype(kernel_fwht_f32<64>) kernel_fwht_t; + +template [[host_name("kernel_fwht_f32_64")]] kernel kernel_fwht_t kernel_fwht_f32<64>; +template [[host_name("kernel_fwht_f32_128")]] kernel kernel_fwht_t kernel_fwht_f32<128>; +template [[host_name("kernel_fwht_f32_256")]] kernel kernel_fwht_t kernel_fwht_f32<256>; +template [[host_name("kernel_fwht_f32_512")]] kernel kernel_fwht_t kernel_fwht_f32<512>; + constant bool FC_flash_attn_ext_pad_has_mask [[function_constant(FC_FLASH_ATTN_EXT_PAD + 0)]]; constant int32_t FC_flash_attn_ext_pad_ncpsg [[function_constant(FC_FLASH_ATTN_EXT_PAD + 25)]]; diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index e7cd6d0cb668..b4061e35a050 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -8793,6 +8793,9 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_mul_mat_hadamard(GGML_TYPE_F32, GGML_TYPE_F32, 512, 1, 512)); test_cases.emplace_back(new test_mul_mat_hadamard(GGML_TYPE_F32, GGML_TYPE_F32, 128, 32, 128)); test_cases.emplace_back(new test_mul_mat_hadamard(GGML_TYPE_F32, GGML_TYPE_F32, 128, 4, 128, {2, 3})); + test_cases.emplace_back(new test_mul_mat_hadamard(GGML_TYPE_F32, GGML_TYPE_F32, 256, 512, 256)); // many rows + test_cases.emplace_back(new test_mul_mat_hadamard(GGML_TYPE_F32, GGML_TYPE_F32, 32, 1, 32)); // too small (N<64) + test_cases.emplace_back(new test_mul_mat_hadamard(GGML_TYPE_F32, GGML_TYPE_F32, 1024, 1, 1024)); // too big (N>512) #if 0 // > 4GB A matrix. Too slow to be enabled by default. @@ -9803,6 +9806,10 @@ static std::vector> make_test_cases_perf() { test_cases.emplace_back(new test_mul_mat_hadamard(GGML_TYPE_F32, GGML_TYPE_F32, 64, 1, 64)); test_cases.emplace_back(new test_mul_mat_hadamard(GGML_TYPE_F32, GGML_TYPE_F32, 256, 1, 256)); test_cases.emplace_back(new test_mul_mat_hadamard(GGML_TYPE_F32, GGML_TYPE_F32, 128, 32, 128)); + test_cases.emplace_back(new test_mul_mat_hadamard(GGML_TYPE_F32, GGML_TYPE_F32, 64, 2048, 64)); + test_cases.emplace_back(new test_mul_mat_hadamard(GGML_TYPE_F32, GGML_TYPE_F32, 128, 2048, 128)); + test_cases.emplace_back(new test_mul_mat_hadamard(GGML_TYPE_F32, GGML_TYPE_F32, 256, 2048, 256)); + test_cases.emplace_back(new test_mul_mat_hadamard(GGML_TYPE_F32, GGML_TYPE_F32, 512, 2048, 512)); test_cases.emplace_back(new test_solve_tri(GGML_TYPE_F32, { 64, 64, 4, 4 }, { 32, 64, 4, 4 })); test_cases.emplace_back(new test_solve_tri(GGML_TYPE_F32, { 128, 128, 4, 2 }, { 32, 128, 4, 2 })); From 9a3bf2b84923a85583b4ee8177b0cca13824bb03 Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Tue, 28 Jul 2026 11:05:16 +0300 Subject: [PATCH 032/190] server : add extra trace log for prompt similarity (#26218) --- tools/server/server-task.cpp | 2 ++ tools/server/server-task.h | 2 +- 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/tools/server/server-task.cpp b/tools/server/server-task.cpp index 1fd7cce27bb3..99e63b05f544 100644 --- a/tools/server/server-task.cpp +++ b/tools/server/server-task.cpp @@ -1755,6 +1755,8 @@ bool server_prompt_cache::load(server_prompt & prompt, const server_tokens & tok const float f_keep_cur = float(lcp_cur) / it->prompt.tokens.size(); const float sim_cur = float(lcp_cur) / tokens_new.size(); + SRV_TRC(" - prompt with length %7zu, lcp = %7d, f_keep = %.3f, sim = %.3f\n", it->prompt.tokens.size(), lcp_cur, f_keep_cur, sim_cur); + // don't trash large prompts if (f_keep_cur < 0.25f) { continue; diff --git a/tools/server/server-task.h b/tools/server/server-task.h index c3eea2ecb81b..411d91807955 100644 --- a/tools/server/server-task.h +++ b/tools/server/server-task.h @@ -650,7 +650,7 @@ struct server_prompt_cache { server_prompt_cache_state * alloc(const server_prompt & prompt, size_t state_size_main, size_t state_size_drft); - bool load(server_prompt & prompt, const server_tokens & tokens_new, llama_context * ctx_main, llama_context * ctx_drft, int32_t id_slot); + bool load(server_prompt & prompt, const server_tokens & tokens_new, llama_context * ctx_tgt, llama_context * ctx_dft, int32_t id_slot); void update(); }; From d6b61ac0d361acbb0b1ac2e160cb8698baaddd6e Mon Sep 17 00:00:00 2001 From: meatposes Date: Tue, 28 Jul 2026 03:37:25 -0500 Subject: [PATCH 033/190] sycl: fix use-after-return of the SDPA scale in the oneDNN flash-attention path (#25880) * sycl: fix use-after-return of the SDPA scale in the oneDNN flash-attention path The scale was uploaded with an async memcpy sourced from a stack local. On the in-order queue that copy is ordered behind the K/V staging kernels; once n_kv is large enough (>= ~26k observed on Arc Pro B70) the staging outlives the host stack frame and the copy reads recycled memory, feeding the SDPA a garbage scale. Output then collapses to a single repeated token and the KV cache is poisoned for the rest of the session. Short contexts win the race by accident, and test-backend-ops caps FLASH_ATTN_EXT at kv=1024, which is why CI never caught it. The previous device_count > 1 wait_and_throw() gate (and reverting it, PR #25741) fixes the symptom only by keeping the frame alive across the copy at the cost of a host sync on every FA call. Fix: cache one device scalar per (device, value) -- the scale is constant per model -- and upload it synchronously once. The single-device fast path (no per-call host sync) is then safe: every device-side hazard already serializes on the in-order queue. The multi-GPU conservative wait is kept unchanged. Also: - GGML_SYCL_FA_ONEDNN_MAX_KV env (0 = unlimited): optional n_kv ceiling that routes very long sequences to the native FA kernel. - test-backend-ops: FLASH_ATTN_EXT F16 cases up to kv=65536 (Qwen3.6-27B geometry hsk=hsv=256 GQA 6, and hsk=128 GQA 4), closing the kv=1024 blind spot. Note the race itself needs a live multi-op pipeline to reproduce; single-op runs pass even on broken builds. Verified on Arc Pro B70 (bmg_g31), Qwen3.6-27B Q4_K, -c 131072: output byte-identical at temp 0 to the native FA path through 32k-deep prefill, with prefill depth-flat at 820-840 t/s (vs 340-350 native at 32k depth). Assisted-by: Claude Fable 5 * sycl: handle GGML_SYCL_FA_ONEDNN_MAX_KV like the other runtime env vars and document it Review feedback on #25880: - read the variable once at backend init into g_ggml_sycl_fa_onednn_max_kv via ggml_sycl_get_env, and print it in the startup env listing (-lv 4 shows it) - document GGML_SYCL_FA_ONEDNN and GGML_SYCL_FA_ONEDNN_MAX_KV in the SYCL.md runtime table Also trim the added FLASH_ATTN_EXT cases to kv={4096,16384}: the 32768/65536 shapes exceed the legacy NMSE threshold on both the oneDNN and native kernels (long-sequence fp16 accumulation drift, present before this PR) and would fail CI for an unrelated reason. Assisted-by: Claude Fable 5 * sycl: clarify GGML_SYCL_FA_ONEDNN_MAX_KV default is disabled Assisted-by: Claude Fable 5 * sycl: state default behavior of GGML_SYCL_FA_ONEDNN_MAX_KV explicitly Assisted-by: Claude Fable 5 * Update ggml/src/ggml-sycl/fattn-onednn.cpp Co-authored-by: Neo Zhang * sycl: write the SDPA scale from a kernel instead of caching it The per-(device, value) scale cache was a function-local static unordered_map with no synchronization, so concurrent backend instances could access and rehash it at the same time. Write the scalar with a single_task instead. The value is captured into the command, so no host memory has to outlive the call -- which is what the use-after-return fix needed in the first place. That removes the shared container, the leaked device allocation and the string key, and it also closes the remaining async-memcpy-from-a-stack-local on the first flash-attention call. Ordering does not rely on timing: the queue is created with sycl::property::queue::in_order and the dnnl stream wraps that same queue, so the write completes before the SDPA reads the scalar. The multi-GPU wait_and_throw() branch is unchanged. Also drop the include, which is unused. Assisted-by: Claude Opus 5 --------- Co-authored-by: Neo Zhang --- docs/backend/SYCL.md | 2 ++ ggml/src/ggml-sycl/common.hpp | 1 + ggml/src/ggml-sycl/fattn-onednn.cpp | 30 ++++++++++++++++++++--------- ggml/src/ggml-sycl/ggml-sycl.cpp | 3 +++ tests/test-backend-ops.cpp | 9 +++++++++ 5 files changed, 36 insertions(+), 9 deletions(-) diff --git a/docs/backend/SYCL.md b/docs/backend/SYCL.md index 0814ceb60f92..c72dc3b3ee02 100644 --- a/docs/backend/SYCL.md +++ b/docs/backend/SYCL.md @@ -794,6 +794,8 @@ use 1 SYCL GPUs: [0] with Max compute units:512 | GGML_SYCL_ENABLE_GRAPH | 0 (default) or 1 | Enable running computations through SYCL Graphs feature. Disabled by default because SYCL Graph is still on development, no better performance. | | GGML_SYCL_USE_LEVEL_ZERO_API | 1 (default) or 0 | Use Level Zero API for device memory allocation instead of SYCL. Reduces system RAM usage on Intel dGPUs by avoiding DMA-buf/TTM host memory staging. Requires GGML_SYCL_SUPPORT_LEVEL_ZERO_API=ON at build time. SYCL backend always runs on Level Zero running time even if it's set as OFF (The SYCL api will be usage for memory allocation).| | GGML_SYCL_ENABLE_DNN | 0 or 1 (default)| Enable running computations through oneDNN and always use oneMKL. | +| GGML_SYCL_FA_ONEDNN | 1 (default) or 0 | Enable the oneDNN fused SDPA (flash-attention) path on supported GPUs. Set to 0 to always use the native SYCL flash-attention kernel. | +| GGML_SYCL_FA_ONEDNN_MAX_KV | 0 (default, disabled) or positive integer | By default (0), all sequences are handled by the oneDNN fused SDPA path, regardless of KV length; a positive value caps that length, past which sequences fall back to the native kernel. If GPU driver watchdog resets (DEVICE_LOST) occur during long-context inference, set this near the context depth where they start, e.g. 24576. | | GGML_SYCL_ENABLE_VMM | 0 or 1 (default) | Enable the virtual-memory device pool. | | GGML_SYCL_ENABLE_FUSION | 0 or 1 (default) | Enable fused-kernel dispatch in graph compute (currently top-k MoE gating). | | ZES_ENABLE_SYSMAN | 0 (default) or 1 | Support to get free memory of GPU by sycl::aspect::ext_intel_free_memory.
Recommended to use when --split-mode = layer | diff --git a/ggml/src/ggml-sycl/common.hpp b/ggml/src/ggml-sycl/common.hpp index e5d9ee89dd86..f27ec5dd6260 100644 --- a/ggml/src/ggml-sycl/common.hpp +++ b/ggml/src/ggml-sycl/common.hpp @@ -65,6 +65,7 @@ extern int g_ggml_sycl_prioritize_dmmv; extern int g_ggml_sycl_enable_flash_attention; extern int g_ggml_sycl_dev2dev_memcpy; extern int g_ggml_sycl_fa_onednn; +extern int g_ggml_sycl_fa_onednn_max_kv; #if defined(__clang__) && __has_builtin(__builtin_expect) diff --git a/ggml/src/ggml-sycl/fattn-onednn.cpp b/ggml/src/ggml-sycl/fattn-onednn.cpp index f2e12ef1aeff..8465e12248f4 100644 --- a/ggml/src/ggml-sycl/fattn-onednn.cpp +++ b/ggml/src/ggml-sycl/fattn-onednn.cpp @@ -38,6 +38,12 @@ bool ggml_sycl_flash_attn_ext_onednn_supported(const ggml_tensor * dst) { if (K->type != GGML_TYPE_F16 || V->type != GGML_TYPE_F16) { return false; } + // Optional KV-length ceiling (GGML_SYCL_FA_ONEDNN_MAX_KV, 0 = unlimited). Escape hatch: + // very long sequences make the fused SDPA slow enough to risk the xe driver watchdog on + // some stacks; past the cap we fall back to the native FA kernel instead. + if (g_ggml_sycl_fa_onednn_max_kv > 0 && K->ne[1] > g_ggml_sycl_fa_onednn_max_kv) { + return false; + } // gate for the following cases // 1. if the oneDNN graph Add node has no input --> skip // 2. types other than f16 need different logical_tensor declaration @@ -208,9 +214,17 @@ void ggml_sycl_flash_attn_ext_onednn(ggml_backend_sycl_context & ctx, ggml_tenso cont_to_f16_sycl((const char *) V->data, Vf.get(), d, seq, Hkv, mb, V->nb[1], V->nb[2], V->nb[3], stream); // divide-by-(1/scale) reproduces ggml's score *= kq_scale on the proven probe graph. + // + // The scale must not be uploaded with an async memcpy from a stack local: on the in-order + // queue that copy waits behind the K/V staging kernels, and once those take long enough + // (n_kv >= ~26k on B70) the host frame is recycled before the copy runs, feeding the SDPA a + // garbage scale (output collapses to a repeated token). Write the scalar from a kernel + // instead -- the value is captured into the command, so no host memory has to outlive the + // call, and the enqueue stays async. const sycl::half scale_h = (sycl::half) (1.0f / kq_scale); ggml_sycl_pool_alloc scbuf(ctx.pool(), 1); - stream->memcpy(scbuf.get(), &scale_h, sizeof(sycl::half)); + sycl::half * const scale_dev = scbuf.get(); + stream->single_task([=]() { *scale_dev = scale_h; }); ggml_sycl_pool_alloc outf(ctx.pool(), (size_t) H * q * d); // f16 contiguous SDPA out [mb,H,q,d] @@ -232,7 +246,7 @@ void ggml_sycl_flash_attn_ext_onednn(ggml_backend_sycl_context & ctx, ggml_tenso if (r == E.id_q) return Qf.get(); if (r == E.id_k) return Kf.get(); if (r == E.id_v) return Vf.get(); - if (r == E.id_scale) return scbuf.get(); + if (r == E.id_scale) return scale_dev; if (r == E.id_mask) return (void *) mask->data; return nullptr; }; @@ -245,14 +259,12 @@ void ggml_sycl_flash_attn_ext_onednn(ggml_backend_sycl_context & ctx, ggml_tenso E.cp.execute(strm, ti, {to}); permute_sdpa_out_sycl(outf.get(), (float *) dst->data, mb, H, q, d, stream); - // Single device: no sync is required, and actually PP perf is ~6% > wait_and_throw() (tested on llama-3.1-8b & qwen3.6-27b, both Q8_0, with Arc B70). - // Any future multi-GPU refactor MUST re-measure this single-device path and keep the best - // single-device PP speed. Otherwise (multiple devices/streams can race the reuse): + // Single device needs no sync: the dnnl stream wraps this same in-order queue, so the SDPA + // serializes with the staging kernels before it and the permute/pool reuse after it. The + // garbage output formerly blamed on the missing sync here was the scale use-after-return + // fixed above. Keep the conservative wait for multi-GPU, where other devices' streams can + // race the pool: if (ggml_sycl_info().device_count > 1) { - // cont_to_f16 -> oneDNN execute -> permute is async on this stream, but the - // pool_alloc*s above free their device buffers at host return. Without this wait the next - // scheduler op re-acquires those bytes while the GPU is still computing the SDPA, turning - // it into garbage and collapsing multi-turn output to a single repeated token ("GGGGG..."). stream->wait_and_throw(); } } diff --git a/ggml/src/ggml-sycl/ggml-sycl.cpp b/ggml/src/ggml-sycl/ggml-sycl.cpp index cb8974eedb75..3b807c7cbd54 100644 --- a/ggml/src/ggml-sycl/ggml-sycl.cpp +++ b/ggml/src/ggml-sycl/ggml-sycl.cpp @@ -85,6 +85,7 @@ int g_ggml_sycl_enable_optimize = 1; int g_ggml_sycl_enable_graph = 0; int g_ggml_sycl_enable_dnn = 1; int g_ggml_sycl_fa_onednn = 1; +int g_ggml_sycl_fa_onednn_max_kv = 0; int g_ggml_sycl_enable_vmm = 1; int g_ggml_sycl_enable_fusion = 1; int g_ggml_sycl_prioritize_dmmv = 0; @@ -287,6 +288,7 @@ static void ggml_check_sycl() try { g_ggml_sycl_enable_graph = ggml_sycl_get_env("GGML_SYCL_ENABLE_GRAPH", 0); g_ggml_sycl_enable_dnn = ggml_sycl_get_env("GGML_SYCL_ENABLE_DNN", 1); g_ggml_sycl_fa_onednn = ggml_sycl_get_env("GGML_SYCL_FA_ONEDNN", 1); + g_ggml_sycl_fa_onednn_max_kv = ggml_sycl_get_env("GGML_SYCL_FA_ONEDNN_MAX_KV", 0); g_ggml_sycl_enable_vmm = ggml_sycl_get_env("GGML_SYCL_ENABLE_VMM", 1); g_ggml_sycl_enable_fusion = ggml_sycl_get_env("GGML_SYCL_ENABLE_FUSION", 1); g_ggml_sycl_prioritize_dmmv = ggml_sycl_get_env("GGML_SYCL_PRIORITIZE_DMMV", 0); @@ -359,6 +361,7 @@ static void ggml_check_sycl() try { GGML_LOG_INFO(" GGML_SYCL_ENABLE_DNN: DNN disabled by compile flag\n"); GGML_LOG_INFO(" GGML_SYCL_FA_ONEDNN: %d\n", g_ggml_sycl_fa_onednn); #endif + GGML_LOG_INFO(" GGML_SYCL_FA_ONEDNN_MAX_KV: %d\n", g_ggml_sycl_fa_onednn_max_kv); #ifdef SYCL_FLASH_ATTN GGML_LOG_INFO(" GGML_SYCL_ENABLE_FLASH_ATTN: %d\n", g_ggml_sycl_enable_flash_attention); #else diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index b4061e35a050..f6b60e52a400 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -9547,6 +9547,15 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_flash_attn_ext(64, 128, 4, {1, 1}, 128, 2, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q4_0, GGML_TYPE_Q1_0)); test_cases.emplace_back(new test_flash_attn_ext(128, 64, 4, {1, 1}, 64, 2, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q1_0, GGML_TYPE_F16)); + // large-KV F16 cases (Qwen3.6-27B geometry and a llama-class control): the upstream matrix + // stops at kv=1024, blind to long-context FA bugs (e.g. the oneDNN SDPA ordering race on BMG). + for (int64_t kv : { 4096, 16384 }) { + test_cases.emplace_back(new test_flash_attn_ext(256, 256, 4, {6, 1}, kv, 512, true, false, 0, 0, + GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); + test_cases.emplace_back(new test_flash_attn_ext(128, 128, 8, {4, 1}, kv, 512, true, false, 0, 0, + GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); + } + test_cases.emplace_back(new test_cross_entropy_loss (GGML_TYPE_F32, { 10, 5, 4, 3})); test_cases.emplace_back(new test_cross_entropy_loss (GGML_TYPE_F32, {30000, 1, 1, 1})); test_cases.emplace_back(new test_cross_entropy_loss_back(GGML_TYPE_F32, { 10, 5, 4, 3})); From 6ba5ef247034cd57201360aed246d98f5a404d92 Mon Sep 17 00:00:00 2001 From: Aldehir Rojas Date: Tue, 28 Jul 2026 04:27:20 -0500 Subject: [PATCH 034/190] common/chat: add specialized minimax m3 parser (#26210) --- common/chat-peg-parser.cpp | 138 ++++++++++ common/chat-peg-parser.h | 12 + common/chat.cpp | 273 +++++++++++++++++++ common/chat.h | 1 + models/templates/MiniMax-M3.jinja | 247 +++++++++++++++++ tests/test-chat.cpp | 429 ++++++++++++++++++++++++++++++ 6 files changed, 1100 insertions(+) create mode 100644 models/templates/MiniMax-M3.jinja diff --git a/common/chat-peg-parser.cpp b/common/chat-peg-parser.cpp index a309f02765b7..f786f5ff2314 100644 --- a/common/chat-peg-parser.cpp +++ b/common/chat-peg-parser.cpp @@ -1056,3 +1056,141 @@ void common_chat_peg_gemma4_mapper::visit(const common_peg_ast_arena & arena, co visit(arena, child_id); } } + +static void minimax_m3_collect(const common_peg_ast_arena & arena, + const common_peg_ast_node & node, + const std::string & tag, + std::vector & out) { + for (auto child_id : node.children) { + const auto & child = arena.get(child_id); + if (child.tag == tag) { + out.push_back(child_id); + } else { + minimax_m3_collect(arena, child, tag, out); + } + } +} + +static common_peg_ast_id minimax_m3_value_of(const common_peg_ast_arena & arena, const common_peg_ast_node & node) { + for (auto child_id : node.children) { + const auto & tag = arena.get(child_id).tag; + if (tag == common_chat_peg_builder::TOOL_ARG_VALUE || + tag == common_chat_peg_builder::TOOL_ARG_STRING_VALUE || + tag == common_chat_peg_minimax_m3_mapper::TOOL_ARG_OBJECT || + tag == common_chat_peg_minimax_m3_mapper::TOOL_ARG_ARRAY) { + return child_id; + } + } + return COMMON_PEG_INVALID_AST_ID; +} + +static std::string minimax_m3_value_to_json(const common_peg_ast_arena & arena, common_peg_ast_id id, bool closed); + +static std::string minimax_m3_member_to_json(const common_peg_ast_arena & arena, const common_peg_ast_node & node) { + auto name_id = arena.find_by_tag(node, common_chat_peg_builder::TOOL_ARG_NAME); + if (name_id == COMMON_PEG_INVALID_AST_ID) { + return ""; + } + + return ordered_json(arena.get(name_id).text).dump() + ":" + + minimax_m3_value_to_json(arena, minimax_m3_value_of(arena, node), !node.is_partial); +} + +static std::string minimax_m3_container_to_json(const common_peg_ast_arena & arena, + const common_peg_ast_node & node, + bool is_object, + bool closed) { + const std::string tag = is_object ? common_chat_peg_builder::TOOL_ARG + : common_chat_peg_minimax_m3_mapper::TOOL_ARG_ITEM; + + std::vector entries; + minimax_m3_collect(arena, node, tag, entries); + + std::string result = is_object ? "{" : "["; + + bool add_comma = false; + for (auto entry_id : entries) { + const auto & entry = arena.get(entry_id); + + std::string text; + if (is_object) { + text = minimax_m3_member_to_json(arena, entry); + } else { + text = minimax_m3_value_to_json(arena, minimax_m3_value_of(arena, entry), !entry.is_partial); + } + + if (text.empty()) { + continue; + } + + if (add_comma) { + result += ","; + } + add_comma = true; + result += text; + } + + if (closed) { + result += is_object ? "}" : "]"; + } + return result; +} + +static std::string minimax_m3_value_to_json(const common_peg_ast_arena & arena, common_peg_ast_id id, bool closed) { + if (id == COMMON_PEG_INVALID_AST_ID) { + return ""; + } + + const auto & node = arena.get(id); + + if (node.tag == common_chat_peg_minimax_m3_mapper::TOOL_ARG_OBJECT) { + return minimax_m3_container_to_json(arena, node, /* is_object = */ true, closed); + } + + if (node.tag == common_chat_peg_minimax_m3_mapper::TOOL_ARG_ARRAY) { + return minimax_m3_container_to_json(arena, node, /* is_object = */ false, closed); + } + + if (node.tag == common_chat_peg_builder::TOOL_ARG_STRING_VALUE) { + return "\"" + escape_json_string_inner(std::string(node.text)) + (closed ? "\"" : ""); + } + + // Numbers and booleans are written verbatim by the template + return std::string(node.text); +} + +void common_chat_peg_minimax_m3_mapper::from_ast(const common_peg_ast_arena & arena, + const common_peg_parse_result & result) { + for (const auto & node : result.nodes) { + visit(arena, node); + } +} + +void common_chat_peg_minimax_m3_mapper::visit(const common_peg_ast_arena & arena, common_peg_ast_id id) { + const auto & node = arena.get(id); + + if (node.tag == common_chat_peg_builder::REASONING) { + result.reasoning_content += std::string(node.text); + return; + } + + if (node.tag == common_chat_peg_builder::CONTENT) { + result.content += std::string(node.text); + return; + } + + if (node.tag == common_chat_peg_builder::TOOL) { + auto name_id = arena.find_by_tag(node, common_chat_peg_builder::TOOL_NAME); + if (name_id != COMMON_PEG_INVALID_AST_ID) { + common_chat_tool_call call; + call.name = std::string(arena.get(name_id).text); + call.arguments = minimax_m3_container_to_json(arena, node, /* is_object = */ true, !node.is_partial); + result.tool_calls.push_back(call); + } + return; + } + + for (auto child_id : node.children) { + visit(arena, child_id); + } +} diff --git a/common/chat-peg-parser.h b/common/chat-peg-parser.h index b3ffd7de2dd8..cd14f2c11750 100644 --- a/common/chat-peg-parser.h +++ b/common/chat-peg-parser.h @@ -40,6 +40,18 @@ class common_chat_peg_gemma4_mapper : public common_chat_peg_mapper { void visit(const common_peg_ast_arena & arena, common_peg_ast_id id); }; +class common_chat_peg_minimax_m3_mapper : public common_chat_peg_mapper { + public: + static constexpr const char * TOOL_ARG_OBJECT = "tool-arg-object"; + static constexpr const char * TOOL_ARG_ARRAY = "tool-arg-array"; + static constexpr const char * TOOL_ARG_ITEM = "tool-arg-item"; + + common_chat_peg_minimax_m3_mapper(common_chat_msg & msg) : common_chat_peg_mapper(msg) {} + virtual void from_ast(const common_peg_ast_arena & arena, const common_peg_parse_result & result); + private: + void visit(const common_peg_ast_arena & arena, common_peg_ast_id id); +}; + struct content_structure; struct tool_call_structure; diff --git a/common/chat.cpp b/common/chat.cpp index 7a6e7238cf33..7740f35c0edc 100644 --- a/common/chat.cpp +++ b/common/chat.cpp @@ -816,6 +816,8 @@ const char * common_chat_format_name(common_chat_format format) { return "peg-native"; case COMMON_CHAT_FORMAT_PEG_GEMMA4: return "peg-gemma4"; + case COMMON_CHAT_FORMAT_PEG_MINIMAX_M3: + return "peg-minimax-m3"; default: throw std::runtime_error("Unknown chat format"); } @@ -2270,6 +2272,264 @@ static common_chat_params common_chat_params_init_cohere2moe(const common_chat_t return data; } +static common_chat_params common_chat_params_init_minimax_m3(const common_chat_template & tmpl, + const autoparser::generation_params & inputs) { + common_chat_params data; + + data.prompt = common_chat_template_direct_apply_impl(tmpl, inputs); + data.generation_prompt = common_chat_template_generation_prompt_impl(tmpl, inputs); + data.format = COMMON_CHAT_FORMAT_PEG_MINIMAX_M3; + data.supports_thinking = true; + data.thinking_start_tag = ""; + data.thinking_end_tags = {""}; + + // M3 prefixes every tool tag with the namespace token "]<]minimax[>["; + // params use the parameter name as the tag (...). + const std::string NS = "]<]minimax[>["; + const std::string THINK_START = ""; + const std::string THINK_END = ""; + const std::string FC_START = NS + ""; + const std::string FC_END = NS + ""; + const std::string INVOKE_END = NS + ""; + + data.preserved_tokens = { + NS, + "", + "", + THINK_START, + THINK_END, + }; + + data.message_delimiters = { + { COMMON_CHAT_ROLE_ASSISTANT, "]~b]ai" }, + { COMMON_CHAT_ROLE_USER, "]~b]user" }, + { COMMON_CHAT_ROLE_TOOL, "]~b]tool" }, + { COMMON_CHAT_ROLE_SYSTEM, "]~b]developer" }, + { COMMON_CHAT_ROLE_SYSTEM, "]~b]system" }, + }; + + auto has_tools = inputs.tools.is_array() && !inputs.tools.empty(); + auto has_response_format = !inputs.json_schema.is_null() && inputs.json_schema.is_object(); + auto extract_reasoning = inputs.reasoning_format != COMMON_REASONING_FORMAT_NONE; + auto include_grammar = has_response_format || (has_tools && inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_NONE); + + const std::string GEN_PROMPT = data.generation_prompt; + + using mm3 = common_chat_peg_minimax_m3_mapper; + + if (inputs.has_continuation()) { + const auto & msg = inputs.continue_msg; + + data.generation_prompt = GEN_PROMPT + THINK_START + msg.reasoning_content; + if (inputs.continue_final_message == COMMON_CHAT_CONTINUATION_CONTENT) { + data.generation_prompt += THINK_END + msg.render_content(); + } + + data.prompt += data.generation_prompt; + } + + auto parser = build_chat_peg_parser([&](common_chat_peg_builder & p) { + auto generation_prompt = p.prefix(GEN_PROMPT, THINK_START); + auto end = p.end(); + + auto reasoning = p.eps(); + if (extract_reasoning) { + auto block = inputs.enable_thinking + ? p.literal(THINK_START) + p.space() + + p.ac(p.reasoning(p.until(THINK_END)) + p.literal(THINK_END), THINK_END) + : p.literal(THINK_START) + p.ac(p.until(THINK_END) + p.literal(THINK_END), THINK_END); + + // A turn without reasoning is prefixed with a bare , written either by the + // generation prompt (thinking_mode = "disabled") or by the model itself. + reasoning = p.optional(p.choice({ block, p.literal(THINK_END) })); + } + + if (has_response_format) { + auto response_format = p.rule("response-format", + p.literal("```json") + p.space() + + p.content(p.schema(p.json(), "response-format-schema", inputs.json_schema)) + + p.space() + p.literal("```")); + return generation_prompt + reasoning + response_format + end; + } + + if (!has_tools || inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_NONE) { + return generation_prompt + reasoning + p.content(p.rest()) + end; + } + + auto alternatives_of = [](const json & schema) -> std::optional { + for (const auto * keyword : { "oneOf", "anyOf" }) { + if (schema.contains(keyword) && schema.at(keyword).is_array() && !schema.at(keyword).empty()) { + return schema.at(keyword); + } + } + return std::nullopt; + }; + + auto tool_choice = p.choice(); + foreach_function(inputs.tools, [&](const json & tool) { + const auto & function = tool.at("function"); + std::string name = function.at("name"); + auto params = function.contains("parameters") ? function.at("parameters") : json::object(); + + auto schema_info = common_schema_info(); + schema_info.resolve_refs(params); + + // The template expands argument values recursively in XML (see the to_xml() macro) + std::function value_of; + std::function members_of; + + auto element_of = [&](const std::string & tag, const json & schema, const std::string & rule_name) { + const std::string close = NS + ""; + return p.rule(rule_name, + p.tool_arg( + p.tool_arg_open( + p.literal(NS + "<") + + p.tool_arg_name(p.literal(tag)) + + p.literal(">")) + + value_of(schema, rule_name, close))); + }; + + value_of = [&](const json & schema, + const std::string & rule_name, + const std::string & close) -> common_peg_parser { + auto close_tag = p.tool_arg_close(p.literal(close)); + + // A string accepts anything, so a union with a string alternative is a string + if (schema_info.resolves_to_string(schema)) { + return p.ac(p.tool_arg_string_value(p.until(close)) + close_tag, close); + } + + if (auto alternatives = alternatives_of(schema)) { + std::vector choices; + + size_t index = 0; + for (const auto & alternative : *alternatives) { + const std::string alt_name = rule_name + "-" + std::to_string(index++); + + // There is a risk that this breaks streaming deltas, but that's a risk we + // assume to provide tool arg streaming. + choices.push_back(value_of(alternative, alt_name, close)); + } + + return p.choice(choices); + } + + const std::string type = schema.contains("type") && schema.at("type").is_string() + ? schema.at("type").get() + : ""; + + if (type == "object" && schema.contains("properties")) { + return p.tag(mm3::TOOL_ARG_OBJECT, members_of(schema, rule_name)) + p.space() + close_tag; + } + + if (type == "array" && schema.contains("items")) { + const std::string item_close = NS + ""; + auto item = p.rule(rule_name + "-item", + p.tag(mm3::TOOL_ARG_ITEM, + p.literal(NS + "") + + value_of(schema.at("items"), rule_name + "-item", item_close))); + return p.tag(mm3::TOOL_ARG_ARRAY, p.repeat(p.space() + item, 0, -1)) + p.space() + close_tag; + } + + return p.tool_arg_json_value(p.schema(p.json(), rule_name + "-schema", schema, false)) + close_tag; + }; + + // Required properties in schema order, then any number of optional ones in any order. + members_of = [&](const json & schema, const std::string & rule_prefix) -> common_peg_parser { + const auto & props = schema.at("properties"); + + std::set required; + if (schema.contains("required")) { + schema.at("required").get_to(required); + } + + std::vector required_elements; + std::vector optional_elements; + for (const auto & [key, key_schema] : props.items()) { + auto element = element_of(key, key_schema, rule_prefix + "-" + key); + if (required.find(key) != required.end()) { + required_elements.push_back(element); + } else { + optional_elements.push_back(element); + } + } + + common_peg_parser members = p.eps(); + for (size_t i = 0; i < required_elements.size(); i++) { + if (i > 0) { + members = members + p.space(); + } + members = members + required_elements[i]; + } + + if (!optional_elements.empty()) { + common_peg_parser any_optional = p.choice(); + for (const auto & element : optional_elements) { + any_optional |= element; + } + members = members + p.repeat(p.space() + any_optional, 0, -1); + } + + return members; + }; + + common_peg_parser invoke_body = + params.contains("properties") ? members_of(params, "tool-" + name + "-arg") : p.eps(); + + auto func_parser = p.tool( + p.tool_open(p.literal(NS + "")) + + p.space() + invoke_body + p.space() + + p.tool_close(p.literal(INVOKE_END))); + + tool_choice |= p.rule("tool-" + name, func_parser); + }); + + auto require_tools = inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED; + + common_peg_parser tool_calls = p.eps(); + if (inputs.parallel_tool_calls) { + tool_calls = p.trigger_rule("tool-call", + p.literal(FC_START) + p.space() + tool_choice + + p.zero_or_more(p.space() + tool_choice) + p.space() + p.literal(FC_END)); + } else { + tool_calls = p.trigger_rule("tool-call", + p.literal(FC_START) + p.space() + tool_choice + p.space() + p.literal(FC_END)); + } + + if (!require_tools) { + tool_calls = p.optional(tool_calls); + } + + auto content_before_tools = p.content(p.until(FC_START)); + return generation_prompt + reasoning + content_before_tools + tool_calls + end; + }); + + data.parser = parser.save(); + + if (include_grammar) { + data.grammar_lazy = !(has_response_format || (has_tools && inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED)); + data.grammar = build_grammar([&](const common_grammar_builder & builder) { + foreach_function(inputs.tools, [&](const json & tool) { + const auto & function = tool.at("function"); + auto schema = function.contains("parameters") ? function.at("parameters") : json::object(); + builder.resolve_refs(schema); + }); + if (has_response_format) { + auto schema = inputs.json_schema; + builder.resolve_refs(schema); + } + parser.build_grammar(builder, data.grammar_lazy); + }); + + data.grammar_triggers = { + { COMMON_GRAMMAR_TRIGGER_TYPE_WORD, FC_START }, + }; + } + + return data; +} + namespace workaround { static void map_developer_role_to_system(json & messages) { @@ -2707,6 +2967,15 @@ std::optional common_chat_try_specialized_template( return common_chat_params_init_gigachat_v3(tmpl, params); } + // MiniMax-M3: the namespace token "]<]minimax[>[" collides with the autoparser's + // markup delimiters, so detect the template and use a dedicated parser. + if (src.find("]<]minimax[>[") != std::string::npos && + src.find("") != std::string::npos && + src.find(" mapper; if (params.format == COMMON_CHAT_FORMAT_PEG_GEMMA4) { mapper = std::make_unique(msg); + } else if (params.format == COMMON_CHAT_FORMAT_PEG_MINIMAX_M3) { + mapper = std::make_unique(msg); } else { mapper = std::make_unique(msg); } @@ -3020,6 +3291,8 @@ common_chat_msg common_chat_peg_parse(const common_peg_arena & src_pars std::unique_ptr mapper; if (params.format == COMMON_CHAT_FORMAT_PEG_GEMMA4) { mapper = std::make_unique(msg); + } else if (params.format == COMMON_CHAT_FORMAT_PEG_MINIMAX_M3) { + mapper = std::make_unique(msg); } else { mapper = std::make_unique(msg); } diff --git a/common/chat.h b/common/chat.h index d79f4ecd773c..6d5b220aebb5 100644 --- a/common/chat.h +++ b/common/chat.h @@ -233,6 +233,7 @@ enum common_chat_format { COMMON_CHAT_FORMAT_PEG_SIMPLE, COMMON_CHAT_FORMAT_PEG_NATIVE, COMMON_CHAT_FORMAT_PEG_GEMMA4, + COMMON_CHAT_FORMAT_PEG_MINIMAX_M3, COMMON_CHAT_FORMAT_COUNT, // Not a format, just the # formats }; diff --git a/models/templates/MiniMax-M3.jinja b/models/templates/MiniMax-M3.jinja new file mode 100644 index 000000000000..93022eb9ceec --- /dev/null +++ b/models/templates/MiniMax-M3.jinja @@ -0,0 +1,247 @@ +{# ---------- special token variables ---------- #} +{%- set ns_token = ']<]minimax[>[' -%} +{%- set bod_token = ']~!b[' -%} +{%- set bos_token = ']~b]' -%} +{%- set eos_token = '[e~[' -%} +{%- set toolcall_begin_token = ns_token ~ '' -%} +{%- set toolcall_end_token = ns_token ~ '' -%} +{%- set think_begin_token = '' -%} +{%- set think_end_token = '' -%} +{%- set image_token = ']<]image[>[' -%} +{%- set video_token = ']<]video[>[' -%} +{#- Thinking mode: "enabled" / "disabled" / "adaptive" / not defined -#} +{#- Recursive XML renderer for tool_call arguments ======================== -#} +{#- None values are intentionally skipped in mapping iteration so that + `null` (which would round-trip to the literal string "null") + never appears in the rendered tool_call. The convention is: omit the + field entirely. The top-level `_args` loop applies the same rule. + The `val is none` branch below is a safety net only — upstream cleaning + (drop_none_in_tool_arguments) should ensure no None ever reaches here. -#} +{%- macro to_xml(val, ns) -%} +{%- if val is mapping -%} +{%- for k, v in val.items() if v is not none -%} +{{ ns }}<{{ k }}>{{ to_xml(v, ns) }}{{ ns }} +{%- endfor -%} +{%- elif val is iterable and val is not string -%} +{%- for item in val -%} +{{ ns }}{{ to_xml(item, ns) }}{{ ns }} +{%- endfor -%} +{%- elif val is none -%} +{#- Should be unreachable when upstream cleaning is applied. -#} +{%- elif val is boolean -%} +{{ val | tojson }} +{%- else -%} +{{ val }} +{%- endif -%} +{%- endmacro -%} +{#- Tool Rendering Functions ============================================== -#} +{%- macro render_tool_namespace(namespace_name, tool_list) -%} +{%- for tool in tool_list -%} +{{ tool.function | tojson(ensure_ascii=False) }} +{% endfor -%} +{%- endmacro -%} +{%- macro visible_text(content) -%} + {%- if content is string -%} + {{ content }} + {%- elif content is iterable and content is not mapping -%} + {%- for item in content -%} + {%- if item is mapping and item.type == 'text' -%} + {{- item.text }} + {%- elif item is mapping and item.type == 'image' -%} + {{- image_token }} + {%- elif item is mapping and item.type == 'video' -%} + {{- video_token}} + {%- elif item is string -%} + {{- item }} + {%- endif -%} + {%- endfor -%} + {%- elif content is none -%} + {{- '' }} + {%- else -%} + {{- content }} + {%- endif -%} +{%- endmacro -%} +{#- System Message Construction ============================================ -#} +{%- macro build_system_message(system_message) -%} + {%- if system_message and system_message.content -%} + {{- visible_text(system_message.content) }} + {%- else -%} + {{- 'Your model version is MiniMax-M3, developed by MiniMax. Knowledge cutoff: January 2026. Founded in early 2022, MiniMax is a global AI foundation model company committed to advancing the frontiers of AI towards AGI.' }} + {%- endif -%} + + {#- Thinking mode instructions -#} + {{- '\n\n\n' }} + {{- 'You have a thinking capability that allows you to reason step by step before responding. When thinking is enabled, wrap your reasoning in ' ~ think_begin_token ~ think_end_token ~ ' tags before your response. When thinking is disabled, begin your response directly after the ' ~ think_end_token ~ ' prefix. When thinking is adaptive, decide on your own whether to think for the current turn.\n' }} + {%- if thinking_mode is defined -%} + {%- if thinking_mode == "enabled" -%} + {{- 'Current thinking mode: enabled. You MUST think step by step before every response, including after receiving function/tool results.\n' }} + {%- elif thinking_mode == "disabled" -%} + {{- 'Current thinking mode: disabled. Do not output any thinking process.\n' }} + {%- elif thinking_mode == "adaptive" -%} + {{- 'Current thinking mode: adaptive. You are encouraged to think for complex decision-making, multi-step reasoning, or when analyzing function/tool results.\n' }} + {%- endif -%} + {%- else -%} + {{- 'Current thinking mode: adaptive. You are encouraged to think for complex decision-making, multi-step reasoning, or when analyzing function/tool results.\n' }} + {%- endif -%} + {{- '' }} +{%- endmacro -%} +{%- macro build_developer_message(developer_message) -%} + {%- if developer_message and developer_message.content -%} + {{- visible_text(developer_message.content) }} + {%- else -%} + {%- if model_identity is not defined -%} + {%- set model_identity = "You are a helpful assistant." -%} + {%- endif -%} + {{- model_identity }} + {%- endif -%} +{%- endmacro -%} +{#- Main Template Logic ================================================= -#} +{#- Role mapping: root -> system sp (high priority), system/developer -> developer sp (low priority) -#} +{%- set system_message = none -%} +{%- set developer_message = none -%} +{%- set conversation_messages = messages -%} +{%- if messages and messages[0].role == "root" -%} + {%- set system_message = messages[0] -%} + {%- set conversation_messages = messages[1:] -%} + {%- if conversation_messages and conversation_messages[0].role in ["system", "developer"] -%} + {%- set developer_message = conversation_messages[0] -%} + {%- set conversation_messages = conversation_messages[1:] -%} + {%- endif -%} +{%- elif messages and messages[0].role in ["system", "developer"] -%} + {%- set developer_message = messages[0] -%} + {%- set conversation_messages = messages[1:] -%} +{%- endif -%} +{#- Render system sp (higher priority, root role only) -#} +{{- bod_token ~ bos_token ~ 'system' ~ '\n' }} +{{- build_system_message(system_message) }} +{{- eos_token ~ '\n' }} + +{#- Render developer sp (lower priority: system/developer role + tools) -#} +{{- bos_token ~ 'developer' ~ '\n' }} +{{- build_developer_message(developer_message) }} +{%- if tools -%} + {{- '\n\n' ~ '# Tools' ~ '\n' ~ 'You may call one or more tools to assist with the user query.\nHere are the tools available in JSONSchema format:' ~ '\n' }} + {{- '\n' ~ '' ~ '\n' }} + {{- render_tool_namespace("functions", tools) }} + {{- '' ~ '\n\n' }} + {{- 'To call tools, wrap all invocations in a single ' ~ toolcall_begin_token ~ toolcall_end_token ~ ' block. Parameter values containing nested objects or arrays are recursively expanded into XML elements. Example:\n' }} + {{- '\n' ~ toolcall_begin_token ~ '\n' }} + {{- ns_token + '' }} + {{- ns_token + 'value-1' + ns_token + '' }} + {{- ns_token + '' }} + {{- ns_token + '' }} + {{- ns_token + 'val-a' + ns_token + '' }} + {{- ns_token + 'val-b' + ns_token + '' }} + {{- ns_token + '' }} + {{- ns_token + '' }} + {{- ns_token + '\n' }} + {{- ns_token + '' }} + {{- ns_token + 'value-1' + ns_token + '' }} + {{- ns_token + '\n' }} + {{- toolcall_end_token }} +{%- endif -%} +{{- eos_token ~ '\n' }} + +{#- Render messages -#} +{%- set last_tool_call = namespace(name=none) -%} +{%- for message in conversation_messages -%} + {%- if message.role == 'assistant' -%} + {{- bos_token ~ 'ai' ~ '\n' }} + + {%- set reasoning_content = '' %} + {%- set content = visible_text(message.content) %} + {%- if message.reasoning_content is string %} + {%- set reasoning_content = message.reasoning_content %} + {%- else %} + {%- if think_end_token in content %} + {%- set reasoning_content = content.split(think_end_token)[0].strip('\n').split(think_begin_token)[-1].strip('\n') %} + {%- set content = content.split(think_end_token)[-1].strip('\n') %} + {%- endif %} + {%- endif %} + + {%- if reasoning_content -%} + {#- Render thinking for every assistant turn (all-turn visible) -#} + {{- think_begin_token ~ reasoning_content ~ think_end_token }} + {%- else -%} + {#- No thinking rendered → prefix with think_end_token -#} + {{- think_end_token }} + {%- endif -%} + + {%- if content -%} + {{- content }} + {%- endif -%} + {%- if message.tool_calls -%} + {{- toolcall_begin_token ~ '\n' }} + + {%- for tool_call in message.tool_calls -%} + {%- if tool_call.function -%} + {%- set tool_call = tool_call.function -%} + {%- endif -%} +{{- ns_token + '' }} +{%- set _args = tool_call.arguments -%} +{%- for k, v in _args.items() if v is not none %} +{{- ns_token + '<' + k + '>' -}} +{{- to_xml(v, ns_token) -}} +{{- ns_token + '' }} +{%- endfor -%} +{{- ns_token + '' ~ '\n' }} + {%- endfor -%} + + {{- toolcall_end_token }} + {%- if message.tool_calls[-1].function -%} + {%- set last_tool_call.name = message.tool_calls[-1].function.name -%} + {%- else -%} + {%- set last_tool_call.name = message.tool_calls[-1].name -%} + {%- endif -%} + {%- else -%} + {%- set last_tool_call.name = none -%} + {%- endif -%} + {{- eos_token ~ '\n' }} + + {%- elif message.role == 'tool' -%} + {%- if last_tool_call.name is none -%} + {{- raise_exception("Message has tool role, but there was no previous assistant message with a tool call!") }} + {%- endif -%} + {%- if loop.first or (conversation_messages[loop.index0 - 1].role != 'tool') -%} + {{- bos_token ~ 'tool' }} + {%- endif -%} + {{- '\n' }} + {%- if message.content is string -%} + {{- message.content }} + {%- else -%} + {%- for tr in message.content -%} + {%- if tr is mapping and tr.type is defined and tr.type == 'image' -%} + {{- image_token }} + {%- elif tr is mapping and tr.type is defined and tr.type == 'video' -%} + {{- video_token }} + {%- else -%} + {{- tr.output if tr.output is defined else (tr.text if tr.type == 'text' and tr.text is defined else tr) }} + {%- endif -%} + {%- endfor -%} + {%- endif -%} + {{- '' }} + {%- if loop.last or (conversation_messages[loop.index0 + 1].role != 'tool') -%} + {{- eos_token ~ '\n' -}} + {%- endif -%} + + {%- elif message.role == 'user' -%} + {{- bos_token ~ 'user' ~ '\n' }} + {{- visible_text(message.content) }} + {{- eos_token ~ '\n' }} + {%- endif -%} +{%- endfor -%} + +{#- Generation prompt -#} +{%- if add_generation_prompt -%} +{{- bos_token ~ 'ai' ~ '\n' }} +{%- if thinking_mode is defined and thinking_mode == "disabled" -%} + {{- think_end_token }} +{%- elif thinking_mode is defined and thinking_mode == "adaptive" -%} + {#- adaptive: no prefix, let model decide -#} +{%- elif thinking_mode is defined and thinking_mode == "enabled" -%} + {#- enabled or not defined: default to think -#} + {{- think_begin_token }} +{%- else -%} + {#- adaptive: no prefix, let model decide -#} +{%- endif -%} +{%- endif -%} diff --git a/tests/test-chat.cpp b/tests/test-chat.cpp index 4dd00efddf73..01b07953a627 100644 --- a/tests/test-chat.cpp +++ b/tests/test-chat.cpp @@ -730,6 +730,71 @@ static common_chat_tool imaginary_number_tool{ })", }; +static common_chat_tool nested_args_tool{ + /* .name = */ "nested_args", + /* .description = */ "Tool with nested array arguments", + /* .parameters = */ R"({ + "type": "object", + "properties": { + "tags": { + "type": "array", + "items": { "type": "string" } + }, + "entries": { + "type": "array", + "items": { + "type": "object", + "properties": { + "id": { "type": "integer" }, + "label": { "type": "string" } + }, + "required": ["id", "label"] + } + } + }, + "required": ["tags", "entries"] + })", +}; + +static common_chat_tool union_args_tool{ + /* .name = */ "union_args", + /* .description = */ "Tool with union arguments", + /* .parameters = */ R"({ + "type": "object", + "properties": { + "filter": { + "anyOf": [ + { "type": "array", "items": { "type": "string" } }, + { + "type": "object", + "properties": { + "field": { "type": "string" }, + "op": { "type": "string" } + }, + "required": ["field", "op"] + } + ] + }, + "label": { + "oneOf": [ + { "type": "string" }, + { "type": "object", "properties": { "text": { "type": "string" } } } + ] + }, + "limit": { + "oneOf": [ + { "type": "integer" }, + { + "type": "object", + "properties": { "max": { "type": "integer" } }, + "required": ["max"] + } + ] + } + } + })", +}; + static common_chat_tool nullable_string_tool{ /* .name = */ "set_nullable_str", /* .description = */ "Set a nullable string value", @@ -4850,6 +4915,370 @@ static void test_template_output_peg_parsers(bool detailed_debug) { .run(); } + // MiniMax-M3 tests - namespaced XML invoke format, the parameter name is the tag + // Format: + // ]<]minimax[>[ + // ]<]minimax[>[]<]minimax[>[Tokyo]<]minimax[>[]<]minimax[>[ + // ]<]minimax[>[ + // Reasoning uses .... The generation prompt is only "]~b]ai\n", so the model + // opens the thinking block itself; a turn without reasoning is prefixed with a bare . + { + auto tst = peg_tester("models/templates/MiniMax-M3.jinja", detailed_debug); + + // Content only (bare prefix) + tst.test("Hello, world!\nWhat's up?") + .enable_thinking(true) + .reasoning_format(COMMON_REASONING_FORMAT_AUTO) + .expect(message_assist) + .expect_reconstruction() + .run(); + + // Thinking + content + tst.test("I'm\nthinkingHello, world!\nWhat's up?") + .enable_thinking(true) + .reasoning_format(COMMON_REASONING_FORMAT_AUTO) + .expect(message_assist_thoughts) + .expect_reconstruction() + .run(); + + // Thinking + tool call (single, string param) + tst.test( + "Let me check the time" + "]<]minimax[>[\n" + "]<]minimax[>[" + "]<]minimax[>[Tokyo]<]minimax[>[" + "]<]minimax[>[\n" + "]<]minimax[>[") + .enable_thinking(true) + .reasoning_format(COMMON_REASONING_FORMAT_AUTO) + .tools({ get_time_tool }) + .expect(message_with_tool_calls_and_reasoning("get_time", R"({"city": "Tokyo"})", "Let me check the time")) + .expect_reconstruction() + .run(); + + // Tool call without reasoning, integer param + tst.test( + "" + "]<]minimax[>[\n" + "]<]minimax[>[" + "]<]minimax[>[1]<]minimax[>[" + "]<]minimax[>[\n" + "]<]minimax[>[") + .enable_thinking(true) + .reasoning_format(COMMON_REASONING_FORMAT_AUTO) + .tools({ special_function_tool }) + .expect(message_assist_call) + .expect_reconstruction() + .run(); + + // Tool call with no parameters + tst.test( + "Let's call a tool:" + "]<]minimax[>[\n" + "]<]minimax[>[" + "]<]minimax[>[\n" + "]<]minimax[>[") + .enable_thinking(true) + .reasoning_format(COMMON_REASONING_FORMAT_AUTO) + .tools({ empty_args_tool }) + .expect(message_with_reasoning_and_tool_call("Let's call a tool:", "empty_args", "{}")) + .expect_reconstruction() + .run(); + + // Multiple parallel tool calls in one block + tst.test( + "Calling both" + "]<]minimax[>[\n" + "]<]minimax[>[" + "]<]minimax[>[Paris]<]minimax[>[" + "]<]minimax[>[\n" + "]<]minimax[>[" + "]<]minimax[>[Paris]<]minimax[>[" + "]<]minimax[>[\n" + "]<]minimax[>[") + .enable_thinking(true) + .reasoning_format(COMMON_REASONING_FORMAT_AUTO) + .parallel_tool_calls(true) + .tools({ get_time_tool, get_weather_tool }) + .expect(message_with_reasoning_content_and_multiple_tool_calls( + "Calling both", "", + { { "get_time", R"({"city": "Paris"})" }, { "get_weather", R"({"city": "Paris"})" } })) + .expect_reconstruction() + .run(); + + // Content before the tool call block + tst.test( + "Thinking about it" + "Let me call the function." + "]<]minimax[>[\n" + "]<]minimax[>[" + "]<]minimax[>[1]<]minimax[>[" + "]<]minimax[>[\n" + "]<]minimax[>[") + .enable_thinking(true) + .reasoning_format(COMMON_REASONING_FORMAT_AUTO) + .tools({ special_function_tool }) + .expect_reasoning("Thinking about it") + .expect_content("Let me call the function.") + .expect_tool_calls({ + { "special_function", R"({"arg1": 1})", {} }, + }) + .expect_reconstruction() + .run(); + + // Negative number + tst.test( + "Test negative" + "]<]minimax[>[\n" + "]<]minimax[>[" + "]<]minimax[>[-14]<]minimax[>[" + "]<]minimax[>[\n" + "]<]minimax[>[") + .enable_thinking(true) + .reasoning_format(COMMON_REASONING_FORMAT_AUTO) + .tools({ magic_int_tool }) + .expect_reasoning("Test negative") + .expect_tool_calls({ + { "magic_int", R"({"ref": -14})", {} }, + }) + .expect_reconstruction() + .run(); + + // Decimal number + tst.test( + "Test decimal" + "]<]minimax[>[\n" + "]<]minimax[>[" + "]<]minimax[>[3.14]<]minimax[>[" + "]<]minimax[>[\n" + "]<]minimax[>[") + .enable_thinking(true) + .reasoning_format(COMMON_REASONING_FORMAT_AUTO) + .tools({ amount_tool }) + .expect_reasoning("Test decimal") + .expect_tool_calls({ + { "amount", R"({"orig": 3.14})", {} }, + }) + .expect_reconstruction() + .run(); + + // Boolean + tst.test( + "Test boolean" + "]<]minimax[>[\n" + "]<]minimax[>[" + "]<]minimax[>[true]<]minimax[>[" + "]<]minimax[>[\n" + "]<]minimax[>[") + .enable_thinking(true) + .reasoning_format(COMMON_REASONING_FORMAT_AUTO) + .tools({ toggle_tool }) + .expect_reasoning("Test boolean") + .expect_tool_calls({ + { "toggle", R"({"enabled": true})", {} }, + }) + .expect_reconstruction() + .run(); + + // Multiple params of mixed types (required int first, then optional string) + tst.test( + "Multi-arg call" + "]<]minimax[>[\n" + "]<]minimax[>[" + "]<]minimax[>[42]<]minimax[>[" + "]<]minimax[>[foo bar]<]minimax[>[" + "]<]minimax[>[\n" + "]<]minimax[>[") + .enable_thinking(true) + .reasoning_format(COMMON_REASONING_FORMAT_AUTO) + .tools({ magic_int_tool }) + .expect_reasoning("Multi-arg call") + .expect_tool_calls({ + { "magic_int", R"({"ref": 42, "name": "foo bar"})", {} }, + }) + .expect_reconstruction() + .run(); + + // Nested object param, expanded into one element per key + tst.test( + "Nested object" + "]<]minimax[>[\n" + "]<]minimax[>[" + "]<]minimax[>[" + "]<]minimax[>[1.5]<]minimax[>[" + "]<]minimax[>[-2.5]<]minimax[>[" + "]<]minimax[>[" + "]<]minimax[>[\n" + "]<]minimax[>[") + .enable_thinking(true) + .reasoning_format(COMMON_REASONING_FORMAT_AUTO) + .tools({ imaginary_number_tool }) + .expect_reasoning("Nested object") + .expect_tool_calls({ + { "imaginary_number", R"({"number": {"real": 1.5, "imaginary": -2.5}})", {} }, + }) + .expect_reconstruction() + .run(); + + // Array params, expanded into elements (of scalars and of objects) + tst.test( + "Nested arrays" + "]<]minimax[>[\n" + "]<]minimax[>[" + "]<]minimax[>[" + "]<]minimax[>[alpha]<]minimax[>[" + "]<]minimax[>[beta]<]minimax[>[" + "]<]minimax[>[" + "]<]minimax[>[" + "]<]minimax[>[" + "]<]minimax[>[1]<]minimax[>[" + "]<]minimax[>[" + "]<]minimax[>[" + "]<]minimax[>[" + "]<]minimax[>[\n" + "]<]minimax[>[") + .enable_thinking(true) + .reasoning_format(COMMON_REASONING_FORMAT_AUTO) + .tools({ nested_args_tool }) + .expect_reasoning("Nested arrays") + .expect_tool_calls({ + { "nested_args", R"({"tags": ["alpha", "beta"], "entries": [{"id": 1, "label": "one"}]})", {} }, + }) + .expect_reconstruction() + .run(); + + // Union params (anyOf/oneOf), expanded as a choice of the alternatives + tst.test( + "Union array" + "]<]minimax[>[\n" + "]<]minimax[>[" + "]<]minimax[>[" + "]<]minimax[>[alpha]<]minimax[>[" + "]<]minimax[>[beta]<]minimax[>[" + "]<]minimax[>[" + "]<]minimax[>[\n" + "]<]minimax[>[") + .enable_thinking(true) + .reasoning_format(COMMON_REASONING_FORMAT_AUTO) + .tools({ union_args_tool }) + .expect_reasoning("Union array") + .expect_tool_calls({ + { "union_args", R"({"filter": ["alpha", "beta"]})", {} }, + }) + .expect_reconstruction() + .run(); + + // oneOf between a scalar and an object + tst.test( + "Union scalar" + "]<]minimax[>[\n" + "]<]minimax[>[" + "]<]minimax[>[5]<]minimax[>[" + "]<]minimax[>[\n" + "]<]minimax[>[") + .enable_thinking(true) + .reasoning_format(COMMON_REASONING_FORMAT_AUTO) + .tools({ union_args_tool }) + .expect_reasoning("Union scalar") + .expect_tool_calls({ + { "union_args", R"({"limit": 5})", {} }, + }) + .expect_reconstruction() + .run(); + + tst.test( + "Union nested" + "]<]minimax[>[\n" + "]<]minimax[>[" + "]<]minimax[>[" + "]<]minimax[>[10]<]minimax[>[" + "]<]minimax[>[" + "]<]minimax[>[\n" + "]<]minimax[>[") + .enable_thinking(true) + .reasoning_format(COMMON_REASONING_FORMAT_AUTO) + .tools({ union_args_tool }) + .expect_reasoning("Union nested") + .expect_tool_calls({ + { "union_args", R"({"limit": {"max": 10}})", {} }, + }) + .expect_reconstruction() + .run(); + + // A union with a string alternative is a string + tst.test( + "Union string" + "]<]minimax[>[\n" + "]<]minimax[>[" + "]<]minimax[>[" + "]<]minimax[>[\n" + "]<]minimax[>[") + .enable_thinking(true) + .reasoning_format(COMMON_REASONING_FORMAT_AUTO) + .tools({ union_args_tool }) + .expect_reasoning("Union string") + .expect_tool_calls({ + { "union_args", R"({"label": "hi"})", {} }, + }) + .expect_reconstruction() + .run(); + + // ... even when the value looks structured + tst.test( + "Union string" + "]<]minimax[>[\n" + "]<]minimax[>[" + "]<]minimax[>[" + "]<]minimax[>[\n" + "]<]minimax[>[") + .enable_thinking(true) + .reasoning_format(COMMON_REASONING_FORMAT_AUTO) + .tools({ union_args_tool }) + .expect_reasoning("Union string") + .expect_tool_calls({ + { "union_args", R"({"label": "]<]minimax[>[hi]<]minimax[>["})", {} }, + }) + .expect_reconstruction() + .run(); + + // Edge case: empty reasoning followed by a tool call + tst.test( + "" + "]<]minimax[>[\n" + "]<]minimax[>[" + "]<]minimax[>[XYZCITY]<]minimax[>[" + "]<]minimax[>[\n" + "]<]minimax[>[") + .enable_thinking(true) + .reasoning_format(COMMON_REASONING_FORMAT_AUTO) + .tools({ get_time_tool }) + .expect(message_with_tool_calls("get_time", R"({"city": "XYZCITY"})")) + .run(); + + // Continuation tests + tst.test("world!\nWhat's up?") + .reasoning_format(COMMON_REASONING_FORMAT_AUTO) + .enable_thinking(true) + .messages({ message_user, message_assist_prefill_content }) + .add_generation_prompt(false) + .continue_final_message(COMMON_CHAT_CONTINUATION_CONTENT) + .expect_reasoning("I'm thinking") + .expect_content("Hello, world!\nWhat's up?") + .run(); + + tst.test(" thinkingHello, world!\nWhat's up?") + .reasoning_format(COMMON_REASONING_FORMAT_AUTO) + .enable_thinking(true) + .messages({ message_user, message_assist_prefill_reasoning }) + .add_generation_prompt(false) + .continue_final_message(COMMON_CHAT_CONTINUATION_REASONING) + .expect_reasoning("I'm thinking") + .expect_content("Hello, world!\nWhat's up?") + .run(); + } + // NVIDIA-Nemotron-Nano-v2 tests - ... format // Format: [{"name": "func", "arguments": {...}}] { From 84075273c82f7681d43436b692073cbd4ab15fe9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E7=8E=8B=E9=87=91=E6=97=AD?= <105263726+wjinxu@users.noreply.github.com> Date: Tue, 28 Jul 2026 19:43:27 +0800 Subject: [PATCH 035/190] spec: add DSpark speculative decoding (#25173) * spec: add DSpark speculative decoding DSpark (DeepSpec, 2026) on top of the merged DFlash drafter. It reuses the DFlash encoder/decoder graph, target feature extraction and KV-cache injection, and the verify/accept path unchanged; the draft model is a new "dspark" arch adding a low-rank Markov head (markov_w1/w2) and an optional (unused here) confidence head. No new public APIs. The proposal is the only change: the block is anchor-first (position 0 already predicts the first draft) and the decoder graph applies a semi-autoregressive, previous-token conditioned logit bias in-graph, chained per block position: logits'(i) = logits(i) + markov_w2 . markov_w1[prev(i)] prev(0) = the block's anchor token, prev(i>0) = argmax(logits'(i-1)) vectorized across all blocks in the batch; the anchors are fed through a dedicated graph input (token 0 of every block). Greedy stays lossless (verify unchanged, same as DFlash). - new arch "dspark" (llama_model_dspark : llama_model_dflash, reuses the graph, loads the markov/confidence tensors; shares the target's embed/lm_head). - Qwen3DSparkModel converter. - new spec type "draft-dspark" (common_speculative_impl_draft_dspark : common_speculative_impl_draft_dflash, overrides draft() only: submits whole anchor-first blocks and greedily reads back the biased logits). * spec: read draft block size in the dflash impl * docs: add DSpark section to speculative.md * spec: keep dspark block size read in the dspark impl * dspark : add TODOs for incomplete parts - confidence head is loaded but not used yet - confidence-scheduled prefix pruning is not implemented - the in-graph Markov chain is greedy-only - only Qwen3 backbones are supported for now (also noted in docs) * spec: fold DSpark into the DFlash arch Address review: drop LLM_ARCH_DSPARK and the dspark.block_size / markov_rank GGUF keys. A DSpark draft now converts to a DFlash GGUF; the Markov head tensors are detected by presence (like eagle3 d2t), block_size is read from the existing dflash.block_size key, and the block anchors are taken as a strided view of the decoder's token input instead of a separate graph input. * spec: add confidence-based draft pruning for DSpark The DSpark confidence head predicts per-position acceptance of the drafted block. --spec-draft-conf-min truncates the block at the first position below the threshold (default 0 = disabled). * fold the dspark impl into dflash, selected by spec type * address review comments * dspark: clean up and improve naming * update readme * remove trailing whitespace * dflash: draft full n_max blocks, defer dp.n_max to the central truncation The DSpark markov head views the draft batch as a uniform [n_seqs x block] grid, but the per-seq dp.n_max clamp could produce blocks of different sizes, silently corrupting the strided views and the resulting logits. Drop the clamp and always draft the full n_max block for every sequence: dp.n_max is already enforced by the central truncation in common_speculative_draft(), the same way eagle3 handles it. Co-authored-by: Zaire404 <3147879462@qq.com> * dflash: assert the markov head block-uniformity invariant, require the conf head With the draft batch always submitting equal-size n_max blocks, a non-divisible token count can only mean the batch was split across ubatches or a caller broke the layout - fail loudly instead of silently dropping the markov bias. The block_drafts > block_size early return stays: worst-case graph reserve passes legitimately build with n_seq_tokens > block_size. Also make conf_proj required when the markov head is present: the confidence head is part of the DSpark checkpoint format, and a missing head would otherwise leave --spec-draft-conf-min silently reading stale embeddings instead of confidences. Co-authored-by: Zaire404 <3147879462@qq.com> * dspark: fold conf_min into p_min p_min and conf_min express the same thing - the minimum predicted survival probability for a drafted position - differing only in how the estimate is obtained: token probability for regular drafters, the trained confidence head for DSpark. The DSpark readback never used p_min, so reuse it for the confidence threshold and drop the separate --spec-draft-conf-min flag. Both defaulted to 0 (disabled), so behavior is unchanged. Co-authored-by: Zaire404 <3147879462@qq.com> * dflash: note the confidence broadcast workaround Requested in review: the ggml_repeat only adapts the [1, n_tok] confidences to the n_embd-wide embd_nextn transport so that llama_get_embeddings_nextn can be reused - not a placeholder. Co-authored-by: Zaire404 <3147879462@qq.com> * cont : clarify [no ci] --------- Co-authored-by: Ruixiang Wang Co-authored-by: Zaire404 <3147879462@qq.com> Co-authored-by: Georgi Gerganov --- common/common.h | 3 +- common/speculative.cpp | 104 +++++++++++++++++++++--------- conversion/__init__.py | 1 + conversion/qwen.py | 20 ++++++ docs/speculative.md | 35 ++++++++++- gguf-py/gguf/constants.py | 11 ++++ gguf-py/gguf/tensor_mapping.py | 12 ++++ src/llama-arch.cpp | 7 +++ src/llama-arch.h | 3 + src/llama-model.h | 6 ++ src/models/dflash.cpp | 112 +++++++++++++++++++++++++++++++++ tools/cli/README.md | 2 +- tools/server/README.md | 2 +- tools/server/server-schema.cpp | 1 + 14 files changed, 286 insertions(+), 33 deletions(-) diff --git a/common/common.h b/common/common.h index 2792521836ae..78b1f416e085 100644 --- a/common/common.h +++ b/common/common.h @@ -173,6 +173,7 @@ enum common_speculative_type { COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3, // Eagle3 speculative decoding COMMON_SPECULATIVE_TYPE_DRAFT_MTP, // Multi-token prediction COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH, // DFlash speculative decoding + COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK, // DSpark speculative decoding (DFlash + Markov head) COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE, // simple self-speculative decoding based on n-grams COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K, // self-speculative decoding with n-gram keys only COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K4V, // self-speculative decoding with n-gram keys and 4 m-gram values @@ -388,7 +389,7 @@ struct common_params_speculative { uint32_t need_n_rs_seq() const { bool needs_rs_seq = std::any_of(types.begin(), types.end(), [&](auto t) { - return t == COMMON_SPECULATIVE_TYPE_DRAFT_MTP || t == COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3 || t == COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH; + return t == COMMON_SPECULATIVE_TYPE_DRAFT_MTP || t == COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3 || t == COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH || t == COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK; }); return needs_rs_seq ? draft.n_max : 0u; diff --git a/common/speculative.cpp b/common/speculative.cpp index 3a6c07368358..5653a90b889c 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -34,6 +34,7 @@ const std::map common_speculative_type_fro {"draft-eagle3", COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3}, {"draft-mtp", COMMON_SPECULATIVE_TYPE_DRAFT_MTP}, {"draft-dflash", COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH}, + {"draft-dspark", COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK}, {"ngram-simple", COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE}, {"ngram-map-k", COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K}, {"ngram-map-k4v", COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K4V}, @@ -928,15 +929,20 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { int32_t block_size = 0; llama_token mask_token_id = 0; + // draft-dspark: the draft carries a Markov head and uses an anchor-first block layout + const bool is_dspark; + const int32_t * target_layer_ids = nullptr; // model_dft's extract layer indices uint32_t target_layer_ids_n = 0; // scratch buffer for concatenated target features [n_tokens, n_embd_enc] std::vector features_buf; - common_speculative_impl_draft_dflash(const common_params_speculative & params, uint32_t n_seq) - : common_speculative_impl(COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH, n_seq) + common_speculative_impl_draft_dflash(const common_params_speculative & params, uint32_t n_seq, + common_speculative_type type = COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH) + : common_speculative_impl(type, n_seq) , params(params.draft) + , is_dspark(type == COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK) { auto * ctx_tgt = this->params.ctx_tgt; auto * ctx_dft = this->params.ctx_dft; @@ -963,16 +969,18 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { } mask_token_id = llama_vocab_mask(llama_model_get_vocab(model_dft)); - LOG_INF("%s: adding speculative implementation 'draft-dflash'\n", __func__); + LOG_INF("%s: adding speculative implementation '%s'\n", __func__, common_speculative_type_to_str(type).c_str()); LOG_INF("%s: - n_max=%d, n_min=%d, p_min=%.2f\n", __func__, this->params.n_max, this->params.n_min, this->params.p_min); LOG_INF("%s: - block_size=%d, mask_token_id=%d, n_extract=%u\n", __func__, block_size, mask_token_id, target_layer_ids_n); - // DFlash input is [id_last, * (block_size-1)], so it can draft at most block_size-1 tokens per step - if (this->params.n_max > block_size - 1 || this->params.n_min > block_size - 1) { - LOG_WRN("%s: requested draft size (n_max=%d, n_min=%d) exceeds the trained DFlash block size %d -- clamping to %d\n", - __func__, this->params.n_max, this->params.n_min, block_size, block_size - 1); - this->params.n_max = std::min(this->params.n_max, block_size - 1); - this->params.n_min = std::min(this->params.n_min, block_size - 1); + // DFlash input is [id_last, * (block_size-1)]: in-place denoising yields at most + // block_size-1 draft tokens, DSpark yield a full block_size draft tokens + const int32_t n_draft_max = is_dspark ? block_size : block_size - 1; + if (this->params.n_max > n_draft_max || this->params.n_min > n_draft_max) { + LOG_WRN("%s: requested draft size (n_max=%d, n_min=%d) exceeds the trained block size %d -- clamping to %d\n", + __func__, this->params.n_max, this->params.n_min, block_size, n_draft_max); + this->params.n_max = std::min(this->params.n_max, n_draft_max); + this->params.n_min = std::min(this->params.n_min, n_draft_max); } batch = llama_batch_init(llama_n_batch(ctx_dft), 0, n_seq); @@ -1136,12 +1144,9 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { const int32_t n = (int32_t) dp.n_past; - int32_t n_draft = params.n_max; - if (dp.n_max > 0) { - n_draft = std::min(n_draft, dp.n_max); - } + const int32_t n_draft = params.n_max; - const int32_t n_block_tokens = n_draft + 1; // id_last + n_draft * + const int32_t n_block_tokens = n_draft + (is_dspark ? 0 : 1); i_block_beg[seq_id] = batch.n_tokens; n_block [seq_id] = n_block_tokens; for (int32_t i = 0; i < n_block_tokens; ++i) { @@ -1173,27 +1178,57 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { auto & result = *dp.result; - // greedily read the predicted block at this sequence's noise positions 1..n_block_tokens-1 - for (int32_t i = 1; i < n_block_tokens; ++i) { - common_sampler_sample(smpl, ctx_dft, beg + i, true); + if (is_dspark) { + // DSpark predicts the next token from position 0 and optionally truncates + // at the first position below the confidence threshold. + const float * conf = params.p_min > 0.0f ? llama_get_embeddings_nextn(ctx_dft) : nullptr; - const auto * cur_p = common_sampler_get_candidates(smpl, true); + for (int32_t i = 0; i < n_block_tokens; ++i) { + const int32_t idx = beg + i; - for (int k = 0; k < std::min(3, (int) cur_p->size); ++k) { - LOG_DBG(" - seq_id %d, draft candidate %3d, pos %3d: %6d (%8.3f) '%s'\n", - seq_id, k, i - 1, cur_p->data[k].id, cur_p->data[k].p, - common_token_to_piece(ctx_dft, cur_p->data[k].id).c_str()); - } + if (conf && conf[(size_t) idx * n_embd_dec] < params.p_min) { + break; + } - const llama_token id = cur_p->data[0].id; + common_sampler_sample(smpl, ctx_dft, idx, true); - if (cur_p->data[0].p < params.p_min) { - break; + const auto * cur_p = common_sampler_get_candidates(smpl, true); + + for (int k = 0; k < std::min(3, (int) cur_p->size); ++k) { + LOG_DBG(" - seq_id %d, draft candidate %3d, pos %3d: %6d (%8.3f) '%s'\n", + seq_id, k, i, cur_p->data[k].id, cur_p->data[k].p, + common_token_to_piece(ctx_dft, cur_p->data[k].id).c_str()); + } + + const llama_token id = cur_p->data[0].id; + + common_sampler_accept(smpl, id, true); + + result.push_back(id); } + } else { + // greedily read the predicted block at this sequence's noise positions 1..n_block_tokens-1 + for (int32_t i = 1; i < n_block_tokens; ++i) { + common_sampler_sample(smpl, ctx_dft, beg + i, true); - common_sampler_accept(smpl, id, true); + const auto * cur_p = common_sampler_get_candidates(smpl, true); - result.push_back(id); + for (int k = 0; k < std::min(3, (int) cur_p->size); ++k) { + LOG_DBG(" - seq_id %d, draft candidate %3d, pos %3d: %6d (%8.3f) '%s'\n", + seq_id, k, i - 1, cur_p->data[k].id, cur_p->data[k].p, + common_token_to_piece(ctx_dft, cur_p->data[k].id).c_str()); + } + + const llama_token id = cur_p->data[0].id; + + if (cur_p->data[0].p < params.p_min) { + break; + } + + common_sampler_accept(smpl, id, true); + + result.push_back(id); + } } if (result.size() < (size_t) params.n_min) { @@ -2155,6 +2190,7 @@ std::string common_speculative_type_to_str(common_speculative_type type) { case COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3: return "draft-eagle3"; case COMMON_SPECULATIVE_TYPE_DRAFT_MTP: return "draft-mtp"; case COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH: return "draft-dflash"; + case COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK: return "draft-dspark"; case COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE: return "ngram-simple"; case COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K: return "ngram-map-k"; case COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K4V: return "ngram-map-k4v"; @@ -2208,6 +2244,7 @@ int32_t common_speculative_n_max(const common_params_speculative * spec) { case COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3: case COMMON_SPECULATIVE_TYPE_DRAFT_MTP: case COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH: + case COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK: n_max = std::max(n_max, std::max(0, spec->draft.n_max)); break; case COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE: @@ -2352,6 +2389,7 @@ common_speculative * common_speculative_init(common_params_speculative & params, bool has_draft_eagle3 = (enabled_configs & (1u << COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3)) && params.draft.ctx_dft != nullptr; bool has_draft_mtp = (enabled_configs & (1u << COMMON_SPECULATIVE_TYPE_DRAFT_MTP)) && params.draft.ctx_dft != nullptr; bool has_draft_dflash = (enabled_configs & (1u << COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH)) && params.draft.ctx_dft != nullptr; + bool has_draft_dspark = (enabled_configs & (1u << COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK)) && params.draft.ctx_dft != nullptr; @@ -2362,7 +2400,7 @@ common_speculative * common_speculative_init(common_params_speculative & params, bool has_ngram_mod = (enabled_configs & (1u << COMMON_SPECULATIVE_TYPE_NGRAM_MOD)); // when adding a new type - update here the logic above - static_assert(COMMON_SPECULATIVE_TYPE_COUNT == 10); + static_assert(COMMON_SPECULATIVE_TYPE_COUNT == 11); // this list here defines the priority of the speculators // the one with highest priority are listed first @@ -2395,6 +2433,9 @@ common_speculative * common_speculative_init(common_params_speculative & params, if (has_draft_dflash) { configs.push_back(common_speculative_config(COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH, params)); } + if (has_draft_dspark) { + configs.push_back(common_speculative_config(COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK, params)); + } } std::vector> impls = {}; @@ -2419,6 +2460,11 @@ common_speculative * common_speculative_init(common_params_speculative & params, impls.push_back(std::make_unique(config.params, n_seq)); break; } + case COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK: { + impls.push_back(std::make_unique( + config.params, n_seq, COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK)); + break; + } case COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE: { common_ngram_map ngram_map = get_common_ngram_map(config.type, config.params.ngram_simple); diff --git a/conversion/__init__.py b/conversion/__init__.py index 45c001b78f96..1a47b851a0e6 100644 --- a/conversion/__init__.py +++ b/conversion/__init__.py @@ -53,6 +53,7 @@ "DeepseekV3ForCausalLM": "deepseek", "DeepseekV32ForCausalLM": "deepseek", "DFlashDraftModel": "qwen", + "Qwen3DSparkModel": "qwen", "DeepseekV4ForCausalLM": "deepseek", "DistilBertForMaskedLM": "bert", "DistilBertForSequenceClassification": "bert", diff --git a/conversion/qwen.py b/conversion/qwen.py index 9bc2b99fde5d..d1127f7431f6 100644 --- a/conversion/qwen.py +++ b/conversion/qwen.py @@ -688,3 +688,23 @@ def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Ca if not name.startswith("model."): name = "model." + name return super().filter_tensors((name, gen)) + + +@ModelBase.register("Qwen3DSparkModel") +class DSparkModel(DFlashModel): + # DSpark = DFlash + a semi-autoregressive Markov head + model_arch = gguf.MODEL_ARCH.DFLASH + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + # normalize the flat DeepSpec schema to DFlash's nested dflash_config + self.hparams.setdefault("dflash_config", { + k: self.hparams[k] for k in ("target_layer_ids", "mask_token_id") if k in self.hparams + }) + + @classmethod + def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None: + name, gen = item + if name.endswith(("embed_tokens.weight", "lm_head.weight")): + return None + return super().filter_tensors((name, gen)) diff --git a/docs/speculative.md b/docs/speculative.md index 4100b92f8f18..3957db85c9c1 100644 --- a/docs/speculative.md +++ b/docs/speculative.md @@ -78,6 +78,38 @@ See: - #22105 +### DSpark (`draft-dspark`) + +DSpark extends DFlash with a semi-autoregressive _Markov head_: the draft still emits a whole +block per forward pass, but each block position's logits are biased by a low-rank term keyed on +the previous token, chained in-graph across the block. This keeps drafting at one decode per +block while recovering some of the left-to-right signal that pure block diffusion loses. + +The draft is a small DeepSpec checkpoint trained for a specific target (for example +[`deepseek-ai/dspark_qwen3_4b_block7`](https://huggingface.co/deepseek-ai/dspark_qwen3_4b_block7) +for `Qwen/Qwen3-4B`). Convert it with `--target-model-dir` so it inherits the target's tokenizer +and token embeddings: + +```bash +python convert_hf_to_gguf.py deepseek-ai/dspark_qwen3_4b_block7 \ + --target-model-dir Qwen/Qwen3-4B --outtype bf16 --outfile Qwen3-4B-DSpark.gguf + +llama-server -m Qwen3-4B.gguf -md Qwen3-4B-DSpark.gguf \ + --spec-type draft-dspark --spec-draft-n-max 7 -fa on --jinja +``` + +`--spec-draft-n-max` is clamped to the draft model's trained block size. + +`--spec-draft-conf-min P` truncates each drafted block at the first position whose predicted +acceptance (from the draft's confidence head, if present) falls below `P` (default 0 = disabled). + +Currently only drafts with a Qwen3 backbone are supported; support for other backbones +(e.g. Gemma4) is planned. + +See: + +- #25173 + ### n-gram Cache (`ngram-cache`) An n-gram is a sequence of n tokens. The n-gram cache implementation maintains statistics about short n-gram sequences. @@ -173,7 +205,7 @@ If a draft model is combined with a draftless decoding the draftless decoding ha ### General Speculative Parameters ``` ---spec-type [none|draft-simple|draft-eagle3|draft-dflash|draft-mtp|ngram-cache|ngram-simple|ngram-map-k|ngram-map-k4v|ngram-mod] +--spec-type [none|draft-simple|draft-eagle3|draft-dflash|draft-dspark|draft-mtp|ngram-cache|ngram-simple|ngram-map-k|ngram-map-k4v|ngram-mod] comma-separated list of types of speculative decoding to use (default: none) (env: LLAMA_ARG_SPEC_TYPE) @@ -314,6 +346,7 @@ Specifies a comma-separated list of speculative decoding types to use. | `draft-simple` | Use a simple draft model for speculation | | `draft-eagle3` | Use an EAGLE-3 draft model that reads the target's hidden states | | `draft-dflash` | Use a DFlash block-diffusion draft model that emits a block per step | +| `draft-dspark` | Use a DSpark draft model (DFlash backbone + semi-autoregressive Markov head) | | `draft-mtp` | Use Multi Token Prediction (MTP) heads from the main model | | `ngram-cache` | Use n-gram cache lookup | | `ngram-simple` | Use simple n-gram pattern matching | diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index 4df94b9a3add..2ebefaa3b422 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -988,6 +988,10 @@ class MODEL_TENSOR(IntEnum): # eagle3 FC = auto() # feature fusion layer D2T = auto() # draft to target vocabulary mapping + # dspark + DSPARK_MARKOV_W1 = auto() # markov head: prev-token embed + DSPARK_MARKOV_W2 = auto() # markov head: bias projection + DSPARK_CONF_PROJ = auto() # confidence head # lfm2 audio A_ENC_NORM_CONV = auto() A_ENC_LINEAR_POS = auto() @@ -1617,6 +1621,9 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.NEXTN_SHARED_HEAD_HEAD: "blk.{bid}.nextn.shared_head_head", MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM: "blk.{bid}.nextn.shared_head_norm", MODEL_TENSOR.FC: "fc", + MODEL_TENSOR.DSPARK_MARKOV_W1: "markov_w1", + MODEL_TENSOR.DSPARK_MARKOV_W2: "markov_w2", + MODEL_TENSOR.DSPARK_CONF_PROJ: "conf_proj", MODEL_TENSOR.D2T: "d2t", } @@ -4362,6 +4369,10 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.FFN_UP, MODEL_TENSOR.FC, MODEL_TENSOR.ENC_OUTPUT_NORM, + # optional DSpark heads + MODEL_TENSOR.DSPARK_MARKOV_W1, + MODEL_TENSOR.DSPARK_MARKOV_W2, + MODEL_TENSOR.DSPARK_CONF_PROJ, ], MODEL_ARCH.MISTRAL4: [ MODEL_TENSOR.TOKEN_EMBD, diff --git a/gguf-py/gguf/tensor_mapping.py b/gguf-py/gguf/tensor_mapping.py index 8299ac25b432..5562d43277a4 100644 --- a/gguf-py/gguf/tensor_mapping.py +++ b/gguf-py/gguf/tensor_mapping.py @@ -1304,6 +1304,18 @@ class TensorNameMap: "model.fc", # dflash ), + MODEL_TENSOR.DSPARK_MARKOV_W1: ( + "model.markov_head.markov_w1", # dspark + ), + + MODEL_TENSOR.DSPARK_MARKOV_W2: ( + "model.markov_head.markov_w2", # dspark + ), + + MODEL_TENSOR.DSPARK_CONF_PROJ: ( + "model.confidence_head.proj", # dspark + ), + MODEL_TENSOR.CLS: ( "classifier", # jina "classifier.dense", # roberta diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index a0945e501d53..e81ff647eee4 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -616,6 +616,9 @@ static const std::map LLM_TENSOR_NAMES = { { LLM_TENSOR_MASKED_EMBD_ORDERING, "masked_embd_ordering" }, { LLM_TENSOR_FC, "fc" }, { LLM_TENSOR_D2T, "d2t" }, + { LLM_TENSOR_DSPARK_MARKOV_W1, "markov_w1" }, + { LLM_TENSOR_DSPARK_MARKOV_W2, "markov_w2" }, + { LLM_TENSOR_DSPARK_CONF_PROJ, "conf_proj" }, }; // declare information about the model weight tensors: @@ -870,6 +873,10 @@ static const std::map LLM_TENSOR_INFOS = { // eagle3 {LLM_TENSOR_FC, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}}, {LLM_TENSOR_D2T, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_GET_ROWS}}, + // dspark + {LLM_TENSOR_DSPARK_MARKOV_W1, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_GET_ROWS}}, + {LLM_TENSOR_DSPARK_MARKOV_W2, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_DSPARK_CONF_PROJ, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}}, }; LLM_KV::LLM_KV(llm_arch arch, const char * suffix) : arch(arch), suffix(suffix) {} diff --git a/src/llama-arch.h b/src/llama-arch.h index fb5eb3920bcd..cbc97085ea79 100644 --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -624,6 +624,9 @@ enum llm_tensor { LLM_TENSOR_MASKED_EMBD_ORDERING, LLM_TENSOR_FC, LLM_TENSOR_D2T, + LLM_TENSOR_DSPARK_MARKOV_W1, + LLM_TENSOR_DSPARK_MARKOV_W2, + LLM_TENSOR_DSPARK_CONF_PROJ, }; diff --git a/src/llama-model.h b/src/llama-model.h index 36d0480e5eb7..d6a40fa30204 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -606,6 +606,12 @@ struct llama_model { struct ggml_tensor * fc = nullptr; // feature fusion layer struct ggml_tensor * d2t = nullptr; // draft to target vocabulary mapping + // dspark + struct ggml_tensor * dspark_markov_w1 = nullptr; + struct ggml_tensor * dspark_markov_w2 = nullptr; + struct ggml_tensor * dspark_conf_proj = nullptr; + struct ggml_tensor * dspark_conf_proj_b = nullptr; + // unified vector to store target-model extracted layer ids in eagle3, dflash, etc. std::vector target_layer_ids; diff --git a/src/models/dflash.cpp b/src/models/dflash.cpp index 427eed4594e3..dcff3aec9fa2 100644 --- a/src/models/dflash.cpp +++ b/src/models/dflash.cpp @@ -37,6 +37,23 @@ void llama_model_dflash::load_arch_tensors(llama_model_loader &) { const int64_t n_embd_inp = hparams.n_embd_inp_enc(); + // DSpark = DFlash + a semi-autoregressive Markov head and Confidence head + // + // TODO: only Qwen3-style backbones are supported for now; other backbones (e.g. Gemma4) + // need their own conversion path and graph tweaks + const struct ggml_tensor * markov_meta = ml->get_tensor_meta("markov_w1.weight"); + if (markov_meta) { + const int64_t dspark_markov_rank = markov_meta->ne[0]; + + dspark_markov_w1 = create_tensor(tn(LLM_TENSOR_DSPARK_MARKOV_W1, "weight"), { dspark_markov_rank, n_vocab }, 0); + dspark_markov_w2 = create_tensor(tn(LLM_TENSOR_DSPARK_MARKOV_W2, "weight"), { dspark_markov_rank, n_vocab }, 0); + + dspark_conf_proj = create_tensor(tn(LLM_TENSOR_DSPARK_CONF_PROJ, "weight"), { n_embd + dspark_markov_rank, 1 }, 0); + dspark_conf_proj_b = create_tensor(tn(LLM_TENSOR_DSPARK_CONF_PROJ, "bias"), { 1 }, TENSOR_NOT_REQUIRED); + + LLAMA_LOG_INFO("%s: DFlash with DSpark markov head (rank = %lld)\n", __func__, (long long) dspark_markov_rank); + } + fc = create_tensor(tn(LLM_TENSOR_FC, "weight"), { n_embd_inp, n_embd }, 0); output_norm_enc = create_tensor(tn(LLM_TENSOR_ENC_OUTPUT_NORM, "weight"), { n_embd }, 0); // encoder hidden_norm (after fc) output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), { n_embd }, 0); // decoder final norm @@ -105,6 +122,94 @@ llama_model_dflash::graph::graph(const llama_model & model, const llm_grap ggml_build_forward_expand(gf, cur); } +// DSpark (DFlash + Markov & Confidence head): Markov bias on the draft logits, chained per block position +static void build_dspark_markov_head(llm_graph_context & g, const llama_model & model, ggml_tensor * tokens) { + ggml_context * ctx0 = g.ctx0; + auto & res = g.res; + + ggml_tensor * w1 = model.dspark_markov_w1; + ggml_tensor * w2 = model.dspark_markov_w2; + GGML_ASSERT(w1 && w2 && model.dspark_conf_proj && "DSpark markov/confidence weights not loaded"); + + ggml_tensor * base = res->t_logits; // [n_vocab, n_tokens] + const int64_t n_vocab = base->ne[0]; + const int64_t n_tok = base->ne[1]; + + const auto it = model.gguf_kv.find("dflash.block_size"); + GGML_ASSERT(it != model.gguf_kv.end() && "DSpark draft requires 'dflash.block_size' in GGUF metadata"); + const int64_t block_size = std::stoi(it->second); + GGML_ASSERT(block_size > 0); + + const int64_t n_blocks = g.ubatch.n_seqs_unq; + GGML_ASSERT(n_blocks > 0 && n_tok % n_blocks == 0 && "DSpark markov head requires equal-size blocks"); + // runtime tokens per block in this ubatch (anchor + drafted positions), bounded by training block_size + const int64_t block_drafts = n_tok / n_blocks; + if (block_drafts > block_size) { + return; + } + + // anchor (committed last) token of every block: token 0 of each block, i.e. a strided view + const size_t token_stride = (size_t) block_drafts * tokens->nb[0]; + const size_t base_stride = (size_t) block_drafts * base->nb[1]; + + ggml_tensor * prev = ggml_view_2d(ctx0, tokens, 1, n_blocks, token_stride, 0); + prev = ggml_cont_1d(ctx0, prev, n_blocks); + + // confidence head input: predicts per-position acceptance + ggml_tensor * conf_inp = res->t_embd; // [n_embd, n_tok] + + ggml_tensor * cat = nullptr; + ggml_tensor * cat_conf = nullptr; + + // TODO: the in-graph chain is greedy (argmax); sampling params affect only the final + // token pick, not the Markov conditioning path + for (int64_t i = 0; i < block_drafts; ++i) { + ggml_tensor * w1_prev = ggml_get_rows(ctx0, w1, prev); // [R, n_blocks] + ggml_tensor * bias = ggml_mul_mat(ctx0, w2, w1_prev); // [n_vocab, n_blocks] + + // position i of every block: strided view [n_vocab, n_blocks] + ggml_tensor * base_i = ggml_view_2d(ctx0, base, n_vocab, n_blocks, base_stride, i*base->nb[1]); + ggml_tensor * col = ggml_add(ctx0, base_i, bias); + + cat = cat ? ggml_concat(ctx0, cat, col, 1) : col; + + // conf(i) = sigmoid(conf_proj . [conf_inp(i); markov_w1[prev(i)]] + b) -- [1, n_blocks] + ggml_tensor * conf_inp_i = ggml_view_2d(ctx0, conf_inp, conf_inp->ne[0], n_blocks, + (size_t) block_drafts * conf_inp->nb[1], i*conf_inp->nb[1]); + ggml_tensor * feat = ggml_concat(ctx0, ggml_cont(ctx0, conf_inp_i), w1_prev, 0); + ggml_tensor * conf = ggml_mul_mat(ctx0, model.dspark_conf_proj, feat); + if (model.dspark_conf_proj_b) { + conf = ggml_add(ctx0, conf, model.dspark_conf_proj_b); + } + conf = ggml_sigmoid(ctx0, conf); + + cat_conf = cat_conf ? ggml_concat(ctx0, cat_conf, conf, 1) : conf; + + if (i + 1 < block_drafts) { + prev = ggml_argmax(ctx0, col); + } + } + + // cat is position-major; restore ubatch block-major order + ggml_tensor * out = ggml_reshape_3d(ctx0, cat, n_vocab, n_blocks, block_drafts); + out = ggml_cont(ctx0, ggml_permute(ctx0, out, 0, 2, 1, 3)); // [n_vocab, block_drafts, n_blocks] + out = ggml_reshape_2d(ctx0, out, n_vocab, n_tok); + + { + ggml_tensor * conf = ggml_reshape_3d(ctx0, cat_conf, 1, n_blocks, block_drafts); + conf = ggml_cont(ctx0, ggml_permute(ctx0, conf, 0, 2, 1, 3)); + conf = ggml_reshape_2d(ctx0, conf, 1, n_tok); + + // note: broadcast the [1, n_tok] confidences to n_embd-wide rows to be able to reuse `llama_get_embeddings_nextn` + conf = ggml_repeat(ctx0, conf, res->t_embd); + res->t_h_nextn = conf; + ggml_build_forward_expand(g.gf, conf); + } + + res->t_logits = out; + ggml_build_forward_expand(g.gf, out); +} + // DFlash decoder, dual-mode by batch type: // * embd batch -> fused target features: project + inject K/V into the cache. // * token batch -> noise-block diffusion: attend over [committed, MASK...] to generate draft tokens @@ -210,6 +315,8 @@ llama_model_dflash::graph::graph(const llama_model & model, const llm_gra inp->tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens); ggml_set_input(inp->tokens); + ggml_tensor * inp_tokens = inp->tokens; + ggml_tensor * inpL = ggml_get_rows(ctx0, tok_embd, inp->tokens); cb(inpL, "inp_noise_embd", -1); @@ -290,4 +397,9 @@ llama_model_dflash::graph::graph(const llama_model & model, const llm_gra res->t_logits = cur; ggml_build_forward_expand(gf, cur); + + // DSpark: bias the draft logits with the Markov head + if (model.dspark_markov_w1) { + build_dspark_markov_head(*this, model, inp_tokens); + } } diff --git a/tools/cli/README.md b/tools/cli/README.md index 972ea04dc7d7..bcddd05702bb 100644 --- a/tools/cli/README.md +++ b/tools/cli/README.md @@ -203,7 +203,7 @@ | `--spec-draft-device, -devd, --device-draft ` | comma-separated list of devices to use for offloading the draft model (none = don't offload)
use --list-devices to see a list of available devices | | `--spec-draft-ngl, -ngld, --gpu-layers-draft, --n-gpu-layers-draft N` | max. number of draft model layers to store in VRAM, either an exact number, 'auto', or 'all' (default: auto)
(env: LLAMA_ARG_N_GPU_LAYERS_DRAFT) | | `--spec-draft-model, -md, --model-draft FNAME` | draft model for speculative decoding (default: unused)
(env: LLAMA_ARG_SPEC_DRAFT_MODEL) | -| `--spec-type none,draft-simple,draft-eagle3,draft-mtp,draft-dflash,ngram-simple,ngram-map-k,ngram-map-k4v,ngram-mod,ngram-cache` | comma-separated list of types of speculative decoding to use (default: none)

(env: LLAMA_ARG_SPEC_TYPE) | +| `--spec-type none,draft-simple,draft-eagle3,draft-mtp,draft-dflash,draft-dspark,ngram-simple,ngram-map-k,ngram-map-k4v,ngram-mod,ngram-cache` | comma-separated list of types of speculative decoding to use (default: none)

(env: LLAMA_ARG_SPEC_TYPE) | | `--spec-ngram-mod-n-min N` | minimum number of ngram tokens to use for ngram-based speculative decoding (default: 48) | | `--spec-ngram-mod-n-max N` | maximum number of ngram tokens to use for ngram-based speculative decoding (default: 64) | | `--spec-ngram-mod-n-match N` | ngram-mod lookup length (default: 24) | diff --git a/tools/server/README.md b/tools/server/README.md index 25aacf9f516f..f45c018972d2 100644 --- a/tools/server/README.md +++ b/tools/server/README.md @@ -259,7 +259,7 @@ For the full list of features, please refer to [server's changelog](https://gith | `--spec-draft-device, -devd, --device-draft ` | comma-separated list of devices to use for offloading the draft model (none = don't offload)
use --list-devices to see a list of available devices | | `--spec-draft-ngl, -ngld, --gpu-layers-draft, --n-gpu-layers-draft N` | max. number of draft model layers to store in VRAM, either an exact number, 'auto', or 'all' (default: auto)
(env: LLAMA_ARG_N_GPU_LAYERS_DRAFT) | | `--spec-draft-model, -md, --model-draft FNAME` | draft model for speculative decoding (default: unused)
(env: LLAMA_ARG_SPEC_DRAFT_MODEL) | -| `--spec-type none,draft-simple,draft-eagle3,draft-mtp,draft-dflash,ngram-simple,ngram-map-k,ngram-map-k4v,ngram-mod,ngram-cache` | comma-separated list of types of speculative decoding to use (default: none)

(env: LLAMA_ARG_SPEC_TYPE) | +| `--spec-type none,draft-simple,draft-eagle3,draft-mtp,draft-dflash,draft-dspark,ngram-simple,ngram-map-k,ngram-map-k4v,ngram-mod,ngram-cache` | comma-separated list of types of speculative decoding to use (default: none)

(env: LLAMA_ARG_SPEC_TYPE) | | `--spec-ngram-mod-n-min N` | minimum number of ngram tokens to use for ngram-based speculative decoding (default: 48) | | `--spec-ngram-mod-n-max N` | maximum number of ngram tokens to use for ngram-based speculative decoding (default: 64) | | `--spec-ngram-mod-n-match N` | ngram-mod lookup length (default: 24) | diff --git a/tools/server/server-schema.cpp b/tools/server/server-schema.cpp index e880f4ca728d..674d3ba337bc 100644 --- a/tools/server/server-schema.cpp +++ b/tools/server/server-schema.cpp @@ -209,6 +209,7 @@ std::vector> make_llama_cmpl_schema(const common_params & ->set_hard_limits(0.0f, 1.0f) ->set_desc("Minimum speculative decoding probability for draft tokens (0 = greedy)")); + add((new field_str("speculative.type")) ->set_desc("Speculative decoding method (for debugging and research purposes)") ->set_handler([&](field_eval_context & ctx, const json & data) { From b62b3509813dd3169663885975c2306e96df2242 Mon Sep 17 00:00:00 2001 From: Bhavik Sharda <10757940+BLSharda@users.noreply.github.com> Date: Tue, 28 Jul 2026 17:33:42 +0530 Subject: [PATCH 036/190] ggml-cuda: add chunked SSD matmul for Mamba-2 prefill acceleration (#22675) * ggml-cuda: add chunked SSD matmul for Mamba-2 prefill acceleration * cuda: added SSD CICD fixes for CUDA / HIP / MUSA / MSVC. * ggml-cuda: review comments fixed. * ggml-cuda: Fuse M matrix materialization into pre_matmul kernel and enabled test. * ggml-cuda: test updates and fixes * ggml-cuda: test updates to remove hardcoding of tensor initialise data limits. * ggml-cuda: ssd minor review comment fixed. * ggml-cuda: ssd minor CICD fixed. * CUDA SSD: Fixes correctness by promoting s0_stride_seq to int64_t, improves memory coalescing in ssm_ssd_prepare_dt_kernel, and boosts efficiency by merging B_weighted and C_scaled; also addresses prior review comments. * cuda: fix sdata read-write race in prepare_dt fallback scan loop --- ggml/src/ggml-cuda/ssm-scan.cu | 481 +++++++++++++++++++++++++++++++++ tests/test-backend-ops.cpp | 21 +- 2 files changed, 499 insertions(+), 3 deletions(-) diff --git a/ggml/src/ggml-cuda/ssm-scan.cu b/ggml/src/ggml-cuda/ssm-scan.cu index 3022249c77d5..f3418c2af83d 100644 --- a/ggml/src/ggml-cuda/ssm-scan.cu +++ b/ggml/src/ggml-cuda/ssm-scan.cu @@ -9,6 +9,21 @@ using namespace cub; #include "ssm-scan.cuh" + +// Minimum number of tokens to use SSD (State Space Duality) matmul path instead of scan path. +// For n_tok <= this threshold, the scan kernel is used (lower overhead for short sequences). +#define SSM_SSD_MIN_TOKENS 128 + +// prepare_dt kernel dimensions: one block per (head, seq), each block handles DT_MAX_ITEMS items. +#define SSM_SSD_DT_BLOCK 256 +#define SSM_SSD_DT_MAX_ITEMS 32 + +// Maximum tokens the SSD path supports, derived from the prepare_dt kernel block capacity. +#define SSM_SSD_MAX_TOKENS (SSM_SSD_DT_BLOCK * SSM_SSD_DT_MAX_ITEMS) + +// Chunk size for chunked SSD. Caps matmul cost at O(chunk^2) per chunk. +#define SSM_SSD_CHUNK_SIZE 256 + // We would like to keep pragma unroll for cases where L_template is not 0, // so we suppress the clang transformation warning. #ifdef __clang__ @@ -316,6 +331,429 @@ static void ssm_scan_f32_cuda(const float * src0, const float * src1, const floa } } +#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) +// ============================================================================ +// SSD (State Space Duality) kernels for Mamba-2 prefill (n_tok > SSM_SSD_MIN_TOKENS) +// +// Instead of a sequential scan, SSD reformulates the output as: +// Y = (L (.) (C @ B^T)) @ (X * dt) + decay * C @ s_init +// where L is a causal decay mask derived from A and dt. +// +// This converts the O(T*N) sequential scan into parallel matmuls. +// ============================================================================ +// Softplus(dt) and inclusive prefix sum per head using CUB BlockScan. +// Grid: (n_head, n_seqs) +template +__global__ void ssm_ssd_prepare_dt_kernel( + const float * __restrict__ dt_raw, + float * __restrict__ dt_sp_out, + float * __restrict__ cs_out, + const int n_head, const int n_tok, + const int dt_stride_tok, // elements between tokens in dt + const int dt_stride_seq) { // elements between sequences in dt + + const int h = blockIdx.x; + const int s = blockIdx.y; + + const float * dt_seq = dt_raw + s * dt_stride_seq; + + float * dt_sp_seq = dt_sp_out + s * n_tok * n_head; + float * cs_seq = cs_out + s * n_tok * n_head; + + const int items_per_thread = (n_tok + BLOCK_SIZE - 1) / BLOCK_SIZE; + + // Phase 1: softplus with interleaved distribution (t = i*BLOCK_SIZE + threadIdx.x). + // Each warp reads BLOCK_SIZE consecutive tokens, giving coalesced dt_raw loads + // (stride n_head between threads vs. items_per_thread*n_head in blocked layout). + float local_vals[MAX_ITEMS]; + for (int i = 0; i < items_per_thread; i++) { + const int t = i * BLOCK_SIZE + threadIdx.x; + if (t < n_tok) { + float val = dt_seq[h + t * dt_stride_tok]; + float sp = (val <= 20.0f) ? log1pf(expf(val)) : val; + local_vals[i] = sp; + dt_sp_seq[t * n_head + h] = sp; + } else { + local_vals[i] = 0.0f; + } + } + + // Phase 2+3: per-step inclusive scan to build cs[] in token order. + // With interleaved distribution the per-thread total scan would not give token-order + // prefix sums, so we scan one BLOCK_SIZE slab at a time and carry a running total. +#ifdef USE_CUB + using BlockScan = cub::BlockScan; + __shared__ typename BlockScan::TempStorage scan_temp; + __shared__ float step_total; + + float running = 0.0f; + for (int i = 0; i < items_per_thread; i++) { + float inclusive; + BlockScan(scan_temp).InclusiveSum(local_vals[i], inclusive); + const int t = i * BLOCK_SIZE + threadIdx.x; + if (t < n_tok) { + cs_seq[t * n_head + h] = running + inclusive; + } + if (threadIdx.x == BLOCK_SIZE - 1) { + step_total = inclusive; + } + __syncthreads(); + running += step_total; + } +#else + // Fallback: sequential prefix scan in shared memory, one slab at a time. + __shared__ float sdata[BLOCK_SIZE]; + float running = 0.0f; + for (int i = 0; i < items_per_thread; i++) { + const int t = i * BLOCK_SIZE + threadIdx.x; + sdata[threadIdx.x] = local_vals[i]; + __syncthreads(); + if (threadIdx.x == 0) { + for (int j = 1; j < BLOCK_SIZE; j++) { + sdata[j] += sdata[j - 1]; + } + } + __syncthreads(); + if (t < n_tok) { + cs_seq[t * n_head + h] = running + sdata[threadIdx.x]; + } + running += sdata[BLOCK_SIZE - 1]; + __syncthreads(); + } +#endif +} + +// Prepare SSD matmul inputs for one chunk: X_dt, B_weighted, C_scaled. +// T_matmul controls precision for X_dt, B_weighted (float or half). +// C_scaled is always float (pairs with float s_cur in step 3c). +// Computation is always FP32; only the final store converts to T_matmul. +// Also materializes the causal M matrix = exp(A*(cs_out - cs_in)) * CB (fused with prep to save a launch). +// Grid: (ceil(max(C*head_dim, d_state*C, chunk_len^2) / BLOCK), n_head, n_seqs) +template +__global__ void ssm_ssd_pre_matmul_kernel( + const float * __restrict__ cs, // {n_tok, n_head} cumulative dt sums + const float * __restrict__ dt_sp, // {n_tok, n_head} softplus(dt) + const float * __restrict__ A, // {1, n_head} + const float * __restrict__ x, // {head_dim, n_head, n_tok, n_seqs} + const float * __restrict__ B, // {d_state, n_group, n_tok, n_seqs} + const float * __restrict__ C_src, // {d_state, n_group, n_tok, n_seqs} + T_matmul * __restrict__ X_dt, // {head_dim, C, n_head} x * dt, d-fastest + T_matmul * __restrict__ B_weighted, // {d_state, C, n_head} B * decay_from_end + float * __restrict__ C_scaled, // {d_state, C, n_head} C * decay_to_pos (always float) + const float * __restrict__ CB, // {chunk_len, chunk_len, n_group, n_seqs} + half * __restrict__ M_out, // {chunk_len, chunk_len, n_head, n_seqs} + const int chunk_len, const int head_dim, const int n_head, const int n_group, + const int d_state, const int A_stride, + const int x_stride_tok, const int x_stride_seq, + const int B_stride_tok, const int B_stride_seq, + const int C_stride_tok, const int C_stride_seq, + const int chunk_offset, + const int n_tok_total) { + + const int h = blockIdx.y; + const int s = blockIdx.z; + const int g = h / (n_head / n_group); + + const float A_h = A[h * A_stride]; + const int idx = blockIdx.x * BLOCK_SIZE + threadIdx.x; + + const int cs_seq_off = s * n_tok_total * n_head; + const float cs_base = (chunk_offset > 0) ? cs[cs_seq_off + (chunk_offset - 1) * n_head + h] : 0.0f; + const float cs_last = cs[cs_seq_off + (chunk_offset + chunk_len - 1) * n_head + h] - cs_base; + + // Prepare X_dt = x * dt, stored d-fastest for coalesced reads and writes. + const int n_xdt = chunk_len * head_dim; + if (idx < n_xdt) { + const int d = idx % head_dim; + const int t = idx / head_dim; + + const float x_val = x[s * x_stride_seq + (chunk_offset + t) * x_stride_tok + d + h * head_dim]; + const float dt_val = dt_sp[cs_seq_off + (chunk_offset + t) * n_head + h]; + + X_dt[d + t * head_dim + h * n_xdt + s * n_xdt * n_head] = (T_matmul)(x_val * dt_val); + } + + // Prepare B_weighted and C_scaled together: both share the same index space (d_state * chunk_len) + // and the same cs_t load, so merging halves the cs[] global memory traffic. + const int n_bw = d_state * chunk_len; + if (idx < n_bw) { + const int n = idx % d_state; + const int t = idx / d_state; + + const float cs_t = cs[cs_seq_off + (chunk_offset + t) * n_head + h] - cs_base; + + const float B_val = B[s * B_stride_seq + (chunk_offset + t) * B_stride_tok + g * d_state + n]; + B_weighted[n + t * d_state + h * n_bw + s * n_bw * n_head] = (T_matmul)(B_val * __expf(A_h * (cs_last - cs_t))); + + const float C_val = C_src[s * C_stride_seq + (chunk_offset + t) * C_stride_tok + g * d_state + n]; + C_scaled[n + t * d_state + h * n_bw + s * n_bw * n_head] = C_val * __expf(A_h * cs_t); + } + + // Materialize M = exp(A*(cs_out - cs_in)) * CB with causal mask. + const int n_M = chunk_len * chunk_len; + if (idx < n_M) { + const int t_out = idx % chunk_len; + const int t_in = idx / chunk_len; + + half val; + if (t_in <= t_out) { + const float cs_out = cs[cs_seq_off + (chunk_offset + t_out) * n_head + h] - cs_base; + const float cs_in = cs[cs_seq_off + (chunk_offset + t_in) * n_head + h] - cs_base; + const float decay = __expf(A_h * (cs_out - cs_in)); + const float * CB_g = CB + (int64_t)s * chunk_len * chunk_len * n_group + + (int64_t)g * chunk_len * chunk_len; + const float cb_val = CB_g[t_out + t_in * chunk_len]; + val = __float2half(decay * cb_val); + } else { + val = __float2half(0.0f); + } + + M_out[(int64_t)s * n_M * n_head + (int64_t)h * n_M + t_in * chunk_len + t_out] = val; + } +} + +// Scale running state in-place: s_cur *= decay_total(chunk). +// Called BEFORE cuBLAS state update (beta=1) to fuse inter-chunk decay. +// Eliminates the s_old buffer and D2D memcpy vs the old approach of: +// memcpy(s_old, s_cur) -> cuBLAS(beta=0) -> s_cur += decay * s_old +// Grid: (ceil(d_state * head_dim / BLOCK), n_head, n_seqs) +template +__global__ void ssm_ssd_scale_state_kernel( + float * __restrict__ s_cur, // {d_state, head_dim, n_head, n_seqs} + const float * __restrict__ cs, // {n_tok, n_head} cumulative dt sums + const float * __restrict__ A, // {1, n_head} + const int d_state, const int head_dim, const int n_head, + const int chunk_offset, const int chunk_len, + const int n_tok_total, const int A_stride) { + + const int h = blockIdx.y; + const int s = blockIdx.z; + const int idx = blockIdx.x * BLOCK_SIZE + threadIdx.x; + const int state_per_head = d_state * head_dim; + if (idx >= state_per_head) return; + + const float A_h = A[h * A_stride]; + const int cs_seq_off = s * n_tok_total * n_head; + const float cs_base = (chunk_offset > 0) ? cs[cs_seq_off + (chunk_offset - 1) * n_head + h] : 0.0f; + const float cs_last = cs[cs_seq_off + (chunk_offset + chunk_len - 1) * n_head + h] - cs_base; + const float decay_total = __expf(A_h * cs_last); + + const int off = s * state_per_head * n_head + h * state_per_head + idx; + s_cur[off] *= decay_total; +} + +// Copy initial state from src0[ids[s]] into s_cur for each sequence. +// Grid: (ceil(d_state * head_dim * n_head / BLOCK), n_seqs) +template +__global__ void ssm_ssd_init_state_kernel( + const float * __restrict__ src0, // {d_state, head_dim, n_head, n_rs} + const int32_t * __restrict__ ids, // {n_seqs} + float * __restrict__ s_cur, // {d_state, head_dim, n_head, n_seqs} + const int state_size, // d_state * head_dim * n_head + const int64_t s0_stride_seq) { // elements between state rows + const int s = blockIdx.y; + const int idx = blockIdx.x * BLOCK_SIZE + threadIdx.x; + if (idx >= state_size) return; + + const float * s_src = src0 + (int64_t)ids[s] * s0_stride_seq; + s_cur[s * state_size + idx] = s_src[idx]; +} + +// SSD (State Space Duality) dispatch for Mamba-2 prefill. +// Chunked matmuls: CB, materialize M + cuBLAS Y, S@C, B@X_dt. +// All strides are in elements (floats), not bytes. +static void ssm_scan_ssd_f32_cuda( + ggml_backend_cuda_context & ctx, + const float * src0_d, const float * src1_d, const float * src2_d, const float * src3_d, + const float * src4_d, const float * src5_d, const int32_t * src6_d, float * dst_d, + const int64_t s0_stride_seq, // state (src0) stride between seqs + const int x_stride_tok, const int x_stride_seq, // x (src1) strides + const int dt_stride_tok, const int dt_stride_seq, // dt (src2) strides + const int A_stride, // A (src3) stride between heads + const int B_stride_tok, const int B_stride_seq, // B (src4) strides + const int C_stride_tok, const int C_stride_seq, // C (src5) strides + const int64_t s_off, const int64_t d_state, const int64_t head_dim, + const int64_t n_head, const int64_t n_group, const int64_t n_tok, const int64_t n_seq) { + + cudaStream_t stream = ctx.stream(); + const int64_t d_inner = head_dim * n_head; + + const int64_t chunk_size = SSM_SSD_CHUNK_SIZE; + const int64_t n_chunks = (n_tok + chunk_size - 1) / chunk_size; + + const int64_t state_per_head = d_state * head_dim; + + using matmul_t = half; + static constexpr cudaDataType_t matmul_dtype = CUDA_R_16F; + + ggml_cuda_pool_alloc dt_sp_buf(ctx.pool(), n_tok * n_head * n_seq); + ggml_cuda_pool_alloc cs_buf(ctx.pool(), n_tok * n_head * n_seq); + ggml_cuda_pool_alloc CB_buf(ctx.pool(), chunk_size * chunk_size * n_group * n_seq); + ggml_cuda_pool_alloc X_dt_buf(ctx.pool(), chunk_size * head_dim * n_head * n_seq); + ggml_cuda_pool_alloc B_w_buf(ctx.pool(), d_state * chunk_size * n_head * n_seq); + ggml_cuda_pool_alloc C_s_buf(ctx.pool(), d_state * chunk_size * n_head * n_seq); + float * dt_sp = dt_sp_buf.get(); + float * cs = cs_buf.get(); + float * CB = CB_buf.get(); + matmul_t * X_dt = X_dt_buf.get(); + matmul_t * B_weighted = B_w_buf.get(); + float * C_scaled = C_s_buf.get(); + float * s_cur = (float *)((char *)dst_d + s_off); // write state directly to dst + + // Step 1: softplus(dt) and parallel prefix sum over full sequence + { + dim3 grid(n_head, n_seq); + ssm_ssd_prepare_dt_kernel<<>>( + src2_d, dt_sp, cs, n_head, n_tok, dt_stride_tok, dt_stride_seq); + CUDA_CHECK(cudaGetLastError()); + } + + // Step 2: initialize running state from src0[ids[s]] + { + constexpr int BLOCK = 256; + const int64_t state_size = d_state * head_dim * n_head; + dim3 grid((state_size + BLOCK - 1) / BLOCK, n_seq); + ssm_ssd_init_state_kernel<<>>( + src0_d, src6_d, s_cur, state_size, s0_stride_seq); + CUDA_CHECK(cudaGetLastError()); + } + + // Step 3: chunked SSD loop + // Per chunk: pre_matmul (incl. M) + 4 cuBLAS (CB, Y, S@C, state update) + scale_state + cublasHandle_t handle = ctx.cublas_handle(); + CUBLAS_CHECK(cublasSetStream(handle, stream)); + const float alpha_one = 1.0f; + const float beta_zero = 0.0f; + const float beta_one = 1.0f; + const int lda_C_src = C_stride_tok; // leading dim for C in CB = C^T @ B + const int ldb_B_src = B_stride_tok; // leading dim for B in CB = C^T @ B + + // Scratch buffer for causal M matrix, reused across chunks (max size at chunk_size) + const int64_t n_M_max = chunk_size * chunk_size; + ggml_cuda_pool_alloc M_buf(ctx.pool(), n_M_max * n_head * n_seq); + half * M_mat = M_buf.get(); + + for (int64_t k = 0; k < n_chunks; k++) { + const int64_t chunk_offset = k * chunk_size; + const int64_t chunk_len = (chunk_offset + chunk_size <= n_tok) ? chunk_size : (n_tok - chunk_offset); + + // 3a: CB = C^T @ B per group + for (int64_t s = 0; s < n_seq; s++) { + const float * C_s = src5_d + s * C_stride_seq + chunk_offset * C_stride_tok; + const float * B_s = src4_d + s * B_stride_seq + chunk_offset * B_stride_tok; + float * CB_s = CB + s * chunk_len * chunk_len * n_group; + + if (n_group == 1) { + CUBLAS_CHECK(cublasSgemm(handle, CUBLAS_OP_T, CUBLAS_OP_N, + chunk_len, chunk_len, d_state, + &alpha_one, C_s, lda_C_src, B_s, ldb_B_src, + &beta_zero, CB_s, (int)chunk_len)); + } else { + CUBLAS_CHECK(cublasGemmStridedBatchedEx(handle, CUBLAS_OP_T, CUBLAS_OP_N, + chunk_len, chunk_len, d_state, + &alpha_one, + C_s, CUDA_R_32F, lda_C_src, d_state, + B_s, CUDA_R_32F, ldb_B_src, d_state, + &beta_zero, + CB_s, CUDA_R_32F, (int)chunk_len, (long long)(chunk_len * chunk_len), + n_group, + CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT)); + } + } + + // 3b: prepare X_dt, B_weighted, C_scaled + materialize causal M matrix + const int64_t n_M = chunk_len * chunk_len; + { + constexpr int BLOCK = 256; + const int64_t n_xdt = chunk_len * head_dim; + const int64_t n_bw = d_state * chunk_len; + int64_t max_work = n_xdt; + if (n_bw > max_work) max_work = n_bw; + if (n_M > max_work) max_work = n_M; + dim3 grid((max_work + BLOCK - 1) / BLOCK, n_head, n_seq); + ssm_ssd_pre_matmul_kernel<<>>( + cs, dt_sp, src3_d, src1_d, src4_d, src5_d, + X_dt, B_weighted, C_scaled, + CB, M_mat, + chunk_len, head_dim, n_head, n_group, d_state, A_stride, + x_stride_tok, x_stride_seq, B_stride_tok, B_stride_seq, C_stride_tok, C_stride_seq, + chunk_offset, n_tok); + CUDA_CHECK(cudaGetLastError()); + } + + // 3c: dst = S_cur^T @ C_scaled (state contribution) + { + const int64_t stride_S = state_per_head; + const int64_t stride_Cs = d_state * chunk_len; + + for (int64_t s = 0; s < n_seq; s++) { + float * dst_chunk = dst_d + s * d_inner * n_tok + chunk_offset * d_inner; + + CUBLAS_CHECK(cublasGemmStridedBatchedEx(handle, CUBLAS_OP_T, CUBLAS_OP_N, + head_dim, chunk_len, d_state, + &alpha_one, + s_cur + s * stride_S * n_head, CUDA_R_32F, d_state, stride_S, + C_scaled + s * stride_Cs * n_head, CUDA_R_32F, d_state, stride_Cs, + &beta_zero, + dst_chunk, CUDA_R_32F, d_inner, head_dim, + n_head, + CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT)); + } + } + + // 3d: dst += X_dt @ M^T (intra-chunk contribution, adds to 3c result) + // M is stored as M[t_out, t_in] (lower-triangular), transpose needed for Y = X @ M^T. + { + const int64_t stride_M = n_M; + const int64_t stride_X_h = (int64_t)chunk_len * head_dim; + + for (int64_t s = 0; s < n_seq; s++) { + float * dst_chunk = dst_d + s * d_inner * n_tok + chunk_offset * d_inner; + CUBLAS_CHECK(cublasGemmStridedBatchedEx(handle, CUBLAS_OP_N, CUBLAS_OP_T, + head_dim, chunk_len, chunk_len, + &alpha_one, + X_dt + s * stride_X_h * n_head, matmul_dtype, head_dim, stride_X_h, + M_mat + s * stride_M * n_head, matmul_dtype, chunk_len, stride_M, + &beta_one, + dst_chunk, CUDA_R_32F, d_inner, head_dim, + n_head, + CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT)); + } + } + + // 3e: s_cur = B_weighted @ X_dt^T + decay_total * s_cur_old (state update) + { + // Scale s_cur in-place by per-head decay_total BEFORE cuBLAS overwrites it + constexpr int BLOCK = 256; + dim3 grid((state_per_head + BLOCK - 1) / BLOCK, n_head, n_seq); + ssm_ssd_scale_state_kernel<<>>( + s_cur, cs, src3_d, + d_state, head_dim, n_head, + chunk_offset, chunk_len, n_tok, A_stride); + CUDA_CHECK(cudaGetLastError()); + + // cuBLAS with beta=1: s_cur = B_weighted @ X_dt^T + 1.0 * s_cur (already scaled) + const int64_t stride_Bw = d_state * chunk_len; + const int64_t stride_X = chunk_len * head_dim; + const int64_t stride_S = state_per_head; + + for (int64_t s = 0; s < n_seq; s++) { + // X_dt is d-fastest {hd, C}, read as OP_T to get {C, hd} + CUBLAS_CHECK(cublasGemmStridedBatchedEx(handle, CUBLAS_OP_N, CUBLAS_OP_T, + d_state, head_dim, chunk_len, + &alpha_one, + B_weighted + s * stride_Bw * n_head, matmul_dtype, d_state, stride_Bw, + X_dt + s * stride_X * n_head, matmul_dtype, head_dim, stride_X, + &beta_one, + s_cur + s * stride_S * n_head, CUDA_R_32F, d_state, stride_S, + n_head, + CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT)); + } + } + } +} +#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) + void ggml_cuda_op_ssm_scan(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const struct ggml_tensor * src0 = dst->src[0]; // s const struct ggml_tensor * src1 = dst->src[1]; // x @@ -357,6 +795,49 @@ void ggml_cuda_op_ssm_scan(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { GGML_ASSERT(src6->type == GGML_TYPE_I32); GGML_ASSERT(dst->type == GGML_TYPE_F32); + // Byte strides are narrowed to int for both scan and SSD paths. + GGML_ASSERT(src0->nb[2] <= (size_t)INT_MAX); + GGML_ASSERT(src0->nb[3] <= (size_t)INT_MAX); + GGML_ASSERT(src1->nb[2] <= (size_t)INT_MAX); + GGML_ASSERT(src1->nb[3] <= (size_t)INT_MAX); + GGML_ASSERT(src2->nb[1] <= (size_t)INT_MAX); + GGML_ASSERT(src2->nb[2] <= (size_t)INT_MAX); + GGML_ASSERT(src3->nb[1] <= (size_t)INT_MAX); + GGML_ASSERT(src4->nb[2] <= (size_t)INT_MAX); + GGML_ASSERT(src4->nb[3] <= (size_t)INT_MAX); + GGML_ASSERT(src5->nb[2] <= (size_t)INT_MAX); + GGML_ASSERT(src5->nb[3] <= (size_t)INT_MAX); + +#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) + // Mamba-2 with scalar A per head: use SSD matmul path for long sequences. + // Requires NVIDIA Turing+ otherwise fallback to scan. + const bool is_mamba2 = (src3->nb[1] == sizeof(float)); + const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; + const bool use_ssd = is_mamba2 && n_t > SSM_SSD_MIN_TOKENS + && n_t <= SSM_SSD_MAX_TOKENS + && GGML_CUDA_CC_IS_NVIDIA(cc) + && cc >= GGML_CUDA_CC_TURING + && nr % 8 == 0; // cuBLAS requires 8-element (16-byte) alignment + + if (use_ssd) { + // ssm_ssd_init_state_kernel uses flat linear indexing within each sequence, + // so src0 must be fully contiguous across all inner dimensions. + // The scan path handles non-contiguous nb[2] via src0_nb2 but does not handle nb[1]. + GGML_ASSERT(src0->nb[1] == nc * sizeof(float)); + GGML_ASSERT(src0->nb[2] == nc * nr * sizeof(float)); + + ssm_scan_ssd_f32_cuda(ctx, + src0_d, src1_d, src2_d, src3_d, src4_d, src5_d, src6_d, dst_d, + (int64_t)(src0->nb[3] / sizeof(float)), + (int)(src1->nb[2] / sizeof(float)), (int)(src1->nb[3] / sizeof(float)), + (int)(src2->nb[1] / sizeof(float)), (int)(src2->nb[2] / sizeof(float)), + (int)(src3->nb[1] / sizeof(float)), + (int)(src4->nb[2] / sizeof(float)), (int)(src4->nb[3] / sizeof(float)), + (int)(src5->nb[2] / sizeof(float)), (int)(src5->nb[3] / sizeof(float)), + s_off, nc, nr, nh, ng, n_t, n_s); + return; + } +#endif ssm_scan_f32_cuda(src0_d, src1_d, src2_d, src3_d, src4_d, src5_d, src6_d, dst_d, src0->nb[2], src0->nb[3], src1->nb[2], src1->nb[3], src2->nb[1], src2->nb[2], src3->nb[1], src4->nb[2], src4->nb[3], src5->nb[2], src5->nb[3], diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index f6b60e52a400..a5b660f47a0a 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -4000,7 +4000,7 @@ struct test_ssm_scan : public test_case { test_ssm_scan(ggml_type type = GGML_TYPE_F32, int64_t d_state = 32, - int64_t head_dim = 1, // non-zero for Mamba-2 + int64_t head_dim = 1, // 1 = Mamba-1; > 1 = Mamba-2 (scalar A per head) int64_t n_head = 32, int64_t n_group = 1, int64_t n_seq_tokens = 32, @@ -4008,6 +4008,11 @@ struct test_ssm_scan : public test_case { bool xbc_overlap = false) : type(type), d_state(d_state), head_dim(head_dim), n_head(n_head), n_group(n_group), n_seq_tokens(n_seq_tokens), n_seqs(n_seqs), xbc_overlap(xbc_overlap) {} + double max_nmse_err() override { + // SSD path (head_dim > 1) uses FP16 intermediates (M matrix, X_dt); Mamba-1 is pure FP32. + return (head_dim > 1) ? 2e-7 : 1e-7; + } + ggml_tensor * build_graph(ggml_context * ctx) override { ggml_tensor * s = ggml_new_tensor_4d(ctx, type, d_state, head_dim, n_head, n_seqs); ggml_tensor * dt = ggml_new_tensor_3d(ctx, type, n_head, n_seq_tokens, n_seqs); @@ -4034,14 +4039,14 @@ struct test_ssm_scan : public test_case { return out; } - // similar to test_mul_mat_id + void initialize_tensors(ggml_context * ctx) override { std::random_device rd; std::default_random_engine rng(rd()); for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) { if (t->type == GGML_TYPE_I32) { if (ggml_is_view_op(t->op)) { continue; } - // ids + // ids: permutation of [0..n_seqs) for (int64_t r = 0; r < ggml_nrows(t); r++) { std::vector data(t->ne[0]); for (int i = 0; i < t->ne[0]; i++) { @@ -4050,6 +4055,11 @@ struct test_ssm_scan : public test_case { std::shuffle(data.begin(), data.end(), rng); ggml_backend_tensor_set(t, data.data(), r * t->nb[1], t->ne[0] * sizeof(int32_t)); } + } else if (ggml_is_view_op(t->op)) { + continue; + } else if (t->ne[1] == n_head && t->ne[2] == 1) { + // A {1 or d_state, n_head}: negative decay (2-D tensor, ne[2]==1 distinguishes from 3-D/4-D tensors) + init_tensor_uniform(t, -1.0f, -0.5f); } else { init_tensor_uniform(t); } @@ -8770,6 +8780,9 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 128, 64, 16, 2, 32, 4)); // Mamba-2 test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 256, 64, 8, 2, 32, 4)); // Falcon-H1 test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 128, 128, 4, 4, 16, 2, true)); // x/B/C overlap + test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 128, 80, 128, 1, 256, 1)); // Nemotron-9B SSD path + test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 128, 80, 128, 1, 512, 1)); // Nemotron-9B SSD multi-chunk (2 aligned chunks) + test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 128, 64, 80, 8, 300, 2)); // Mamba-2 SSD multi-chunk (partial 2nd chunk, 2 seqs) test_cases.emplace_back(new test_rwkv_wkv6(GGML_TYPE_F32, 32, 64, 1, 1)); test_cases.emplace_back(new test_rwkv_wkv6(GGML_TYPE_F32, 32, 64, 32, 1)); @@ -9990,6 +10003,8 @@ static std::vector> make_test_cases_perf() { test_cases.emplace_back(new test_ssm_conv_bias_silu(GGML_TYPE_F32, {4, 3328, 1, 1}, {4, 3328, 1, 1}, true)); // generate test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 128, 64, 48, 1, 512, 1)); // prefill test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 128, 64, 48, 1, 1, 1)); // generate + test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 128, 80, 128, 1, 512, 1)); // Nemotron-9B prefill + test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 128, 80, 128, 1, 1, 1)); // Nemotron-9B generate // acc test_cases.emplace_back(new test_acc(GGML_TYPE_F32, {256, 17, 1, 1}, {256, 16, 1, 1}, -1)); From 81616410050cb5d8b733b39863807f4852591d8a Mon Sep 17 00:00:00 2001 From: Jeff Bolz Date: Tue, 28 Jul 2026 13:06:03 +0100 Subject: [PATCH 037/190] vulkan: add iq4_nl support back to FA (#24585) * vulkan: add iq4_nl support back to FA I was originally concerned about wasting shared memory on the LUT, but it's small and unlikely to matter in practice. Also support q1_0 for non-coopmat2. Fixes #23681 * remove q1_0 FA support --- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 45 +++++++++++-------- .../vulkan-shaders/flash_attn.comp | 4 +- .../vulkan-shaders/flash_attn_base.glsl | 11 ++++- .../vulkan-shaders/flash_attn_cm1.comp | 4 +- .../vulkan-shaders/flash_attn_cm2.comp | 40 ++++++++++------- .../vulkan-shaders/flash_attn_dequant.glsl | 15 +++++++ .../vulkan-shaders/vulkan-shaders-gen.cpp | 2 + 7 files changed, 82 insertions(+), 39 deletions(-) diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index 5dcf4503bbee..49fd6e0b42f6 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -3490,7 +3490,7 @@ struct vk_fa_tuning_params { }; static bool ggml_vk_flash_attn_scalar_shmem_support(const vk_device& device, const vk_fa_tuning_params& params, uint32_t hsk, uint32_t hsv, bool f32acc, ggml_type k_type, ggml_type v_type); -static bool ggml_vk_flash_attn_coopmat_shmem_support(const vk_device& device, const vk_fa_tuning_params& params, uint32_t hsk, uint32_t hsv, bool f32acc, ggml_type k_type = GGML_TYPE_F16); +static bool ggml_vk_flash_attn_coopmat_shmem_support(const vk_device& device, const vk_fa_tuning_params& params, uint32_t hsk, uint32_t hsv, bool f32acc, ggml_type k_type = GGML_TYPE_F16, ggml_type v_type = GGML_TYPE_F16); static vk_fa_tuning_params get_fa_tuning_params_scalar(const vk_device& device, uint32_t hsk, uint32_t hsv, uint32_t n_rows, uint32_t n_kv, ggml_type k_type, ggml_type v_type, bool f32acc) { @@ -3646,7 +3646,7 @@ static vk_fa_tuning_params get_fa_tuning_params(const vk_device& device, uint32_ bool shape_ok = (f32acc && device->coopmat_support_16x16x16_f32acc) || (!f32acc && device->coopmat_support_16x16x16_f16acc); const vk_fa_tuning_params params = get_fa_tuning_params_coopmat1(device, hsk, hsv, n_rows, n_kv, k_type, v_type, f32acc); - bool shmem_ok = ggml_vk_flash_attn_coopmat_shmem_support(device, params, hsk, hsv, f32acc, k_type); + bool shmem_ok = ggml_vk_flash_attn_coopmat_shmem_support(device, params, hsk, hsv, f32acc, k_type, v_type); if (!shape_ok || !shmem_ok) { path = FA_SCALAR; @@ -3658,11 +3658,6 @@ static vk_fa_tuning_params get_fa_tuning_params(const vk_device& device, uint32_ path = FA_SCALAR; } - // Q1_0 K/V is only implemented on coopmat2 (flash_attn_cm2); there is no scalar FA shader for it. - if ((k_type == GGML_TYPE_Q1_0 || v_type == GGML_TYPE_Q1_0) && device->coopmat2) { - path = FA_COOPMAT2; - } - switch (path) { case FA_SCALAR: return get_fa_tuning_params_scalar(device, hsk, hsv, n_rows, n_kv, k_type, v_type, f32acc); @@ -3904,16 +3899,27 @@ static uint32_t get_subgroup_size(const std::string &pipeline_name, const vk_dev return 0; // If no matching configuration is found } -// Whether scalar flash attention will use the MMQ path for the given k_type. -static bool ggml_vk_fa_scalar_uses_mmq(const vk_device& device, ggml_type k_type) { +// Whether scalar flash attention will use the MMQ path for the given K/V types. +static bool ggml_vk_fa_type_needs_shmem(ggml_type type) { + switch (type) { + case GGML_TYPE_IQ4_NL: + return true; + default: + return false; + } +} + +static bool ggml_vk_fa_scalar_uses_mmq(const vk_device& device, ggml_type k_type, ggml_type v_type) { #if defined(GGML_VULKAN_INTEGER_DOT_GLSLC_SUPPORT) return device->integer_dot_product && device->subgroup_clustered && + !ggml_vk_fa_type_needs_shmem(v_type) && (k_type == GGML_TYPE_Q4_0 || k_type == GGML_TYPE_Q4_1 || k_type == GGML_TYPE_Q5_0 || k_type == GGML_TYPE_Q5_1 || k_type == GGML_TYPE_Q8_0); #else GGML_UNUSED(device); GGML_UNUSED(k_type); + GGML_UNUSED(v_type); return false; #endif } @@ -4246,7 +4252,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { const bool fa_ds = fa.first.subgroup_size == 0; const bool bf16_kv = fa.first.k_type == GGML_TYPE_BF16; - const bool use_mmq = ggml_vk_fa_scalar_uses_mmq(device, fa.first.k_type); + const bool use_mmq = ggml_vk_fa_scalar_uses_mmq(device, fa.first.k_type, fa.first.v_type); const void * spv_data = nullptr; size_t spv_size = 0; const char *name = nullptr; @@ -10380,7 +10386,6 @@ static void ggml_vk_mul_mat_id(ggml_backend_vk_context * ctx, vk_context& subctx static bool ggml_vk_flash_attn_scalar_shmem_support(const vk_device& device, const vk_fa_tuning_params& params, uint32_t hsk, uint32_t hsv, bool f32acc, ggml_type k_type, ggml_type v_type) { GGML_UNUSED(f32acc); - GGML_UNUSED(v_type); // Needs to be kept up to date on shader changes const uint32_t wg_size = params.workgroup_size; const uint32_t Br = params.block_rows; @@ -10389,13 +10394,15 @@ static bool ggml_vk_flash_attn_scalar_shmem_support(const vk_device& device, con // BF16 uses the fp32 shader (FLOAT_TYPE=float) const uint32_t float_type_size = (device->fp16 && k_type != GGML_TYPE_BF16) ? sizeof(ggml_fp16_t) : sizeof(float); - const bool mmq = ggml_vk_fa_scalar_uses_mmq(device, k_type); + const bool mmq = ggml_vk_fa_scalar_uses_mmq(device, k_type, v_type); // tmpsh is overestimated slightly const uint32_t tmpsh = wg_size * sizeof(float); const uint32_t tmpshv4 = wg_size * 4 * float_type_size; const uint32_t masksh = Bc * (Br + 1) * float_type_size; + // DATA_A_IQ4_NL is compiled into the FA shaders unconditionally, so its shared table is always allocated. + const uint32_t iq_shmem = 16 * float_type_size; uint32_t Qf, kvsh, kblocksh_size; if (mmq) { @@ -10420,7 +10427,7 @@ static bool ggml_vk_flash_attn_scalar_shmem_support(const vk_device& device, con kblocksh_size = 0; } - const uint32_t total_size = tmpsh + tmpshv4 + masksh + Qf + kvsh + kblocksh_size; + const uint32_t total_size = tmpsh + tmpshv4 + masksh + iq_shmem + Qf + kvsh + kblocksh_size; const bool supported = total_size <= device->properties.limits.maxComputeSharedMemorySize; VK_LOG_DEBUG("ggml_vk_flash_attn_scalar_shmem_support(HSK=" << hsk << ", HSV=" << hsv << ", mmq=" << mmq << ", total_size=" << total_size << ", supported=" << supported); @@ -10428,7 +10435,8 @@ static bool ggml_vk_flash_attn_scalar_shmem_support(const vk_device& device, con return supported; } -static bool ggml_vk_flash_attn_coopmat_shmem_support(const vk_device& device, const vk_fa_tuning_params& params, uint32_t hsk, uint32_t hsv, bool f32acc, ggml_type k_type) { +static bool ggml_vk_flash_attn_coopmat_shmem_support(const vk_device& device, const vk_fa_tuning_params& params, uint32_t hsk, uint32_t hsv, bool f32acc, ggml_type k_type, ggml_type v_type) { + GGML_UNUSED(v_type); // Needs to be kept up to date on shader changes const uint32_t Br = params.block_rows; const uint32_t Bc = params.block_cols; @@ -10444,6 +10452,8 @@ static bool ggml_vk_flash_attn_coopmat_shmem_support(const vk_device& device, co const uint32_t f16vec4 = 8; const uint32_t tmpsh = (Bc / MatBc) * sizeof(float); + // DATA_A_IQ4_NL is compiled into the FA shaders unconditionally, so its shared table is always allocated. + const uint32_t iq_shmem = 16 * sizeof(ggml_fp16_t); const uint32_t qstride = hsk_pad / 4 + 2; const uint32_t Qf = Br * qstride * f16vec4; @@ -10465,7 +10475,7 @@ static bool ggml_vk_flash_attn_coopmat_shmem_support(const vk_device& device, co const uint32_t slope = Br * acctype; - const uint32_t total_size = tmpsh + Qf + Psh + sfsh + ksh + pvsh + slope; + const uint32_t total_size = tmpsh + iq_shmem + Qf + Psh + sfsh + ksh + pvsh + slope; const bool supported = total_size <= device->properties.limits.maxComputeSharedMemorySize; VK_LOG_DEBUG("ggml_vk_flash_attn_coopmat_shmem_support(HSK=" << hsk << ", HSV=" << hsv << ", f32acc=" << f32acc << ", total_size=" << total_size << ", supported=" << supported); @@ -17617,7 +17627,7 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm if (op->src[3] && op->src[3]->type != GGML_TYPE_F16) { return false; } - auto fa_kv_ok = [coopmat2](ggml_type t) { + auto fa_kv_ok = [](ggml_type t) { switch (t) { case GGML_TYPE_F32: case GGML_TYPE_F16: @@ -17627,9 +17637,8 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm case GGML_TYPE_Q5_0: case GGML_TYPE_Q4_1: case GGML_TYPE_Q4_0: + case GGML_TYPE_IQ4_NL: return true; - case GGML_TYPE_Q1_0: - return coopmat2; default: return false; } diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn.comp index 3192130ccf57..6c264c78619f 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn.comp @@ -80,7 +80,9 @@ shared vec4 occupancy_limiter[LIMIT_OCCUPANCY_SHMEM > 0 ? LIMIT_OCCUPANCY_SHMEM void main() { #ifdef NEEDS_INIT_IQ_SHMEM - init_iq_shmem(gl_WorkGroupSize); + if (fa_type_needs_shmem(FaTypeK) || fa_type_needs_shmem(FaTypeV)) { + init_iq_shmem(gl_WorkGroupSize); + } #endif init_indices(); diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_base.glsl b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_base.glsl index 66dcf6102190..3c64f91dad36 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_base.glsl +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_base.glsl @@ -97,8 +97,8 @@ layout (binding = 6) readonly buffer MO {uint32_t data_mask_opt[];}; #define FA_TYPE_Q5_0 6u #define FA_TYPE_Q5_1 7u #define FA_TYPE_Q8_0 8u +#define FA_TYPE_IQ4_NL 20u #define FA_TYPE_BF16 30u -#define FA_TYPE_Q1_0 41u #if defined(BFLOAT16) #define O_TYPE float @@ -120,8 +120,8 @@ uint fa_block_elems(uint ty) { case FA_TYPE_Q5_0: return uint(QUANT_K_Q5_0); case FA_TYPE_Q5_1: return uint(QUANT_K_Q5_1); case FA_TYPE_Q8_0: return uint(QUANT_K_Q8_0); + case FA_TYPE_IQ4_NL: return uint(QUANT_K_IQ4_NL); case FA_TYPE_BF16: return 1u; - case FA_TYPE_Q1_0: return uint(QUANT_K_Q1_0); // cm2-only, harmless elsewhere default: return 1u; } } @@ -140,6 +140,13 @@ uint fa_quant_r_mmq(uint ty) { } } +bool fa_type_needs_shmem(uint ty) { + switch (ty) { + case FA_TYPE_IQ4_NL: return true; + default: return false; + } +} + // These can't be `const` globals because GLSL forbids function calls in global // const initializers, even when the spec constants would let the driver fold // them. Macros expand at the use site and fold after specialization. diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp index 16178e577024..057ed739aa8d 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp @@ -64,7 +64,9 @@ shared ACC_TYPE slope[Br]; void main() { #ifdef NEEDS_INIT_IQ_SHMEM - init_iq_shmem(gl_WorkGroupSize); + if (fa_type_needs_shmem(FaTypeK) || fa_type_needs_shmem(FaTypeV)) { + init_iq_shmem(gl_WorkGroupSize); + } #endif init_indices(); diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm2.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm2.comp index b9c03fe499d4..317411153087 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm2.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm2.comp @@ -46,7 +46,7 @@ float16_t faDecodeK(const decodeBufFA_K bl_in, const uint blockCoords[2], const case FA_TYPE_Q5_0: return dequantFuncQ5_0(decodeBufQ5_0(bl_in), blockCoords, coordInBlock); case FA_TYPE_Q5_1: return dequantFuncQ5_1(decodeBufQ5_1(bl_in), blockCoords, coordInBlock); case FA_TYPE_Q8_0: return dequantFuncQ8_0(decodeBufQ8_0(bl_in), blockCoords, coordInBlock); - case FA_TYPE_Q1_0: return dequantFuncQ1_0(decodeBufQ1_0(bl_in), blockCoords, coordInBlock); + case FA_TYPE_IQ4_NL: return dequantFuncIQ4_NL(decodeBufIQ4_NL(bl_in), blockCoords, coordInBlock); default: return float16_t(0); } } @@ -59,7 +59,7 @@ float16_t faDecodeV(const decodeBufFA_V bl_in, const uint blockCoords[2], const case FA_TYPE_Q5_0: return dequantFuncQ5_0(decodeBufQ5_0(bl_in), blockCoords, coordInBlock); case FA_TYPE_Q5_1: return dequantFuncQ5_1(decodeBufQ5_1(bl_in), blockCoords, coordInBlock); case FA_TYPE_Q8_0: return dequantFuncQ8_0(decodeBufQ8_0(bl_in), blockCoords, coordInBlock); - case FA_TYPE_Q1_0: return dequantFuncQ1_0(decodeBufQ1_0(bl_in), blockCoords, coordInBlock); + case FA_TYPE_IQ4_NL: return dequantFuncIQ4_NL(decodeBufIQ4_NL(bl_in), blockCoords, coordInBlock); default: return float16_t(0); } } @@ -67,26 +67,26 @@ float16_t faDecodeV(const decodeBufFA_V bl_in, const uint blockCoords[2], const // V=4 vector decode for K/V; dispatches to per-format _v decoders. f16vec4 faDecodeKVector(const decodeBufFA_K bl_in, const uint blockCoords[2], const uint coordInBlock[2]) { switch (FaTypeK) { - case 0u: return f16vec4(decodeBufF32(bl_in).block); - case 2u: return dequantFuncQ4_0_v(decodeBufQ4_0(bl_in), blockCoords, coordInBlock); - case 3u: return dequantFuncQ4_1_v(decodeBufQ4_1(bl_in), blockCoords, coordInBlock); - case 6u: return dequantFuncQ5_0_v(decodeBufQ5_0(bl_in), blockCoords, coordInBlock); - case 7u: return dequantFuncQ5_1_v(decodeBufQ5_1(bl_in), blockCoords, coordInBlock); - case 8u: return dequantFuncQ8_0_v(decodeBufQ8_0(bl_in), blockCoords, coordInBlock); - case 41u: return dequantFuncQ1_0_v(decodeBufQ1_0(bl_in), blockCoords, coordInBlock); + case FA_TYPE_F32: return f16vec4(decodeBufF32(bl_in).block); + case FA_TYPE_Q4_0: return dequantFuncQ4_0_v(decodeBufQ4_0(bl_in), blockCoords, coordInBlock); + case FA_TYPE_Q4_1: return dequantFuncQ4_1_v(decodeBufQ4_1(bl_in), blockCoords, coordInBlock); + case FA_TYPE_Q5_0: return dequantFuncQ5_0_v(decodeBufQ5_0(bl_in), blockCoords, coordInBlock); + case FA_TYPE_Q5_1: return dequantFuncQ5_1_v(decodeBufQ5_1(bl_in), blockCoords, coordInBlock); + case FA_TYPE_Q8_0: return dequantFuncQ8_0_v(decodeBufQ8_0(bl_in), blockCoords, coordInBlock); + case FA_TYPE_IQ4_NL: return dequantFuncIQ4_NL_v(decodeBufIQ4_NL(bl_in), blockCoords, coordInBlock); default: return f16vec4(0); } } f16vec4 faDecodeVVector(const decodeBufFA_V bl_in, const uint blockCoords[2], const uint coordInBlock[2]) { switch (FaTypeV) { - case 0u: return f16vec4(decodeBufF32(bl_in).block); - case 2u: return dequantFuncQ4_0_v(decodeBufQ4_0(bl_in), blockCoords, coordInBlock); - case 3u: return dequantFuncQ4_1_v(decodeBufQ4_1(bl_in), blockCoords, coordInBlock); - case 6u: return dequantFuncQ5_0_v(decodeBufQ5_0(bl_in), blockCoords, coordInBlock); - case 7u: return dequantFuncQ5_1_v(decodeBufQ5_1(bl_in), blockCoords, coordInBlock); - case 8u: return dequantFuncQ8_0_v(decodeBufQ8_0(bl_in), blockCoords, coordInBlock); - case 41u: return dequantFuncQ1_0_v(decodeBufQ1_0(bl_in), blockCoords, coordInBlock); + case FA_TYPE_F32: return f16vec4(decodeBufF32(bl_in).block); + case FA_TYPE_Q4_0: return dequantFuncQ4_0_v(decodeBufQ4_0(bl_in), blockCoords, coordInBlock); + case FA_TYPE_Q4_1: return dequantFuncQ4_1_v(decodeBufQ4_1(bl_in), blockCoords, coordInBlock); + case FA_TYPE_Q5_0: return dequantFuncQ5_0_v(decodeBufQ5_0(bl_in), blockCoords, coordInBlock); + case FA_TYPE_Q5_1: return dequantFuncQ5_1_v(decodeBufQ5_1(bl_in), blockCoords, coordInBlock); + case FA_TYPE_Q8_0: return dequantFuncQ8_0_v(decodeBufQ8_0(bl_in), blockCoords, coordInBlock); + case FA_TYPE_IQ4_NL: return dequantFuncIQ4_NL_v(decodeBufIQ4_NL(bl_in), blockCoords, coordInBlock); default: return f16vec4(0); } } @@ -169,6 +169,12 @@ ACC_TYPE perElemOpNonGqaSplitKStoreCol0(const in uint32_t r, const in uint32_t c } void main() { +#ifdef NEEDS_INIT_IQ_SHMEM + if (fa_type_needs_shmem(FaTypeK) || fa_type_needs_shmem(FaTypeV)) { + init_iq_shmem(gl_WorkGroupSize); + } +#endif + init_indices(); tensorLayoutNV<2, gl_CooperativeMatrixClampModeConstantNV> tensorLayoutQ = createTensorLayoutNV(2, gl_CooperativeMatrixClampModeConstantNV); @@ -302,7 +308,7 @@ void main() { coopmat K_T; uint32_t k_offset = ik2*p.nb12 + ik3*p.nb13; - // F16: bs_k==1 (direct load). F32: bs_k==4 (vec4 / dequantFuncF32). Q4/Q8 family: bs_k==32. Q1_0: bs_k==128. + // F16: bs_k==1 (direct load). F32: bs_k==4 (vec4 / dequantFuncF32). Quantized types: bs_k==32. #if defined(BFLOAT16) coopMatLoadTensorNV(K_T, data_k, k_offset, sliceTensorLayoutNV(tensorLayoutK, j * Bc, Bc, 0, HSK_pad), tensorViewTranspose); #else diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_dequant.glsl b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_dequant.glsl index 8704479d9600..8ba4725f3342 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_dequant.glsl +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_dequant.glsl @@ -27,6 +27,8 @@ layout (binding = 1) readonly buffer K_PACKED_Q5_1 { block_q5_1_packed16 data[]; layout (binding = 2) readonly buffer V_PACKED_Q5_1 { block_q5_1_packed16 data[]; } v_packed_q5_1; layout (binding = 1) readonly buffer K_PACKED_Q8_0 { block_q8_0_packed16 data[]; } k_packed_q8_0; layout (binding = 2) readonly buffer V_PACKED_Q8_0 { block_q8_0_packed16 data[]; } v_packed_q8_0; +layout (binding = 1) readonly buffer K_PACKED_IQ4_NL { block_iq4_nl_packed16 data[]; } k_packed_iq4_nl; +layout (binding = 2) readonly buffer V_PACKED_IQ4_NL { block_iq4_nl_packed16 data[]; } v_packed_iq4_nl; layout (binding = 1) readonly buffer K_PACKED_BF16 { u16vec4 data[]; } k_packed_bf16; layout (binding = 2) readonly buffer V_PACKED_BF16 { u16vec4 data[]; } v_packed_bf16; @@ -102,6 +104,17 @@ layout (binding = 1) readonly buffer K_PACKED_Q5_1_P32 { block_q5_1_packed32 dat return FLOAT_TYPE(BUF.data[a_offset + ib].d) * FLOAT_TYPEV4(v0.x, v0.y, v1.x, v1.y); \ } +#define FA_DEQUANT4_IQ4_NL(BUF) { \ + const uint shift = (iqs & 0x10) >> 2; \ + const uint qs_i = (iqs & 0xC) >> 1; \ + const uint qsw = uint(BUF.data[a_offset + ib].qs[qs_i]) \ + | (uint(BUF.data[a_offset + ib].qs[qs_i + 1u]) << 16); \ + const FLOAT_TYPE d = FLOAT_TYPE(BUF.data[a_offset + ib].d); \ + const u8vec4 q = unpack8((qsw >> shift) & 0x0F0F0F0Fu); \ + return d * FLOAT_TYPEV4(kvalues_iq4nl[q.x], kvalues_iq4nl[q.y], \ + kvalues_iq4nl[q.z], kvalues_iq4nl[q.w]); \ +} + #define FA_DEQUANT4_BF16(BUF) \ return FLOAT_TYPEV4(bf16_to_fp32(uvec4(BUF.data[(a_offset + ib) / 4]))); @@ -114,6 +127,7 @@ FLOAT_TYPEV4 dequantize4(uint ib, uint iqs, uint a_offset, uint binding_idx) { case FA_TYPE_Q5_0: FA_DEQUANT4_Q5_0(k_packed_q5_0) case FA_TYPE_Q5_1: FA_DEQUANT4_Q5_1(k_packed_q5_1) case FA_TYPE_Q8_0: FA_DEQUANT4_Q8_0(k_packed_q8_0) + case FA_TYPE_IQ4_NL: FA_DEQUANT4_IQ4_NL(k_packed_iq4_nl) case FA_TYPE_BF16: FA_DEQUANT4_BF16(k_packed_bf16) } } else { @@ -124,6 +138,7 @@ FLOAT_TYPEV4 dequantize4(uint ib, uint iqs, uint a_offset, uint binding_idx) { case FA_TYPE_Q5_0: FA_DEQUANT4_Q5_0(v_packed_q5_0) case FA_TYPE_Q5_1: FA_DEQUANT4_Q5_1(v_packed_q5_1) case FA_TYPE_Q8_0: FA_DEQUANT4_Q8_0(v_packed_q8_0) + case FA_TYPE_IQ4_NL: FA_DEQUANT4_IQ4_NL(v_packed_iq4_nl) case FA_TYPE_BF16: FA_DEQUANT4_BF16(v_packed_bf16) } } diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp b/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp index 58d347bc547d..3e6b39576e6d 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp @@ -673,6 +673,8 @@ void process_shaders() { fa_base_dict["ACC_TYPE"] = fp16 && f16acc ? "float16_t" : "float"; fa_base_dict["ACC_TYPEV2"] = fp16 && f16acc ? "f16vec2" : "vec2"; fa_base_dict["ACC_TYPEV4"] = fp16 && f16acc ? "f16vec4" : "vec4"; + // Compile IQ4_NL support into all FA variants so its shared LUT is available when K or V uses it. + fa_base_dict["DATA_A_IQ4_NL"] = "1"; if (fp16 && f16acc) { fa_base_dict["ACC_TYPE_MAX"] = "float16_t(65504.0)"; } From da5b448622ce8f8265bed15a7f80c5cf17894511 Mon Sep 17 00:00:00 2001 From: Aman Gupta Date: Tue, 28 Jul 2026 21:23:24 +0800 Subject: [PATCH 038/190] ggml : set output of view src (#25729) * llama-graph: set_outputs to t->view_src * change set_output to GGML_ASSERT about views not being outputs * sampler : avoid views in outputs * cont : fix dist sampler * cont : consistent logits handling * ggml : set output of view src * graph : simplify set_outputs() * cont : cleanup Co-authored-by: Gaurav Garg --------- Co-authored-by: Georgi Gerganov Co-authored-by: Gaurav Garg --- ggml/src/ggml.c | 4 ++- src/llama-sampler.cpp | 57 +++++++++++++++++++++++++++---------------- 2 files changed, 39 insertions(+), 22 deletions(-) diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index a7d1fe7d94be..59191c663eb0 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -7854,7 +7854,9 @@ void ggml_set_input(struct ggml_tensor * tensor) { } void ggml_set_output(struct ggml_tensor * tensor) { - tensor->flags |= GGML_TENSOR_FLAG_OUTPUT; + for (struct ggml_tensor * cur = tensor; cur != NULL; cur = cur->view_src) { + cur->flags |= GGML_TENSOR_FLAG_OUTPUT; + } } void ggml_set_param(struct ggml_tensor * tensor) { diff --git a/src/llama-sampler.cpp b/src/llama-sampler.cpp index 6520e4181e61..a9cb6bee5fd7 100644 --- a/src/llama-sampler.cpp +++ b/src/llama-sampler.cpp @@ -993,7 +993,9 @@ static void llama_sampler_greedy_backend_apply( GGML_UNUSED(gf); GGML_UNUSED(smpl); - struct ggml_tensor * curl = ggml_argmax(ctx, data->logits); + struct ggml_tensor * logits = ggml_reshape_1d(ctx, data->logits, ggml_nelements(data->logits)); + + struct ggml_tensor * curl = ggml_argmax(ctx, logits); ggml_set_name(curl, "greedy_argmax"); data->sampled = curl; @@ -1158,7 +1160,10 @@ static void llama_sampler_dist_backend_apply( ggml_set_name (sctx->inp_uniform, "uniform"); ggml_set_input(sctx->inp_uniform); - struct ggml_tensor * probs = ggml_soft_max(ctx, data->logits); + // flatten + struct ggml_tensor * logits = ggml_reshape_1d(ctx, data->logits, ggml_nelements(data->logits)); + + struct ggml_tensor * probs = ggml_soft_max(ctx, logits); ggml_set_name(probs, "dist_probs"); struct ggml_tensor * cumsum = ggml_cumsum(ctx, probs); @@ -1289,22 +1294,22 @@ static void llama_sampler_top_k_backend_apply( struct llama_sampler_data * data) { auto * sctx = (llama_sampler_top_k *) smpl->ctx; - struct ggml_tensor * top_k = ggml_top_k(ctx, data->logits, sctx->k); + struct ggml_tensor * logits = ggml_reshape_1d(ctx, data->logits, ggml_nelements(data->logits)); + + struct ggml_tensor * top_k = ggml_top_k(ctx, logits, sctx->k); ggml_set_name(top_k, "top_k"); if (data->candidates) { struct ggml_tensor * candidates_rows = ggml_reshape_2d(ctx, data->candidates, 1, data->candidates->ne[0]); data->candidates = ggml_get_rows(ctx, candidates_rows, top_k); - data->candidates = ggml_reshape_1d(ctx, data->candidates, sctx->k); ggml_set_name(data->candidates, "top_k_candidates"); } else { data->candidates = top_k; } - struct ggml_tensor * logits_rows = ggml_reshape_2d(ctx, data->logits, 1, data->logits->ne[0]); - struct ggml_tensor * top_k_rows = ggml_get_rows(ctx, logits_rows, top_k); - data->logits = ggml_reshape_1d(ctx, top_k_rows, sctx->k); - ggml_set_name(top_k_rows, "top_k_rows"); + struct ggml_tensor * logits_rows = ggml_reshape_2d(ctx, logits, 1, logits->ne[0]); + data->logits = ggml_get_rows(ctx, logits_rows, top_k); + ggml_set_name(data->logits, "top_k_rows"); GGML_UNUSED(gf); } @@ -1435,21 +1440,25 @@ static void llama_sampler_top_p_backend_apply( struct llama_sampler_data * data) { auto * sctx = (llama_sampler_top_p *) smpl->ctx; + // flatten + struct ggml_tensor * logits = ggml_reshape_1d(ctx, data->logits, ggml_nelements(data->logits)); + auto ggml_sort = [ctx](struct ggml_tensor * a, struct ggml_tensor * b) { GGML_ASSERT(ggml_nrows(a) == 1); struct ggml_tensor * a_reshaped = ggml_reshape_2d(ctx, a, 1, a->ne[0]); struct ggml_tensor * a_sorted = ggml_get_rows(ctx, a_reshaped, b); - return ggml_reshape_1d(ctx, a_sorted, a->ne[0]); + return a_sorted; }; // Get the sorted logits in descending order. - struct ggml_tensor * sorted_idx = ggml_argsort(ctx, data->logits, GGML_SORT_ORDER_DESC); + struct ggml_tensor * sorted_idx = ggml_argsort(ctx, logits, GGML_SORT_ORDER_DESC); ggml_set_name(sorted_idx, "top_p_sorted_idx"); // Do the sorting via reshape + get_rows - struct ggml_tensor * sorted_logits = ggml_sort(data->logits, sorted_idx); + struct ggml_tensor * sorted_logits = ggml_sort(logits, sorted_idx); ggml_set_name(sorted_logits, "top_p_sorted_logits"); + sorted_logits = ggml_reshape_1d(ctx, sorted_logits, ggml_nelements(sorted_logits)); struct ggml_tensor * softmax = ggml_soft_max(ctx, sorted_logits); ggml_set_name(softmax, "top_p_softmax"); @@ -1626,10 +1635,12 @@ static void llama_sampler_min_p_backend_apply( struct llama_sampler_data * data) { auto * sctx = (llama_sampler_min_p *) smpl->ctx; - struct ggml_tensor * max_idx = ggml_argmax(ctx, data->logits); + struct ggml_tensor * logits = ggml_reshape_1d(ctx, data->logits, ggml_nelements(data->logits)); + + struct ggml_tensor * max_idx = ggml_argmax(ctx, logits); ggml_set_name(max_idx, "max_idx"); - struct ggml_tensor * logits_rows = ggml_reshape_2d(ctx, data->logits, 1, data->logits->ne[0]); + struct ggml_tensor * logits_rows = ggml_reshape_2d(ctx, logits, 1, logits->ne[0]); ggml_set_name(logits_rows, "logits_rows"); struct ggml_tensor * max_logit = ggml_get_rows(ctx, logits_rows, max_idx); @@ -1640,7 +1651,7 @@ static void llama_sampler_min_p_backend_apply( ggml_set_name(threshold, "min_p_threshold"); // Subtract the threshold from logits. - struct ggml_tensor * sub = ggml_sub(ctx, data->logits, threshold); + struct ggml_tensor * sub = ggml_sub(ctx, logits, threshold); // Create a mask where logits below the threshold are 0 (discard), // and others are 1 (keep). @@ -1652,7 +1663,7 @@ static void llama_sampler_min_p_backend_apply( struct ggml_tensor * min_p_bias = ggml_log(ctx, mask); ggml_set_name(min_p_bias, "min_p_bias"); - data->logits = ggml_add(ctx, data->logits, min_p_bias); + data->logits = ggml_add(ctx, logits, min_p_bias); ggml_set_name(data->logits, "min_p_logits"); GGML_UNUSED(gf); @@ -1829,18 +1840,20 @@ static void llama_sampler_backend_temp_sampling( struct llama_sampler_data * data, float temp) { if (temp <= 0.0f) { + struct ggml_tensor * logits = ggml_reshape_1d(ctx, data->logits, ggml_nelements(data->logits)); + // Find the most probable token index. - struct ggml_tensor * max_idx = ggml_argmax(ctx, data->logits); + struct ggml_tensor * max_idx = ggml_argmax(ctx, logits); ggml_set_name(max_idx, "temp_max_idx"); if (data->candidates) { - struct ggml_tensor * candidates_rows = ggml_reshape_2d(ctx, data->candidates, 1, data->candidates->ne[0]); + struct ggml_tensor * candidates_rows = ggml_reshape_2d(ctx, data->candidates, 1, ggml_nelements(data->candidates)); data->candidates = ggml_get_rows(ctx, candidates_rows, max_idx); } else { data->candidates = max_idx; } - struct ggml_tensor * logits_rows = ggml_reshape_2d(ctx, data->logits, 1, data->logits->ne[0]); + struct ggml_tensor * logits_rows = ggml_reshape_2d(ctx, logits, 1, ggml_nelements(logits)); data->logits = ggml_get_rows(ctx, logits_rows, max_idx); return; @@ -2019,13 +2032,15 @@ static void llama_sampler_temp_ext_backend_apply( return; } + struct ggml_tensor * logits = ggml_reshape_1d(ctx, data->logits, ggml_nelements(data->logits)); + // Calculate min_temp, max_temp, and max_entropy. const float min_temp = std::max(0.0f, sctx->temp - sctx->delta); const float max_temp = sctx->temp + sctx->delta; - const float max_entropy = logf(data->logits->ne[0]); + const float max_entropy = logf(logits->ne[0]); // Calculate the probabilities. - struct ggml_tensor * probs = ggml_soft_max(ctx, data->logits); + struct ggml_tensor * probs = ggml_soft_max(ctx, logits); ggml_set_name(probs, "temp_ext_softmax_probs"); // Clamp probabilities to avoid log(0) which would give -inf @@ -2063,7 +2078,7 @@ static void llama_sampler_temp_ext_backend_apply( ggml_set_name(dyn_temp, "temp_ext_dyn_temp"); // Scale the logits by the dynamic temperature - struct ggml_tensor * scaled_logits = ggml_div(ctx, data->logits, dyn_temp); + struct ggml_tensor * scaled_logits = ggml_div(ctx, logits, dyn_temp); ggml_set_name(scaled_logits, "temp_ext_scaled_logits"); data->logits = scaled_logits; From ee3d1b54c1e47fce4cea9d63a60a1e45795eaf27 Mon Sep 17 00:00:00 2001 From: Xuan-Son Nguyen Date: Tue, 28 Jul 2026 16:35:20 +0200 Subject: [PATCH 039/190] server: abstract llama_memory calls to common_memory (#26221) --- common/common.cpp | 32 ++++++++++++++++-- common/common.h | 15 ++++++--- tools/server/server-context.cpp | 59 ++++++++++----------------------- 3 files changed, 58 insertions(+), 48 deletions(-) diff --git a/common/common.cpp b/common/common.cpp index 82dd780fd8b3..9d5d2834ac88 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -1518,23 +1518,49 @@ common_context_seq_rm_type common_context_can_seq_rm(llama_context * ctx) { return res; } -void common_context_seq_rm(llama_context * ctx, llama_seq_id seq_id, llama_pos p0, llama_pos p1) { +static void common_context_seq_rm(llama_context * ctx, llama_seq_id seq_id, llama_pos p0, llama_pos p1) { auto * mem = llama_get_memory(ctx); if (!llama_memory_seq_rm(mem, seq_id, p0, p1)) { GGML_ABORT("%s", string_format("failed to remove sequence %d with p0=%d, p1=%d\n", seq_id, p0, p1).c_str()); } } -void common_context_seq_cp(llama_context * ctx, llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1) { +static void common_context_seq_cp(llama_context * ctx, llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1) { auto * mem = llama_get_memory(ctx); llama_memory_seq_cp(mem, seq_id_src, seq_id_dst, p0, p1); } -void common_context_seq_add(llama_context * ctx, llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos delta) { +static void common_context_seq_add(llama_context * ctx, llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos delta) { auto * mem = llama_get_memory(ctx); llama_memory_seq_add(mem, seq_id, p0, p1, delta); } +void common_memory::init(llama_context * ctx_tgt, llama_context * ctx_dft) { + this->ctx_tgt = ctx_tgt; + this->ctx_dft = ctx_dft; +} + +void common_memory::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1) const { + common_context_seq_rm(ctx_tgt, seq_id, p0, p1); + if (ctx_dft) { + common_context_seq_rm(ctx_dft, seq_id, p0, p1); + } +} + +void common_memory::seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1) const { + common_context_seq_cp(ctx_tgt, seq_id_src, seq_id_dst, p0, p1); + if (ctx_dft) { + common_context_seq_cp(ctx_dft, seq_id_src, seq_id_dst, p0, p1); + } +} + +void common_memory::seq_add(llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos delta) const { + common_context_seq_add(ctx_tgt, seq_id, p0, p1, delta); + if (ctx_dft) { + common_context_seq_add(ctx_dft, seq_id, p0, p1, delta); + } +} + void common_set_adapter_lora(struct llama_context * ctx, std::vector & lora) { std::vector loras; std::vector scales; diff --git a/common/common.h b/common/common.h index 78b1f416e085..e7c55ae925fc 100644 --- a/common/common.h +++ b/common/common.h @@ -949,10 +949,17 @@ enum common_context_seq_rm_type { // note: clears the memory of the context common_context_seq_rm_type common_context_can_seq_rm(llama_context * ctx); -// aborts execution on failure -void common_context_seq_rm (llama_context * ctx, llama_seq_id seq_id, llama_pos p0, llama_pos p1); -void common_context_seq_add(llama_context * ctx, llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos delta); -void common_context_seq_cp (llama_context * ctx, llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1); +struct common_memory { + llama_context * ctx_tgt = nullptr; + llama_context * ctx_dft = nullptr; + + void init(llama_context * ctx_tgt, llama_context * ctx_dft = nullptr); + + // aborts execution on failure + void seq_rm (llama_seq_id seq_id, llama_pos p0, llama_pos p1) const; + void seq_add(llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos delta) const; + void seq_cp (llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1) const; +}; // // Batch utils diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 744593c760a3..dba15d426fd4 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -164,6 +164,8 @@ struct server_slot { llama_context * ctx_tgt = nullptr; llama_context * ctx_dft = nullptr; + common_memory mem; + // multimodal mtmd_context * mctx = nullptr; mtmd::batch_ptr mbatch = nullptr; @@ -253,10 +255,7 @@ struct server_slot { void prompt_clear() { SLT_TRC(*this, "clearing prompt with %zu tokens\n", prompt.tokens.size()); - common_context_seq_rm(ctx_tgt, id, -1, -1); - if (ctx_dft) { - common_context_seq_rm(ctx_dft, id, -1, -1); - } + mem.seq_rm(id, -1, -1); prompt.clear(); } @@ -668,13 +667,8 @@ struct server_slot { void copy_state_to(server_slot & other) const { GGML_ASSERT(state == SLOT_STATE_DONE_PROMPT); - common_context_seq_rm(ctx_tgt, other.id, -1, -1); - common_context_seq_cp(ctx_tgt, id, other.id, -1, -1); - - if (ctx_dft) { - common_context_seq_rm(ctx_dft, other.id, -1, -1); - common_context_seq_cp(ctx_dft, id, other.id, -1, -1); - } + mem.seq_rm(other.id, -1, -1); + mem.seq_cp(id, other.id, -1, -1); other.n_decoded = n_decoded; other.n_remaining = n_remaining; @@ -1302,6 +1296,7 @@ struct server_context_impl { slot.id = i; slot.ctx_tgt = ctx_tgt; slot.ctx_dft = ctx_dft; + slot.mem.init(ctx_tgt, ctx_dft); slot.spec = spec.get(); slot.n_ctx = n_ctx_slot; @@ -2881,13 +2876,8 @@ struct server_context_impl { SLT_WRN(slot, "slot context shift, n_keep = %d, n_left = %d, n_discard = %d\n", n_keep, n_left, n_discard); - common_context_seq_rm (ctx_tgt, slot.id, n_keep , n_keep + n_discard); - common_context_seq_add(ctx_tgt, slot.id, n_keep + n_discard, slot.prompt.n_tokens(), -n_discard); - - if (ctx_dft) { - common_context_seq_rm (ctx_dft, slot.id, n_keep , n_keep + n_discard); - common_context_seq_add(ctx_dft, slot.id, n_keep + n_discard, slot.prompt.tokens.pos_next(), -n_discard); - } + slot.mem.seq_rm (slot.id, n_keep , n_keep + n_discard); + slot.mem.seq_add(slot.id, n_keep + n_discard, slot.prompt.tokens.pos_next(), -n_discard); // add generated tokens to cache // ref: https://github.com/ggml-org/llama.cpp/pull/16818#discussion_r2473269481 @@ -2998,7 +2988,9 @@ struct server_context_impl { ckpt.load_dft(ctx_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); } - common_context_seq_rm(ctx_dft, slot.id, ckpt.pos_max + 1, -1); + if (!llama_memory_seq_rm(llama_get_memory(ctx_dft), slot.id, ckpt.pos_max + 1, -1)) { + GGML_ABORT("failed to remove sequence %d\n", slot.id); + } } if (!draft.empty()) { @@ -3201,13 +3193,8 @@ struct server_context_impl { const int64_t kv_shift = (int64_t) head_p - (int64_t) head_c; - common_context_seq_rm (ctx_tgt, slot.id, head_p, head_c); - common_context_seq_add(ctx_tgt, slot.id, head_c, head_c + n_match, kv_shift); - - if (ctx_dft) { - common_context_seq_rm (ctx_dft, slot.id, head_p, head_c); - common_context_seq_add(ctx_dft, slot.id, head_c, head_c + n_match, kv_shift); - } + slot.mem.seq_rm (slot.id, head_p, head_c); + slot.mem.seq_add(slot.id, head_c, head_c + n_match, kv_shift); for (size_t i = 0; i < n_match; i++) { slot.prompt.tokens.set_token(head_p + i, slot.prompt.tokens[head_c + i]); @@ -3379,10 +3366,7 @@ struct server_context_impl { SLT_TRC(slot, "cached n_tokens = %d, memory_seq_rm [%d, end)\n", slot.prompt.n_tokens(), p0); - common_context_seq_rm(ctx_tgt, slot.id, p0, -1); - if (ctx_dft) { - common_context_seq_rm(ctx_dft, slot.id, p0, -1); - } + slot.mem.seq_rm(slot.id, p0, -1); // If using an alora, there may be uncached tokens that come // before the invocation sequence. When this happens, the @@ -3837,18 +3821,14 @@ struct server_context_impl { SLT_DBG(slot, "restoring speculative checkpoint (pos_min = %d, pos_max = %d, size = %zu)\n", ckpt.pos_min, ckpt.pos_max, ckpt.size()); - { - ckpt.load_tgt(slot.ctx_tgt, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); - - common_context_seq_rm(slot.ctx_tgt, slot.id, ckpt.pos_max + 1, -1); - } + ckpt.load_tgt(slot.ctx_tgt, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); if (slot.ctx_dft) { ckpt.load_dft(slot.ctx_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); - - common_context_seq_rm(slot.ctx_dft, slot.id, ckpt.pos_max + 1, -1); } + slot.mem.seq_rm(slot.id, ckpt.pos_max + 1, -1); + slot.prompt.tokens.keep_first(ckpt.n_tokens); slot.smpl = std::move(smpl_save); @@ -3889,10 +3869,7 @@ struct server_context_impl { slot.sampled = ids.back(); // last accepted token SLT_DBG(slot, "add accepted tokens: sampled=%d, ids.size=%zu, n_draft=%zu\n", slot.sampled, ids.size(), n_draft); - common_context_seq_rm(slot.ctx_tgt, slot.id, slot.prompt.tokens.pos_next(), -1); - if (slot.ctx_dft) { - common_context_seq_rm(slot.ctx_dft, slot.id, slot.prompt.tokens.pos_next(), -1); - } + slot.mem.seq_rm(slot.id, slot.prompt.tokens.pos_next(), -1); for (size_t i = 0; i < ids.size(); ++i) { completion_token_output result; From ad77bd31a609f3117ebc1754ac278ff38975725c Mon Sep 17 00:00:00 2001 From: Julien Jerphanion Date: Tue, 28 Jul 2026 16:51:20 +0200 Subject: [PATCH 040/190] docs: Adapt conda-forge package name (#26229) Co-authored-by: dev-tinker --- docs/install.md | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/docs/install.md b/docs/install.md index 7198e61bf35b..b36b0be26736 100644 --- a/docs/install.md +++ b/docs/install.md @@ -16,22 +16,22 @@ conda-forge provides builds for: - Apple Metal (macOS) ```sh -conda install -c conda-forge llama-cpp +conda install -c conda-forge llama.cpp ``` ```sh -mamba install -c conda-forge llama-cpp +mamba install -c conda-forge llama.cpp ``` ```sh # Project-local installation -pixi add llama-cpp +pixi add llama.cpp # Global installation -pixi global install llama-cpp +pixi global install llama.cpp ``` -This distribution is managed on [`conda-forge/llama-cpp-feedstock`](https://github.com/conda-forge/llama.cpp-feedstock/). +This distribution is managed on [`conda-forge/llama.cpp-feedstock`](https://github.com/conda-forge/llama.cpp-feedstock/). Shall you have any problems, please open an issue on [its issue tracker](https://github.com/conda-forge/llama.cpp-feedstock/issues). From 6e2bc65fb2a822a8082a7515858a7c180115dd42 Mon Sep 17 00:00:00 2001 From: Aleksander Grygier Date: Tue, 28 Jul 2026 17:13:25 +0200 Subject: [PATCH 041/190] ui: rendering performance follow-up (#26097) --- .../ChatMessageToolCallBlockDefault.svelte | 8 +- .../ChatMessageToolCallBlockEditFile.svelte | 1 - ...atMessageToolCallBlockSearchResults.svelte | 2 +- .../ui/src/lib/constants/latex-protection.ts | 82 ++++++++++ tools/ui/src/lib/utils/agentic.ts | 85 ++++++++-- tools/ui/src/lib/utils/code.ts | 28 +++- tools/ui/src/lib/utils/latex-protection.ts | 152 +++++++++++------- .../src/lib/utils/parse-partial-json-args.ts | 128 +++++++++------ tools/ui/src/lib/utils/search-results.ts | 40 ++++- .../tests/unit/parse-toolcalls-memo.test.ts | 146 +++++++++++++++++ 10 files changed, 538 insertions(+), 134 deletions(-) create mode 100644 tools/ui/tests/unit/parse-toolcalls-memo.test.ts diff --git a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/ChatMessageToolCallBlockDefault.svelte b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/ChatMessageToolCallBlockDefault.svelte index acf2de12ac64..92652a0a86ab 100644 --- a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/ChatMessageToolCallBlockDefault.svelte +++ b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/ChatMessageToolCallBlockDefault.svelte @@ -11,8 +11,7 @@ classifyToolResult, formatJsonPretty, parseToolResultWithImages, - type AgenticSection, - type ToolResultLine + type AgenticSection } from '$lib/utils'; import { getBuiltinToolUi } from '$lib/constants/built-in-tools'; import type { DatabaseMessageExtra } from '$lib/types'; @@ -29,11 +28,10 @@ let { section, open, isStreaming, attachments, onToggle }: Props = $props(); const title = $derived(getBuiltinToolUi(section.toolName)?.label ?? section.toolName ?? ''); - - const parsedLines: ToolResultLine[] = $derived( + const outputKind = $derived(classifyToolResult(section.toolResult)); + const parsedLines = $derived( section.toolResult ? parseToolResultWithImages(section.toolResult, attachments) : [] ); - const outputKind = $derived(classifyToolResult(section.toolResult)); diff --git a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/ChatMessageToolCallBlockEditFile.svelte b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/ChatMessageToolCallBlockEditFile.svelte index 6f30060f54c4..b990c3898b23 100644 --- a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/ChatMessageToolCallBlockEditFile.svelte +++ b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/ChatMessageToolCallBlockEditFile.svelte @@ -15,7 +15,6 @@ let { section, open, isStreaming, onToggle }: Props = $props(); const editFileMeta = $derived(parseEditFileMeta(section)); - const editDiffs = $derived( (editFileMeta?.edits ?? []).map((edit) => computeLineDiff(edit.oldText, edit.newText)) ); diff --git a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/ChatMessageToolCallBlockSearchResults.svelte b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/ChatMessageToolCallBlockSearchResults.svelte index e4b4adf151d3..60862dd06353 100644 --- a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/ChatMessageToolCallBlockSearchResults.svelte +++ b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/ChatMessageToolCallBlockSearchResults.svelte @@ -27,7 +27,7 @@ const isStreamingCall = $derived(section.type === AgenticSectionType.TOOL_CALL_STREAMING); const showSpinner = $derived(isPending || (isStreamingCall && isStreaming)); - const results: SearchResult[] = $derived(extractSearchResults(section.toolResult)); + const results = $derived(extractSearchResults(section.toolResult)); const query = $derived(extractSearchQuery(section.toolArgs)); // Same icon-resolution chain as ChatMessageToolCallBlockDefault so diff --git a/tools/ui/src/lib/constants/latex-protection.ts b/tools/ui/src/lib/constants/latex-protection.ts index da27c008373b..c42aec41acf3 100644 --- a/tools/ui/src/lib/constants/latex-protection.ts +++ b/tools/ui/src/lib/constants/latex-protection.ts @@ -28,6 +28,18 @@ export const LATEX_MATH_AND_CODE_PATTERN = /** Regex to capture the content of a $$...\\\\...$$ block (display-formula with line-break) */ export const LATEX_LINEBREAK_REGEXP = /\$\$([\s\S]*?\\\\[\s\S]*?)\$\$/; +/** + * Matches the unescaped `\[...\]` display-math delimiter and surrounding + * context so callers can insert line-breaks around the placeholder or convert + * to inline when the formula has a non-empty trailing context (e.g. a table + * cell that opens with `\[` and closes with content after `\]`). + * + * group 1: prefix before `\[` + * group 2: formula body + * group 3: trailing context after `\]` + */ +export const LATEX_DISPLAY_BLOCK_REGEXP = /([\S].*?)\\\[([\s\S]*?)\\\](.*)/g; + /** * Cheap gate for `preprocessLaTeX`. Every transformation it performs is triggered * by a `$` (inline/display math, currency escaping) or a backslash escape @@ -36,6 +48,76 @@ export const LATEX_LINEBREAK_REGEXP = /\$\$([\s\S]*?\\\\[\s\S]*?)\$\$/; */ export const LATEX_TRIGGER_REGEXP = /[$\\]/; +/** Inline LaTeX math delimiter (the dollar sign). */ +export const LATEX_INLINE_DELIMITER = '$'; + +/** Display LaTeX math delimiter (paired dollar signs). */ +export const LATEX_DISPLAY_DELIMITER = '$$'; + +/** Matches a single non-whitespace character. */ +export const LATEX_NON_WHITESPACE_REGEXP = /\S/; + +/** Matches a character that may appear adjacent to `$`, indicating a non-TeX + * context such as an identifier (`var$`, `$var`), currency ($5), or code. */ +export const LATEX_NEIGHBOR_CHAR_REGEXP = /[A-Za-z0-9_$-]/; + +/** Matches a single digit (used to detect currency-like `$5`). */ +export const LATEX_DIGIT_REGEXP = /[0-9]/; + +/** Matches the leading blockquote prefix (`> ` or `>`) on a markdown line. */ +export const LATEX_BLOCKQUOTE_PREFIX_REGEXP = /^(>\s*)/; + +/** Matches the placeholder inserted by the protect/restore pipeline for a + * protected LaTeX expression. Group 1 is the index into `latexExpressions`. */ +export const LATEX_PLACEHOLDER_REGEXP = /<>/g; + +/** Matches the placeholder inserted by the protect/restore pipeline for a + * protected code block. Group 1 is the index into `codeBlocks`. */ +export const CODE_BLOCK_PLACEHOLDER_REGEXP = /<>/g; + +/** Matches a `$` immediately followed by a digit, which is treated as a + * currency amount (e.g. `$5`) and escaped to `\$5` so it isn't parsed as math. */ +export const LATEX_CURRENCY_DOLLAR_REGEXP = /\$(?=\d)/g; + +/** Captures remaining `$$...$$`, `\[...\]`, `\(...\)` (only unescaped via + * `(?(); + for (const tm of toolMessages) { + if (tm.toolCallId && !toolMsgById.has(tm.toolCallId)) { + toolMsgById.set(tm.toolCallId, tm); + } + } + for (const tc of toolCalls) { - const resultMsg = toolMessages.find((m) => m.toolCallId === tc.id); + const resultMsg = tc.id ? toolMsgById.get(tc.id) : undefined; // Only show as pending/loading if we're actively streaming; otherwise it's just a tool call without result const type = resultMsg ? AgenticSectionType.TOOL_CALL @@ -112,9 +121,10 @@ function deriveSingleTurnSections( } // 4. Streaming tool calls (not yet persisted - currently being received) + const persistedIds = new Set(toolCalls.map((t) => t.id).filter(Boolean)); for (const tc of streamingToolCalls) { // Skip if already in persisted tool calls - if (tc.id && toolCalls.find((t) => t.id === tc.id)) continue; + if (tc.id && persistedIds.has(tc.id)) continue; sections.push({ type: AgenticSectionType.TOOL_CALL_STREAMING, content: '', @@ -281,15 +291,31 @@ export function splitSearchSummaryList( return { lines }; } +/** Bounded cache for parseToolResultWithImages results. */ +const TOOL_RESULT_LINES_CACHE_MAX_SIZE = 32; +const toolResultLinesCache = new Map(); + /** * Parse tool result text into lines, matching image attachments by name. + * Memoized: called per render during streaming on unchanged tool result + * strings with unchanged extras. */ export function parseToolResultWithImages( toolResult: string, extras?: DatabaseMessageExtra[] ): ToolResultLine[] { + // Cache key includes image attachment names so we recompute when + // attachments change, even if the count stays the same. + const imageNames = (extras ?? []) + .filter((e): e is DatabaseMessageExtraImageFile => e.type === AttachmentType.IMAGE) + .map((e) => e.name) + .join(NEWLINE); + const cacheKey = `${imageNames}:${toolResult}`; + const cached = toolResultLinesCache.get(cacheKey); + if (cached !== undefined) return cached; + const lines = toolResult.split(NEWLINE); - return lines.map((line) => { + const result = lines.map((line) => { const match = line.match(ATTACHMENT_SAVED_REGEX); if (!match || !extras) return { text: line }; @@ -301,8 +327,19 @@ export function parseToolResultWithImages( return { text: line, image }; }); + + if (toolResultLinesCache.size >= TOOL_RESULT_LINES_CACHE_MAX_SIZE) { + toolResultLinesCache.delete(toolResultLinesCache.keys().next().value!); + } + toolResultLinesCache.set(cacheKey, result); + + return result; } +/** Bounded cache for classifyToolResult results. */ +const CLASSIFY_CACHE_MAX_SIZE = 32; +const classifyCache = new Map(); + /** * Pick a renderer tier for a tool's result content. * @@ -312,25 +349,39 @@ export function parseToolResultWithImages( * through MarkdownContent for proper formatting. * text - everything else, rendered as plain text lines (with image * attachment resolution as a side effect). + * Memoized: called per render during streaming on unchanged content. */ export function classifyToolResult(content: string | undefined): ToolResultKind { if (!content) return ToolResultKind.TEXT; + + const cached = classifyCache.get(content); + if (cached !== undefined) return cached; + const trimmed = content.trim(); if (!trimmed) return ToolResultKind.TEXT; + let result: ToolResultKind = ToolResultKind.TEXT; + // Strongest signal: JSON object/array round-trips through JSON.parse. if (TOOL_RESULT_JSON_OPEN_REGEX.test(trimmed)) { try { JSON.parse(trimmed); - return ToolResultKind.JSON; + result = ToolResultKind.JSON; } catch (error) { console.error('[agentic] tool result looked like JSON but failed to parse:', error); } } - if (looksLikeMarkdown(trimmed)) return ToolResultKind.MARKDOWN; + if (result === ToolResultKind.TEXT && looksLikeMarkdown(trimmed)) { + result = ToolResultKind.MARKDOWN; + } + + if (classifyCache.size >= CLASSIFY_CACHE_MAX_SIZE) { + classifyCache.delete(classifyCache.keys().next().value!); + } + classifyCache.set(content, result); - return ToolResultKind.TEXT; + return result; } /** @@ -370,19 +421,35 @@ function looksLikeMarkdown(content: string): boolean { return false; } +/** Bounded cache for parsed tool-call JSON blobs. */ +const TOOL_CALLS_CACHE_MAX_SIZE = 64; +const toolCallsParseCache = new Map(); + /** * Safely parse the toolCalls JSON string from a DatabaseMessage. + * Memoized: the same JSON string is re-parsed on every render during + * streaming, which is wasted CPU since tool calls don't change mid-stream. */ function parseToolCalls(toolCallsJson?: string): ApiChatCompletionToolCall[] { if (!toolCallsJson) return []; + const cached = toolCallsParseCache.get(toolCallsJson); + if (cached) return cached; + + let result: ApiChatCompletionToolCall[]; try { const parsed = JSON.parse(toolCallsJson); - - return Array.isArray(parsed) ? parsed : []; + result = Array.isArray(parsed) ? parsed : []; } catch { - return []; + result = []; + } + + if (toolCallsParseCache.size >= TOOL_CALLS_CACHE_MAX_SIZE) { + toolCallsParseCache.delete(toolCallsParseCache.keys().next().value!); } + toolCallsParseCache.set(toolCallsJson, result); + + return result; } /** diff --git a/tools/ui/src/lib/utils/code.ts b/tools/ui/src/lib/utils/code.ts index 44b9b48415c4..35f3877f4aaf 100644 --- a/tools/ui/src/lib/utils/code.ts +++ b/tools/ui/src/lib/utils/code.ts @@ -34,6 +34,10 @@ function escapeCode(code: string): string { return code.replace(AMPERSAND_REGEX, '&').replace(LT_REGEX, '<').replace(GT_REGEX, '>'); } +/** Bounded cache for highlightCode results. */ +const HIGHLIGHT_CACHE_MAX_SIZE = 64; +const highlightCache = new Map(); + /** * Highlights code using highlight.js * @param code - The code to highlight @@ -47,23 +51,37 @@ function escapeCode(code: string): string { export function highlightCode(code: string, language: string, autoDetect = true): string { if (!code) return ''; + // Cache key includes language and autoDetect flag since results differ. + // During streaming, the same code string may be highlighted repeatedly + // (e.g., when text after a code block changes but the code itself doesn't). + const cacheKey = `${language}:${autoDetect}:${code}`; + const cached = highlightCache.get(cacheKey); + if (cached) return cached; + const trimmed = trimCodePadding(code); + let result: string; try { const lang = language.toLowerCase(); const isSupported = hljs.getLanguage(lang); if (isSupported) { - return hljs.highlight(trimmed, { language: lang }).value; + result = hljs.highlight(trimmed, { language: lang }).value; } else if (autoDetect) { - return hljs.highlightAuto(trimmed).value; + result = hljs.highlightAuto(trimmed).value; } else { - return escapeCode(trimmed); + result = escapeCode(trimmed); } } catch { - // Fallback to escaped plain text - return escapeCode(trimmed); + result = escapeCode(trimmed); } + + if (highlightCache.size >= HIGHLIGHT_CACHE_MAX_SIZE) { + highlightCache.delete(highlightCache.keys().next().value!); + } + highlightCache.set(cacheKey, result); + + return result; } export { trimCodePadding }; diff --git a/tools/ui/src/lib/utils/latex-protection.ts b/tools/ui/src/lib/utils/latex-protection.ts index 573eb9297807..bbeed825006d 100644 --- a/tools/ui/src/lib/utils/latex-protection.ts +++ b/tools/ui/src/lib/utils/latex-protection.ts @@ -1,9 +1,31 @@ import { + CODE_BLOCK_PLACEHOLDER_REGEXP, CODE_BLOCK_REGEXP, + LATEX_BACKSLASH, + LATEX_BLOCKQUOTE_PREFIX_REGEXP, + LATEX_CURRENCY_DOLLAR_REGEXP, + LATEX_CURRENCY_ESCAPE, + LATEX_DIGIT_REGEXP, + LATEX_DISPLAY_BLOCK_REGEXP, + LATEX_DISPLAY_CLOSE, + LATEX_DISPLAY_CONVERT_REGEXP, + LATEX_DISPLAY_DELIMITER, + LATEX_DISPLAY_OPEN, + LATEX_INLINE_CLOSE, + LATEX_INLINE_CONVERT_REGEXP, + LATEX_INLINE_DELIMITER, + LATEX_INLINE_OPEN, LATEX_MATH_AND_CODE_PATTERN, + LATEX_MHCHEM_CE, + LATEX_MHCHEM_PU, LATEX_LINEBREAK_REGEXP, + LATEX_NEIGHBOR_CHAR_REGEXP, + LATEX_NON_WHITESPACE_REGEXP, + LATEX_PLACEHOLDER_REGEXP, + LATEX_PROTECT_REGEXP, LATEX_TRIGGER_REGEXP, - MHCHEM_PATTERN_MAP + MHCHEM_PATTERN_MAP, + NEWLINE } from '$lib/constants'; /** @@ -20,13 +42,13 @@ import { * @returns The processed string with LaTeX replaced by placeholders. */ export function maskInlineLaTeX(content: string, latexExpressions: string[]): string { - if (!content.includes('$')) { + if (!content.includes(LATEX_INLINE_DELIMITER)) { return content; } return content - .split('\n') + .split(NEWLINE) .map((line) => { - if (line.indexOf('$') == -1) { + if (line.indexOf(LATEX_INLINE_DELIMITER) == -1) { return line; } @@ -34,7 +56,7 @@ export function maskInlineLaTeX(content: string, latexExpressions: string[]): st let currentPosition = 0; while (currentPosition < line.length) { - const openDollarIndex = line.indexOf('$', currentPosition); + const openDollarIndex = line.indexOf(LATEX_INLINE_DELIMITER, currentPosition); if (openDollarIndex == -1) { processedLine += line.slice(currentPosition); @@ -42,7 +64,7 @@ export function maskInlineLaTeX(content: string, latexExpressions: string[]): st } // Is there a next $-sign? - const closeDollarIndex = line.indexOf('$', openDollarIndex + 1); + const closeDollarIndex = line.indexOf(LATEX_INLINE_DELIMITER, openDollarIndex + 1); if (closeDollarIndex == -1) { processedLine += line.slice(currentPosition); @@ -62,14 +84,14 @@ export function maskInlineLaTeX(content: string, latexExpressions: string[]): st shouldSkipAsNonLatex = true; } - if (/[A-Za-z0-9_$-]/.test(charBeforeOpen)) { + if (LATEX_NEIGHBOR_CHAR_REGEXP.test(charBeforeOpen)) { // Character, digit, $, _ or - before first '$', no TeX. shouldSkipAsNonLatex = true; } if ( - /[0-9]/.test(charAfterOpen) && - (/[A-Za-z0-9_$-]/.test(charAfterClose) || ' ' == charBeforeClose) + LATEX_DIGIT_REGEXP.test(charAfterOpen) && + (LATEX_NEIGHBOR_CHAR_REGEXP.test(charAfterClose) || ' ' == charBeforeClose) ) { // First $ seems to belong to an amount. shouldSkipAsNonLatex = true; @@ -92,7 +114,7 @@ export function maskInlineLaTeX(content: string, latexExpressions: string[]): st return processedLine; }) - .join('\n'); + .join(NEWLINE); } function escapeBrackets(text: string): string { @@ -107,9 +129,9 @@ function escapeBrackets(text: string): string { if (codeBlock != null) { return codeBlock; } else if (squareBracket != null) { - return `$$${squareBracket}$$`; + return `${LATEX_DISPLAY_DELIMITER}${squareBracket}${LATEX_DISPLAY_DELIMITER}`; } else if (roundBracket != null) { - return `$${roundBracket}$`; + return `${LATEX_INLINE_DELIMITER}${roundBracket}${LATEX_INLINE_DELIMITER}`; } return match; @@ -145,32 +167,49 @@ const doEscapeMhchem = false; * preprocessLaTeX("Price: $10. The equation is \\(x^2\\).") * // → "Price: $10. The equation is $x^2$." */ +/** Bounded cache for preprocessLaTeX results. */ +const LATEX_CACHE_MAX_SIZE = 64; +const latexCache = new Map(); + export function preprocessLaTeX(content: string): string { // See also: // https://github.com/danny-avila/LibreChat/blob/main/client/src/utils/latex.ts + // Memoize on the input string. During streaming the prefix before an + // incomplete code block stays the same across multiple tokens, so the + // full protect/restore pipeline would re-run unnecessarily. + const cached = latexCache.get(content); + if (cached !== undefined) return cached; + + // Save original before the function mutates `content` through steps 0-8 + const originalContent = content; + // Every step below keys off a `$` or a backslash escape (\[ \] \( \) \ce{ \pu{). // With neither present the protect/restore passes round-trip the input // unchanged, so skip them: the step 2 scan is O(n^2) in line length and costs // ~90ms on a 26KB single-line message that contains no math at all. This // matters during streaming, where the whole message is reprocessed per frame. if (!LATEX_TRIGGER_REGEXP.test(content)) { + if (latexCache.size >= LATEX_CACHE_MAX_SIZE) { + latexCache.delete(latexCache.keys().next().value!); + } + latexCache.set(originalContent, content); return content; } // Step 0: Temporarily remove blockquote markers (>) to process LaTeX correctly // Store the structure so we can restore it later const blockquoteMarkers: Map = new Map(); - const lines = content.split('\n'); + const lines = content.split(NEWLINE); const processedLines = lines.map((line, index) => { - const match = line.match(/^(>\s*)/); + const match = line.match(LATEX_BLOCKQUOTE_PREFIX_REGEXP); if (match) { blockquoteMarkers.set(index, match[1]); return line.slice(match[1].length); } return line; }); - content = processedLines.join('\n'); + content = processedLines.join(NEWLINE); // Step 1: Protect code blocks const codeBlocks: string[] = []; @@ -187,58 +226,52 @@ export function preprocessLaTeX(content: string): string { // Match \S...\[...\] and protect them and insert a line-break. // Guarded: with no `\[` present this pattern still probes every start offset, // expanding `.*?` to the end of each line before failing - O(n^2) for nothing. - if (content.includes('\\[')) { - content = content.replace( - /([\S].*?)\\\[([\s\S]*?)\\\](.*)/g, - (match, group1, group2, group3) => { - // Check if there are characters following the formula (display-formula in a table-cell?) - if (group1.endsWith('\\')) { - return match; // Backslash before \[, do nothing. - } - const hasSuffix = /\S/.test(group3); - let optBreak; - - if (hasSuffix) { - latexExpressions.push(`\\(${group2.trim()}\\)`); // Convert into inline. - optBreak = ''; - } else { - latexExpressions.push(`\\[${group2}\\]`); - optBreak = '\n'; - } - - return `${group1}${optBreak}<>${optBreak}${group3}`; + if (content.includes(LATEX_DISPLAY_OPEN)) { + content = content.replace(LATEX_DISPLAY_BLOCK_REGEXP, (match, group1, group2, group3) => { + // Check if there are characters following the formula (display-formula in a table-cell?) + if (group1.endsWith(LATEX_BACKSLASH)) { + return match; // Backslash before \[, do nothing. } - ); + const hasSuffix = LATEX_NON_WHITESPACE_REGEXP.test(group3); + let optBreak; + + if (hasSuffix) { + latexExpressions.push(`${LATEX_INLINE_OPEN}${group2.trim()}${LATEX_INLINE_CLOSE}`); // Convert into inline. + optBreak = ''; + } else { + latexExpressions.push(`${LATEX_DISPLAY_OPEN}${group2}${LATEX_DISPLAY_CLOSE}`); + optBreak = NEWLINE; + } + + return `${group1}${optBreak}<>${optBreak}${group3}`; + }); } // Match \(...\), \[...\], $$...$$ and protect them - content = content.replace( - /(\$\$[\s\S]*?\$\$|(? { - latexExpressions.push(match); + content = content.replace(LATEX_PROTECT_REGEXP, (match) => { + latexExpressions.push(match); - return `<>`; - } - ); + return `<>`; + }); // Protect inline $...$ but NOT if it looks like money (e.g., $10, $3.99) content = maskInlineLaTeX(content, latexExpressions); // Step 3: Escape standalone $ before digits (currency like $5 → \$5) // (Now that inline math is protected, this will only escape dollars not already protected) - content = content.replace(/\$(?=\d)/g, '\\$'); + content = content.replace(LATEX_CURRENCY_DOLLAR_REGEXP, LATEX_CURRENCY_ESCAPE); // Step 4: Restore protected LaTeX expressions (they are valid) - content = content.replace(/<>/g, (_, index) => { + content = content.replace(LATEX_PLACEHOLDER_REGEXP, (_, index) => { let expr = latexExpressions[parseInt(index)]; const match = expr.match(LATEX_LINEBREAK_REGEXP); if (match) { // Katex: The $$-delimiters should be in their own line // if there are \\-line-breaks. const formula = match[1]; - const prefix = formula.startsWith('\n') ? '' : '\n'; - const suffix = formula.endsWith('\n') ? '' : '\n'; - expr = '$$' + prefix + formula + suffix + '$$'; + const prefix = formula.startsWith(NEWLINE) ? '' : NEWLINE; + const suffix = formula.endsWith(NEWLINE) ? '' : NEWLINE; + expr = LATEX_DISPLAY_DELIMITER + prefix + formula + suffix + LATEX_DISPLAY_DELIMITER; } return expr; }); @@ -247,7 +280,7 @@ export function preprocessLaTeX(content: string): string { // This must happen BEFORE restoring code blocks to avoid affecting code content content = escapeBrackets(content); - if (doEscapeMhchem && (content.includes('\\ce{') || content.includes('\\pu{'))) { + if (doEscapeMhchem && (content.includes(LATEX_MHCHEM_CE) || content.includes(LATEX_MHCHEM_PU))) { content = escapeMhchem(content); } @@ -257,31 +290,38 @@ export function preprocessLaTeX(content: string): string { // Using the look‑behind pattern `(? { + return `${LATEX_INLINE_DELIMITER}${formula}${LATEX_INLINE_DELIMITER}`; + }) // inline .replace( // Using the look‑behind pattern `(? { - return `$$${content}$$`; + LATEX_DISPLAY_CONVERT_REGEXP, // display, see also PR #16599 + (_, formula: string) => { + return `${LATEX_DISPLAY_DELIMITER}${formula}${LATEX_DISPLAY_DELIMITER}`; } ); // Step 7: Restore code blocks // This happens AFTER all LaTeX conversions to preserve code content - content = content.replace(/<>/g, (_, index) => { + content = content.replace(CODE_BLOCK_PLACEHOLDER_REGEXP, (_, index) => { return codeBlocks[parseInt(index)]; }); // Step 8: Restore blockquote markers if (blockquoteMarkers.size > 0) { - const finalLines = content.split('\n'); + const finalLines = content.split(NEWLINE); const restoredLines = finalLines.map((line, index) => { const marker = blockquoteMarkers.get(index); return marker ? marker + line : line; }); - content = restoredLines.join('\n'); + content = restoredLines.join(NEWLINE); + } + + if (latexCache.size >= LATEX_CACHE_MAX_SIZE) { + latexCache.delete(latexCache.keys().next().value!); } + latexCache.set(originalContent, content); return content; } diff --git a/tools/ui/src/lib/utils/parse-partial-json-args.ts b/tools/ui/src/lib/utils/parse-partial-json-args.ts index 09c11228906e..58439bd6e301 100644 --- a/tools/ui/src/lib/utils/parse-partial-json-args.ts +++ b/tools/ui/src/lib/utils/parse-partial-json-args.ts @@ -14,70 +14,94 @@ const JSON_ARRAY_CLOSE = ']'; // comma when the model cut off mid-key. const TRAILING_JSON_PUNCTUATION_REGEX = /,?\s*$/; +/** Bounded cache for parsePartialJsonArgs results. */ +const PARTIAL_JSON_CACHE_MAX_SIZE = 32; +const partialJsonCache = new Map | null>(); + +function cacheResult(input: string, result: Record | null): void { + if (partialJsonCache.size >= PARTIAL_JSON_CACHE_MAX_SIZE) { + partialJsonCache.delete(partialJsonCache.keys().next().value!); + } + partialJsonCache.set(input, result); +} + // Parse partial tool-arg JSON streamed token-by-token. Closes any // unterminated string and dangling open containers (in reverse order), // so parsers can still surface keys already received while the call -// is still in flight. +// is still in flight. Memoized: the char-by-char scanner runs on every +// render during streaming even when toolArgs hasn't changed. export function parsePartialJsonArgs(toolArgsString: string): Record | null { + const cached = partialJsonCache.get(toolArgsString); + if (cached !== undefined) return cached; + + let result: Record | null; + try { const parsed: unknown = JSON.parse(toolArgsString); - if (parsed && typeof parsed === 'object' && !Array.isArray(parsed)) { - return parsed as Record; - } - return null; + result = + parsed && typeof parsed === 'object' && !Array.isArray(parsed) + ? (parsed as Record) + : null; } catch { - let inString = false; - let escape = false; - const stack: ('{' | '[')[] = []; + result = scanPartialJson(toolArgsString); + } - for (let i = 0; i < toolArgsString.length; i++) { - const ch = toolArgsString[i]; - if (escape) { - escape = false; - continue; - } - if (ch === JSON_BACKSLASH && inString) { - escape = true; - continue; - } - if (ch === JSON_QUOTE) { - inString = !inString; - continue; - } - if (inString) continue; - if (ch === JSON_OBJECT_OPEN) stack.push(JSON_OBJECT_OPEN); - else if (ch === JSON_OBJECT_CLOSE) { - if (stack.length === 0 || stack[stack.length - 1] !== JSON_OBJECT_OPEN) return null; - stack.pop(); - } else if (ch === JSON_ARRAY_OPEN) stack.push(JSON_ARRAY_OPEN); - else if (ch === JSON_ARRAY_CLOSE) { - if (stack.length === 0 || stack[stack.length - 1] !== JSON_ARRAY_OPEN) return null; - stack.pop(); - } - } + cacheResult(toolArgsString, result); + return result; +} + +/** Char-by-char scanner for unterminated partial JSON. */ +function scanPartialJson(toolArgsString: string): Record | null { + let inString = false; + let escape = false; + const stack: ('{' | '[')[] = []; - let completed = toolArgsString; + for (let i = 0; i < toolArgsString.length; i++) { + const ch = toolArgsString[i]; if (escape) { - // Dangling escape at end of partial JSON: escape the trailing - // backslash as a literal so we can close the string cleanly. - completed += JSON_BACKSLASH; + escape = false; + continue; } - if (inString) completed += JSON_QUOTE; - if (!inString) completed = completed.replace(TRAILING_JSON_PUNCTUATION_REGEX, ''); - - // Close in reverse nesting order: innermost container first. - for (let i = stack.length - 1; i >= 0; i--) { - completed += stack[i] === JSON_OBJECT_OPEN ? JSON_OBJECT_CLOSE : JSON_ARRAY_CLOSE; + if (ch === JSON_BACKSLASH && inString) { + escape = true; + continue; } - - try { - const parsed: unknown = JSON.parse(completed); - if (parsed && typeof parsed === 'object' && !Array.isArray(parsed)) { - return parsed as Record; - } - return null; - } catch { - return null; + if (ch === JSON_QUOTE) { + inString = !inString; + continue; + } + if (inString) continue; + if (ch === JSON_OBJECT_OPEN) stack.push(JSON_OBJECT_OPEN); + else if (ch === JSON_OBJECT_CLOSE) { + if (stack.length === 0 || stack[stack.length - 1] !== JSON_OBJECT_OPEN) return null; + stack.pop(); + } else if (ch === JSON_ARRAY_OPEN) stack.push(JSON_ARRAY_OPEN); + else if (ch === JSON_ARRAY_CLOSE) { + if (stack.length === 0 || stack[stack.length - 1] !== JSON_ARRAY_OPEN) return null; + stack.pop(); } } + + let completed = toolArgsString; + if (escape) { + // Dangling escape at end of partial JSON: escape the trailing + // backslash as a literal so we can close the string cleanly. + completed += JSON_BACKSLASH; + } + if (inString) completed += JSON_QUOTE; + if (!inString) completed = completed.replace(TRAILING_JSON_PUNCTUATION_REGEX, ''); + + // Close in reverse nesting order: innermost container first. + for (let i = stack.length - 1; i >= 0; i--) { + completed += stack[i] === JSON_OBJECT_OPEN ? JSON_OBJECT_CLOSE : JSON_ARRAY_CLOSE; + } + + try { + const parsed: unknown = JSON.parse(completed); + return parsed && typeof parsed === 'object' && !Array.isArray(parsed) + ? (parsed as Record) + : null; + } catch { + return null; + } } diff --git a/tools/ui/src/lib/utils/search-results.ts b/tools/ui/src/lib/utils/search-results.ts index 861b0d2d3200..d090e6dc4bfb 100644 --- a/tools/ui/src/lib/utils/search-results.ts +++ b/tools/ui/src/lib/utils/search-results.ts @@ -155,41 +155,71 @@ function parseChunk(chunk: string): SearchResult | null { return result; } +/** Bounded cache for extractSearchResults results. */ +const SEARCH_RESULTS_CACHE_MAX_SIZE = 32; +const searchResultsCache = new Map(); + /** * Extract a SearchResult[] from a tool-result string. Returns `[]` when * the input does not match the expected shape — useful for branching * between dedicated search-results rendering and the generic tool-call - * block. + * block. Memoized: called per render during streaming on unchanged + * tool result strings. */ export function extractSearchResults(text: string | undefined | null): SearchResult[] { if (!text) return []; + const cached = searchResultsCache.get(text); + if (cached) return cached; + const results: SearchResult[] = []; for (const chunk of splitChunks(text)) { const parsed = parseChunk(chunk); if (parsed) results.push(parsed); } + + if (searchResultsCache.size >= SEARCH_RESULTS_CACHE_MAX_SIZE) { + searchResultsCache.delete(searchResultsCache.keys().next().value!); + } + searchResultsCache.set(text, results); + return results; } +/** Bounded cache for extractSearchQuery results. */ +const SEARCH_QUERY_CACHE_MAX_SIZE = 32; +const searchQueryCache = new Map(); + /** * Best-effort extraction of the search query out of a tool call's JSON * argument blob. Currently looks for a `query` field (the convention * used by Exa and most web-search MCP servers); returns an empty string - * if it cannot be located. + * if it cannot be located. Memoized: called per render during streaming + * on unchanged tool args strings. */ export function extractSearchQuery(toolArgs: string | undefined | null): string { if (!toolArgs) return ''; + + const cached = searchQueryCache.get(toolArgs); + if (cached !== undefined) return cached; + + let result = ''; try { const parsed: unknown = JSON.parse(toolArgs); if (parsed && typeof parsed === 'object' && !Array.isArray(parsed)) { const candidate = (parsed as Record)[SEARCH_TOOL_QUERY_FIELD]; - if (typeof candidate === 'string') return candidate.trim(); + if (typeof candidate === 'string') result = candidate.trim(); } } catch { - return ''; + result = ''; + } + + if (searchQueryCache.size >= SEARCH_QUERY_CACHE_MAX_SIZE) { + searchQueryCache.delete(searchQueryCache.keys().next().value!); } - return ''; + searchQueryCache.set(toolArgs, result); + + return result; } /** diff --git a/tools/ui/tests/unit/parse-toolcalls-memo.test.ts b/tools/ui/tests/unit/parse-toolcalls-memo.test.ts new file mode 100644 index 000000000000..85187febcba4 --- /dev/null +++ b/tools/ui/tests/unit/parse-toolcalls-memo.test.ts @@ -0,0 +1,146 @@ +// Tests for the memoized parseToolCalls and O(1) tool message lookup in +// deriveAgenticSections. These were added to prevent regressions where +// streaming text tokens trigger redundant JSON.parse calls on unchanged +// tool call data. + +import { describe, it, expect, vi } from 'vitest'; +import { deriveAgenticSections } from '$lib/utils/agentic'; +import type { ApiChatCompletionToolCall } from '$lib/types/api'; +import type { DatabaseMessage } from '$lib/types/database'; +import { MessageRole, AgenticSectionType } from '$lib/enums'; + +function makeMessage(overrides: Partial): DatabaseMessage { + return { + id: 'm1', + convId: 'c1', + type: 'text', + timestamp: 0, + role: MessageRole.ASSISTANT, + content: '', + parent: null, + children: [], + ...overrides + } as DatabaseMessage; +} + +describe('parseToolCalls memoization', () => { + it('returns the same array reference for the same JSON string', () => { + // parseToolCalls is not exported, but deriveAgenticSections uses it + // internally. We verify memoization through behavior: calling + // deriveAgenticSections twice with the same toolCalls should not + // re-parse (which we verify by checking the returned sections + // are equivalent). + const toolCallsJson = JSON.stringify([ + { id: 'call_1', type: 'function', function: { name: 'test', arguments: '{}' } } + ]); + + const msg = makeMessage({ content: 'hello', toolCalls: toolCallsJson }); + const sections1 = deriveAgenticSections(msg, [], [], false); + const sections2 = deriveAgenticSections(msg, [], [], false); + + expect(sections1).toHaveLength(sections2.length); + expect(sections1[0].type).toBe(sections2[0].type); + }); + + it('does not re-parse JSON on cache hit', () => { + const toolCallsJson = JSON.stringify([ + { id: 'call_1', type: 'function', function: { name: 'test', arguments: '{}' } } + ]); + + const msg = makeMessage({ content: 'hello', toolCalls: toolCallsJson }); + const spy = vi.spyOn(JSON, 'parse'); + + deriveAgenticSections(msg, [], [], false); + const callsAfterFirst = spy.mock.calls.length; + + deriveAgenticSections(msg, [], [], false); + expect(spy.mock.calls.length).toBe(callsAfterFirst); + + spy.mockRestore(); + }); + + it('handles empty/undefined toolCalls without error', () => { + const msg = makeMessage({ content: 'hello' }); + const sections = deriveAgenticSections(msg, [], [], false); + + expect(sections).toHaveLength(1); + expect(sections[0].type).toBe(AgenticSectionType.TEXT); + }); + + it('handles invalid JSON gracefully', () => { + const msg = makeMessage({ content: 'hello', toolCalls: '{invalid' }); + const sections = deriveAgenticSections(msg, [], [], false); + + // Should return just the text section, no tool call sections + expect(sections).toHaveLength(1); + expect(sections[0].type).toBe(AgenticSectionType.TEXT); + }); +}); + +describe('deriveAgenticSections O(1) tool message lookup', () => { + it('matches tool messages to tool calls by toolCallId', () => { + const toolCallsJson = JSON.stringify([ + { id: 'call_1', type: 'function', function: { name: 'test_1', arguments: '{}' } }, + { id: 'call_2', type: 'function', function: { name: 'test_2', arguments: '{}' } } + ]); + + const toolMessages = [ + makeMessage({ role: MessageRole.TOOL, toolCallId: 'call_1', content: 'result_1' }), + makeMessage({ role: MessageRole.TOOL, toolCallId: 'call_2', content: 'result_2' }) + ]; + + const msg = makeMessage({ content: 'hello', toolCalls: toolCallsJson }); + const sections = deriveAgenticSections(msg, toolMessages, [], false); + + // Expect: TEXT + 2 TOOL_CALL sections + const toolCallSections = sections.filter((s) => s.type === AgenticSectionType.TOOL_CALL); + expect(toolCallSections).toHaveLength(2); + expect(toolCallSections[0].toolResult).toBe('result_1'); + expect(toolCallSections[1].toolResult).toBe('result_2'); + }); + + it('handles missing tool messages (pending calls during streaming)', () => { + const toolCallsJson = JSON.stringify([ + { id: 'call_1', type: 'function', function: { name: 'test', arguments: '{}' } } + ]); + + const msg = makeMessage({ content: '', toolCalls: toolCallsJson }); + const sections = deriveAgenticSections(msg, [], [], true); + + const toolCallSection = sections.find((s) => s.type === AgenticSectionType.TOOL_CALL_PENDING); + expect(toolCallSection).toBeDefined(); + expect(toolCallSection?.content).toBe(''); + }); + + it('scales with many tool calls (no O(n^2) blowup)', () => { + const N = 100; + const toolCalls = Array.from( + { length: N }, + (_, i): ApiChatCompletionToolCall => ({ + id: `call_${i}`, + type: 'function', + function: { name: `tool_${i}`, arguments: '{}' } + }) + ); + const toolCallsJson = JSON.stringify(toolCalls); + + const toolMessages = Array.from({ length: N }, (_, i) => + makeMessage({ + role: MessageRole.TOOL, + toolCallId: `call_${i}`, + content: `result_${i}` + }) + ); + + const msg = makeMessage({ content: 'hello', toolCalls: toolCallsJson }); + + // If the lookup were still O(n^2), this would be noticeably slow + const start = Date.now(); + const sections = deriveAgenticSections(msg, toolMessages, [], false); + const elapsed = Date.now() - start; + + const toolCallSections = sections.filter((s) => s.type === AgenticSectionType.TOOL_CALL); + expect(toolCallSections).toHaveLength(N); + expect(elapsed).toBeLessThan(100); // Should be fast with O(1) lookup + }); +}); From 7e1e28cae36d41fe7bbe9dae7c9625de6565c063 Mon Sep 17 00:00:00 2001 From: Daniel Bevenius Date: Tue, 28 Jul 2026 17:20:25 +0200 Subject: [PATCH 042/190] mtmd : add Nemotron 3 Nano Omni support (parakeet) (#22520) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * mtmd : add Nemotron 3 Nano Omni support (parakeet) This commit adds support for the subsampling and encoder part of Nemotron Nemo 3 omni model. The Parakeet subsampling/encoder were taken from parakeet.cpp which is currently a pull request against whisper.cpp. I've tried to copy the code a close as possible to hopefully enable easy patching between the these two project later. Refs: https://github.com/ggml-org/whisper.cpp/pull/3735 * mtmd : generate rel pos tensor in graph instead of in conversion [no ci] This commit removes the generation of the relative positional tensor in the model conversion script and instead computes it in the encoder graph. This is only done for the window of positions required for the current audio sample. * mtmd : add clip_get_model to clip API [no ci] This commit adds a function to get access to the clip_model. It also removes the two functions clip_get_mel_filter_tensor, and clip_get_window_tensor(const struct clip_ctx * ctx) which can now use clip_get_model to access the model tensors that it needs. * mtmd : read mel_filters and window into hparams * mtmd : use set_input_f32 lambda [no ci] * mtmd : add better asserts for mel_filters and hann window [no ci] * mtmd : add missing size_t cast * mtmd : change type of pad to size_t * mtmd : zero initialize samples_padded * mtmd : remove unsued ctx member from parakeet preprocessor * mtmd : make log_mel_spectrogram_parakeet_worker_thread private static * mtmd : sync/update parakeeet impl with latest whisper.cpp This commit updates the parakeet code in mtmd to reflect the latest updates to parakeet.cpp in whisper.cpp. A follow up commit will address the currently hardcoded dw_pad and see if we can add n_conv_kernel as a model metadata field. * mtmd : add audio_conv_kernel_size to model conversion This commit updates the model conversion to read the conv_kernel_size field from the sound_config section of the models config.json file. It then uses this field instead of the hardcoded values in parakeet.cpp. * mtmd : cleanup [no ci] * conversion : call super().filter_tensors [no ci] * do not discard result of super filter_tensors * mtmd : use build_mm instead of ggml_mul_mat * mtmd : use build_ffn * mtmd : move and reuse get_vector lambda * mtmd : use build_inp_raw for parakeet * mtmd : throw exception in get_scalar instead of assert * mtmd : fix std::min call * mtmt : use .c_str in throw clause in get_vector * mtmd : check for F32 type and non-empty tensor in get_vector The get_vector lambda is used by get_scalar but also standalone to read in the mel_filters and the window data. Therefor we are not checking for 1D tensors but allowing multiple dimensions. We do have a check in get_scalar to verify the size of the vector. * mtmd : replace hardcoded 1101 for n_tokens_real * mtmd : assert subsampling_factor is 8 This commit adds an assert of the parakeet subsampling factor to check that it is 8. The motivation for this is that this model currently has three convolutions with a stride of 2. If the underlying model updates the subsampling factor these convolution operations will need to be updated and this will produce and error if this occurs. * mtmd : remove unused ggml_tensors attn_pos_w and mm_norm_w * mtmd : remove single thread path This commit removes the single thread path which was a left over from the original parakeet.cpp where n_threads is configurable. * fix some security issues --------- Co-authored-by: Sigbjørn Skjæret <1629204+CISC@users.noreply.github.com> Co-authored-by: Xuan Son Nguyen --- conversion/nemotron.py | 57 ++++- gguf-py/gguf/constants.py | 14 ++ gguf-py/gguf/gguf_writer.py | 3 + gguf-py/gguf/tensor_mapping.py | 40 ++++ tools/mtmd/CMakeLists.txt | 1 + tools/mtmd/clip-impl.h | 9 + tools/mtmd/clip-model.h | 24 +- tools/mtmd/clip.cpp | 210 +++++++++++++++- tools/mtmd/models/models.h | 5 + tools/mtmd/models/parakeet.cpp | 421 +++++++++++++++++++++++++++++++++ tools/mtmd/mtmd-audio.cpp | 203 ++++++++++++++++ tools/mtmd/mtmd-audio.h | 15 ++ tools/mtmd/mtmd.cpp | 4 + 13 files changed, 985 insertions(+), 21 deletions(-) create mode 100644 tools/mtmd/models/parakeet.cpp diff --git a/conversion/nemotron.py b/conversion/nemotron.py index e44688a78807..0572b42ca2a1 100644 --- a/conversion/nemotron.py +++ b/conversion/nemotron.py @@ -39,28 +39,48 @@ def get_vision_config(self) -> dict[str, Any] | None: } return vision_config + def get_audio_config(self) -> dict[str, Any] | None: + return self.global_config.get("sound_config") + def set_gguf_parameters(self): if "image_mean" not in self.preprocessor_config: self.preprocessor_config["image_mean"] = [0.485, 0.456, 0.406] if "image_std" not in self.preprocessor_config: self.preprocessor_config["image_std"] = [0.229, 0.224, 0.225] + if self.hparams_audio is not None: + self.has_vision_encoder = True + self.has_audio_encoder = True + self.gguf_writer.add_audio_num_mel_bins(self.hparams_audio["num_mel_bins"]) + self.gguf_writer.add_audio_attention_layernorm_eps(1e-5) + self.gguf_writer.add_audio_subsampling_factor(self.hparams_audio["subsampling_factor"]) + self.gguf_writer.add_audio_conv_kernel_size(self.hparams_audio["conv_kernel_size"]) + self.gguf_writer.add_clip_audio_projector_type(gguf.VisionProjectorType.PARAKEET) + self.gguf_writer.add_clip_vision_projector_type(gguf.VisionProjectorType.NEMOTRON_V2_VL) + else: + self.gguf_writer.add_clip_projector_type(gguf.VisionProjectorType.NEMOTRON_V2_VL) + super().set_gguf_parameters() hparams = self.global_config - self.gguf_writer.add_clip_projector_type(gguf.VisionProjectorType.NEMOTRON_V2_VL) self.gguf_writer.add_vision_attention_layernorm_eps(1e-6) self.gguf_writer.add_vision_use_gelu(True) downsample_ratio = hparams.get("downsample_ratio", 0.5) self.gguf_writer.add_vision_projector_scale_factor(int(1.0 / downsample_ratio)) def tensor_force_quant(self, name, new_name, bid, n_dims): - if ".position_embd." in new_name or "pos_embed" in new_name: - return gguf.GGMLQuantizationType.F32 + if "sound_encoder" in name or new_name.startswith("mm.a."): + if "bias" in new_name or "norm" in new_name: + return gguf.GGMLQuantizationType.F32 + if "conv" in new_name and "weight" in new_name: + return gguf.GGMLQuantizationType.F32 + return super().tensor_force_quant(name, new_name, bid, n_dims) @classmethod def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None: - name, gen = item + if (titem := super().filter_tensors(item)) is None: + return None + name, gen = titem if "input_conditioner" in name: return None @@ -69,14 +89,18 @@ def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Ca if "radio_model.model.patch_generator.video_embedder" in name: return None - if not name.startswith("vision_model.radio_model.model.") and not name.startswith("mlp1."): + if not name.startswith(("vision_model.radio_model.model.", "mlp1.", "sound_encoder.", "sound_projection.")): return None if "patch_generator.pos_embed" in name: if not name.endswith(".weight"): name += ".weight" - return super().filter_tensors((name, gen)) + # num_batches is only used for training not inference. + if "conv.norm" in name and "num_batches" in name: + return None + + return name, gen def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: # RADIO's pos_embed doesn't have .weight suffix, but clip.cpp expects it @@ -104,7 +128,26 @@ def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iter n_embd = self.hparams["hidden_size"] data_torch = data_torch.reshape(n_embd, 3, patch_size, patch_size) - yield from super().modify_tensors(data_torch, name, bid) + if "depthwise_conv.weight" in name: + data_torch = data_torch.unsqueeze(-1) + data_torch = data_torch.permute(3, 1, 0, 2).contiguous() + + if "pointwise_conv" in name and name.endswith(".weight"): + if len(data_torch.shape) == 3 and data_torch.shape[2] == 1: + data_torch = data_torch.reshape(data_torch.shape[0], data_torch.shape[1]) + + if "subsampling.layers" in name and name.endswith(".bias"): + if len(data_torch.shape) == 1: + data_torch = data_torch.reshape(1, -1, 1, 1) + + if "pointwise_conv" in name and name.endswith(".bias"): + if len(data_torch.shape) == 1: + data_torch = data_torch.reshape(1, -1, 1, 1) + + for mapped_name, tensor in super().modify_tensors(data_torch, name, bid): + if name.startswith("sound_projection.") and mapped_name.startswith("mm.model.mlp."): + mapped_name = mapped_name.replace("mm.model.mlp.", "mm.a.mlp.") + yield mapped_name, tensor @ModelBase.register("NemotronForCausalLM") diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index 2ebefaa3b422..124ea28b0616 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -373,6 +373,7 @@ class ClipAudio: FEED_FORWARD_LENGTH = "clip.audio.feed_forward_length" PROJECTION_DIM = "clip.audio.projection_dim" BLOCK_COUNT = "clip.audio.block_count" + SUBSAMPLING_FACTOR = "clip.audio.subsampling_factor" CHUNK_SIZE = "clip.audio.chunk_size" CONV_KERNEL_SIZE = "clip.audio.conv_kernel_size" MAX_POS_EMB = "clip.audio.max_pos_emb" @@ -1002,6 +1003,10 @@ class MODEL_TENSOR(IntEnum): A_ENC_CONV_NORM = auto() # SSM conv A_ENC_CONV_PW1 = auto() A_ENC_CONV_PW2 = auto() + A_ENC_CONV_NORM_MEAN = auto() # parakeet + A_ENC_CONV_NORM_VAR = auto() # parakeet + A_ENC_MEL_FILTERS = auto() # parakeet + A_ENC_WINDOW = auto() # parakeet A_CTC_OUT = auto() A_CTC_OUT_MID = auto() A_ENC_ATTN_REL_POS_EMB = auto() @@ -1591,6 +1596,10 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.A_ENC_CONV_NORM: "a.blk.{bid}.conv_norm", MODEL_TENSOR.A_ENC_CONV_PW1: "a.blk.{bid}.conv_pw1", MODEL_TENSOR.A_ENC_CONV_PW2: "a.blk.{bid}.conv_pw2", + MODEL_TENSOR.A_ENC_CONV_NORM_MEAN: "a.blk.{bid}.conv_norm_mean", + MODEL_TENSOR.A_ENC_CONV_NORM_VAR: "a.blk.{bid}.conv_norm_var", + MODEL_TENSOR.A_ENC_MEL_FILTERS: "a.mel_filters", + MODEL_TENSOR.A_ENC_WINDOW: "a.window", MODEL_TENSOR.A_CTC_OUT: "a.enc_ctc_out", MODEL_TENSOR.A_CTC_OUT_MID: "a.enc_ctc_out_mid", MODEL_TENSOR.A_ENC_ATTN_REL_POS_EMB: "a.blk.{bid}.attn_rel_pos_emb", @@ -1810,6 +1819,10 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.A_ENC_CONV_NORM, MODEL_TENSOR.A_ENC_CONV_PW1, MODEL_TENSOR.A_ENC_CONV_PW2, + MODEL_TENSOR.A_ENC_CONV_NORM_MEAN, + MODEL_TENSOR.A_ENC_CONV_NORM_VAR, + MODEL_TENSOR.A_ENC_MEL_FILTERS, + MODEL_TENSOR.A_ENC_WINDOW, MODEL_TENSOR.A_MM_INP_PROJ, MODEL_TENSOR.A_MM_SOFT_EMB_NORM, MODEL_TENSOR.A_MM_EMBEDDING, @@ -4861,6 +4874,7 @@ class VisionProjectorType: YOUTUVL = "youtuvl" NEMOTRON_V2_VL = "nemotron_v2_vl" HUNYUANVL = "hunyuanvl" + PARAKEET = "parakeet" # audio MINIMAXM3 = "minimax_m3" MINICPMV4_6 = "minicpmv4_6" GRANITE_SPEECH = "granite_speech" # audio diff --git a/gguf-py/gguf/gguf_writer.py b/gguf-py/gguf/gguf_writer.py index 657ed69b6895..3aa4f049f2bb 100644 --- a/gguf-py/gguf/gguf_writer.py +++ b/gguf-py/gguf/gguf_writer.py @@ -1374,6 +1374,9 @@ def add_audio_local_group_size(self, value: int) -> None: def add_audio_stack_factor(self, value: int) -> None: self.add_uint32(Keys.ClipAudio.Projector.STACK_FACTOR, value) + def add_audio_subsampling_factor(self, value: int) -> None: + self.add_uint32(Keys.ClipAudio.SUBSAMPLING_FACTOR, value) + def add_audio_chunk_size(self, value: int) -> None: self.add_uint32(Keys.ClipAudio.CHUNK_SIZE, value) diff --git a/gguf-py/gguf/tensor_mapping.py b/gguf-py/gguf/tensor_mapping.py index 5562d43277a4..1e991b873cea 100644 --- a/gguf-py/gguf/tensor_mapping.py +++ b/gguf-py/gguf/tensor_mapping.py @@ -2107,6 +2107,7 @@ class TensorNameMap: "conformer.pre_encode.conv.{bid}", # lfm2 "model.audio_tower.subsample_conv_projection.conv_{bid}.conv", # gemma3n "conformer.subsample_conv_projection.layer{bid}.conv", # gemma4 + "sound_encoder.encoder.subsampling.layers.{bid}", # parakeet "encoder.conv{bid}", # mimo-audio-tokenizer ), @@ -2140,6 +2141,7 @@ class TensorNameMap: "conformer.layers.{bid}.self_attn.linear_q", # lfm2 "conformer.layers.{bid}.attention.attn.q_proj", # gemma3n "conformer.layers.{bid}.self_attn.q_proj", # gemma4 + "sound_encoder.encoder.layers.{bid}.self_attn.q_proj", # parakeet "encoder.layers.{bid}.attn.to_q", # granite_speech "encoder.layers.{bid}.self_attn.q_proj", # mimo-audio-tokenizer ), @@ -2149,6 +2151,7 @@ class TensorNameMap: "conformer.layers.{bid}.self_attn.linear_k", # lfm2 "conformer.layers.{bid}.attention.attn.k_proj", # gemma3n "conformer.layers.{bid}.self_attn.k_proj", # gemma4 + "sound_encoder.encoder.layers.{bid}.self_attn.k_proj", # parakeet "encoder.layers.{bid}.attn.to_k", # granite_speech (split from to_kv) "encoder.layers.{bid}.self_attn.k_proj", # mimo-audio-tokenizer ), @@ -2158,6 +2161,7 @@ class TensorNameMap: "conformer.layers.{bid}.self_attn.linear_v", # lfm2 "conformer.layers.{bid}.attention.attn.v_proj", # gemma3n "conformer.layers.{bid}.self_attn.v_proj", # gemma4 + "sound_encoder.encoder.layers.{bid}.self_attn.v_proj", # parakeet "encoder.layers.{bid}.attn.to_v", # granite_speech (split from to_kv) "encoder.layers.{bid}.self_attn.v_proj", # mimo-audio-tokenizer ), @@ -2187,6 +2191,7 @@ class TensorNameMap: "audio_tower.layers.{bid}.self_attn_layer_norm", # ultravox "conformer.layers.{bid}.norm_self_att", # lfm2 "conformer.layers.{bid}.attention.pre_attn_norm", # gemma3n + "sound_encoder.encoder.layers.{bid}.norm_self_att", # parakeet "encoder.layers.{bid}.attn.pre_norm", # granite_speech "encoder.layers.{bid}.self_attn_layer_norm", # mimo-audio-tokenizer ), @@ -2196,6 +2201,7 @@ class TensorNameMap: "conformer.layers.{bid}.self_attn.linear_out", # lfm2 "conformer.layers.{bid}.attention.post", # gemma3n "conformer.layers.{bid}.self_attn.post", # gemma4 + "sound_encoder.encoder.layers.{bid}.self_attn.o_proj", # parakeet "encoder.layers.{bid}.attn.to_out", # granite_speech "encoder.layers.{bid}.self_attn.out_proj", # mimo-audio-tokenizer ), @@ -2204,6 +2210,7 @@ class TensorNameMap: "audio_tower.layers.{bid}.final_layer_norm", # ultravox "conformer.layers.{bid}.norm_out", # lfm2 "conformer.layers.{bid}.attention.post_norm", # gemma3n + "sound_encoder.encoder.layers.{bid}.norm_out", # parakeet "encoder.layers.{bid}.post_norm", # granite_speech "encoder.layers.{bid}.final_layer_norm", # mimo-audio-tokenizer ), @@ -2212,6 +2219,7 @@ class TensorNameMap: "conformer.layers.{bid}.norm_feed_forward1", # lfm2 "conformer.layers.{bid}.ffw_layer_start.pre_layer_norm", # gemma3n "conformer.layers.{bid}.feed_forward1.pre_layer_norm", # gemma4 + "sound_encoder.encoder.layers.{bid}.norm_feed_forward1", # parakeet "encoder.layers.{bid}.ff1.pre_norm", # granite_speech ), @@ -2229,6 +2237,7 @@ class TensorNameMap: "conformer.layers.{bid}.feed_forward1.linear1", # lfm2 "conformer.layers.{bid}.ffw_layer_start.ffw_layer_1", # gemma3n "conformer.layers.{bid}.feed_forward1.ffw_layer_1", # gemma4 + "sound_encoder.encoder.layers.{bid}.feed_forward1.linear1", # parakeet "encoder.layers.{bid}.ff1.up_proj", # granite_speech "encoder.layers.{bid}.fc1", # mimo-audio-tokenizer ), @@ -2240,6 +2249,7 @@ class TensorNameMap: "conformer.layers.{bid}.feed_forward1.linear2", # lfm2 "conformer.layers.{bid}.ffw_layer_start.ffw_layer_2", # gemma3n "conformer.layers.{bid}.feed_forward1.ffw_layer_2", # gemma4 + "sound_encoder.encoder.layers.{bid}.feed_forward1.linear2", # parakeet "encoder.layers.{bid}.ff1.down_proj", # granite_speech "encoder.layers.{bid}.fc2", # mimo-audio-tokenizer ), @@ -2248,6 +2258,7 @@ class TensorNameMap: "conformer.layers.{bid}.feed_forward2.linear1", # lfm2 "conformer.layers.{bid}.ffw_layer_end.ffw_layer_1", # gemma3n "conformer.layers.{bid}.feed_forward2.ffw_layer_1", # gemma4 + "sound_encoder.encoder.layers.{bid}.feed_forward2.linear1", # parakeet "encoder.layers.{bid}.ff2.up_proj", # granite_speech ), @@ -2255,6 +2266,7 @@ class TensorNameMap: "conformer.layers.{bid}.feed_forward2.linear2", # lfm2 "conformer.layers.{bid}.ffw_layer_end.ffw_layer_2", # gemma3n "conformer.layers.{bid}.feed_forward2.ffw_layer_2", # gemma4 + "sound_encoder.encoder.layers.{bid}.feed_forward2.linear2", # parakeet "encoder.layers.{bid}.ff2.down_proj", # granite_speech ), @@ -2262,6 +2274,7 @@ class TensorNameMap: "conformer.layers.{bid}.norm_feed_forward2", # lfm2 "conformer.layers.{bid}.ffw_layer_end.pre_layer_norm", # gemma3n "conformer.layers.{bid}.feed_forward2.pre_layer_norm", # gemma4 + "sound_encoder.encoder.layers.{bid}.norm_feed_forward2", # parakeet "encoder.layers.{bid}.ff2.pre_norm", # granite_speech ), @@ -2290,20 +2303,24 @@ class TensorNameMap: MODEL_TENSOR.A_ENC_LINEAR_POS: ( "conformer.layers.{bid}.self_attn.linear_pos", # lfm2 "conformer.layers.{bid}.attention.attn.relative_position_embedding.pos_proj", # gemma3n + "sound_encoder.encoder.layers.{bid}.self_attn.relative_k_proj", # parakeet ), MODEL_TENSOR.A_ENC_POS_BIAS_U: ( "conformer.layers.{bid}.self_attn.pos_bias_u", # lfm2 + "sound_encoder.encoder.layers.{bid}.self_attn.bias_u", # parakeet ), MODEL_TENSOR.A_ENC_POS_BIAS_V: ( "conformer.layers.{bid}.self_attn.pos_bias_v", # lfm2 + "sound_encoder.encoder.layers.{bid}.self_attn.bias_v", # parakeet ), MODEL_TENSOR.A_ENC_OUT: ( "conformer.pre_encode.out", # lfm2 "model.audio_tower.subsample_conv_projection.input_proj_linear", # gemma3n (note: it should be A_ENC_INP_PROJ, this is a mistake; it should be corrected in C++ code when it's supported) "conformer.output_proj", # gemma4 + "sound_encoder.encoder.subsampling.linear", # parakeet ), # note: some tensors below has "audio." pseudo-prefix, to prevent conflicts with vision tensors @@ -2313,6 +2330,7 @@ class TensorNameMap: "audio.multi_modal_projector.linear_{bid}", # ultravox, meralion "audio_adapter.model.{bid}", # lfm2 "audio_tower.proj{bid}", # qwen3omni + "sound_projection.linear{bid}", # parakeet (linear1, linear2) ), MODEL_TENSOR.A_MMPROJ_FC: ( @@ -2323,6 +2341,7 @@ class TensorNameMap: MODEL_TENSOR.A_MM_NORM_PRE: ( "audio.multi_modal_projector.ln_pre", # ultravox + "sound_projection.norm", # parakeet ), MODEL_TENSOR.A_MM_NORM_MID: ( @@ -2368,30 +2387,43 @@ class TensorNameMap: MODEL_TENSOR.A_ENC_CONV_DW: ( "conformer.layers.{bid}.conv.depthwise_conv", # lfm2 "conformer.layers.{bid}.lconv1d.depthwise_conv1d", # gemma3n + "sound_encoder.encoder.layers.{bid}.conv.depthwise_conv", # parakeet "encoder.layers.{bid}.conv.depth_conv.conv", # granite_speech ), MODEL_TENSOR.A_ENC_CONV_NORM: ( "conformer.layers.{bid}.conv.batch_norm", # lfm2 "conformer.layers.{bid}.lconv1d.pre_layer_norm", # gemma3n + "sound_encoder.encoder.layers.{bid}.conv.norm", # parakeet + ), + + MODEL_TENSOR.A_ENC_CONV_NORM_MEAN: ( + "sound_encoder.encoder.layers.{bid}.conv.norm.running_mean", # parakeet + ), + + MODEL_TENSOR.A_ENC_CONV_NORM_VAR: ( + "sound_encoder.encoder.layers.{bid}.conv.norm.running_var", # parakeet "encoder.layers.{bid}.conv.batch_norm", # granite_speech ), MODEL_TENSOR.A_ENC_CONV_PW1: ( "conformer.layers.{bid}.conv.pointwise_conv1", # lfm2 "conformer.layers.{bid}.lconv1d.linear_start", # gemma3n + "sound_encoder.encoder.layers.{bid}.conv.pointwise_conv1", # parakeet "encoder.layers.{bid}.conv.up_conv", # granite_speech ), MODEL_TENSOR.A_ENC_CONV_PW2: ( "conformer.layers.{bid}.conv.pointwise_conv2", # lfm2 "conformer.layers.{bid}.lconv1d.linear_end", # gemma3n + "sound_encoder.encoder.layers.{bid}.conv.pointwise_conv2", # parakeet "encoder.layers.{bid}.conv.down_conv", # granite_speech ), MODEL_TENSOR.A_ENC_NORM_CONV: ( "conformer.layers.{bid}.norm_conv", # lfm2 "conformer.layers.{bid}.lconv1d.conv_norm", # gemma3n + "sound_encoder.encoder.layers.{bid}.norm_conv", # parakeet "encoder.layers.{bid}.conv.norm", # granite_speech ), @@ -2403,6 +2435,14 @@ class TensorNameMap: "conformer.layers.{bid}.attention.attn.per_dim_scale", # gemma4 ), + MODEL_TENSOR.A_ENC_MEL_FILTERS: ( + "sound_encoder.encoder.feature_extractor.featurizer.fb", # parakeet + ), + + MODEL_TENSOR.A_ENC_WINDOW: ( + "sound_encoder.encoder.feature_extractor.featurizer.window", # parakeet + ), + MODEL_TENSOR.A_MM_EMBEDDING: ( "model.embed_audio.embedding", # gemma3n ), diff --git a/tools/mtmd/CMakeLists.txt b/tools/mtmd/CMakeLists.txt index 18a8288ba048..15040e4af5f9 100644 --- a/tools/mtmd/CMakeLists.txt +++ b/tools/mtmd/CMakeLists.txt @@ -60,6 +60,7 @@ add_library(mtmd models/mobilenetv5.cpp models/youtuvl.cpp models/yasa2.cpp + models/parakeet.cpp ) set_target_properties(mtmd PROPERTIES diff --git a/tools/mtmd/clip-impl.h b/tools/mtmd/clip-impl.h index 09204113801f..589fc724ed08 100644 --- a/tools/mtmd/clip-impl.h +++ b/tools/mtmd/clip-impl.h @@ -88,6 +88,7 @@ #define KEY_A_ATTN_WINDOW_SIZE "clip.audio.window_size" // mimo-audio-tokenizer: sliding-window radius #define KEY_A_LOCAL_BLOCK_COUNT "clip.audio.local_block_count" // mimo-v2.5: input_local_transformer layer count #define KEY_A_LOCAL_GROUP_SIZE "clip.audio.local_group_size" // mimo-v2.5: input_local_transformer grouping size +#define KEY_AUDIO_SUBSAMPLING_FACTOR "clip.audio.subsampling_factor" // // tensor name constants @@ -338,6 +339,12 @@ #define TN_YASA_STAGE_DOWN_CONV "v.stage.%d.down.conv.%s" #define TN_YASA_STAGE_BLK "v.stage.%d.blk.%d.%s.%s" +// parakeet +#define TN_MEL_FILTERS "a.mel_filters" +#define TN_WINDOW "a.window" +#define TN_CONV_NORM_MEAN "%s.blk.%d.conv_norm_mean" +#define TN_CONV_NORM_VAR "%s.blk.%d.conv_norm_var" + // align x to upper multiple of n #define CLIP_ALIGN(x, n) ((((x) + (n) - 1) / (n)) * (n)) @@ -392,6 +399,7 @@ enum projector_type { PROJECTOR_TYPE_KIMIK25, PROJECTOR_TYPE_NEMOTRON_V2_VL, PROJECTOR_TYPE_HUNYUANVL, + PROJECTOR_TYPE_PARAKEET, PROJECTOR_TYPE_EXAONE4_5, PROJECTOR_TYPE_MINICPMV4_6, PROJECTOR_TYPE_GRANITE_SPEECH, @@ -455,6 +463,7 @@ static std::map PROJECTOR_TYPE_NAMES = { { PROJECTOR_TYPE_MINIMAX_M3, "minimax_m3"}, { PROJECTOR_TYPE_GRANITE4_VISION, "granite4_vision"}, { PROJECTOR_TYPE_MIMO_AUDIO, "mimo_audio"}, + { PROJECTOR_TYPE_PARAKEET, "parakeet"}, }; static projector_type clip_projector_type_from_string(const std::string & str) { diff --git a/tools/mtmd/clip-model.h b/tools/mtmd/clip-model.h index 8dc87549766e..146eabce23b7 100644 --- a/tools/mtmd/clip-model.h +++ b/tools/mtmd/clip-model.h @@ -110,6 +110,8 @@ struct clip_hparams { // audio int32_t n_mel_bins = 0; // whisper preprocessor int32_t proj_stack_factor = 0; // ultravox + int32_t subsampling_factor = 0; // parakeet + int32_t audio_chunk_size = 0; int32_t audio_conv_kernel_size = 0; int32_t audio_max_pos_emb = 0; @@ -124,6 +126,10 @@ struct clip_hparams { int32_t audio_window_len = -1; int32_t audio_hop_len = -1; + // parakeet + std::vector mel_filters; + std::vector window; + // mimo-audio-tokenizer: residual vector quantizer int32_t rvq_num_quantizers = 0; std::vector rvq_codebook_size; // per-quantizer bin count (ragged, e.g. 1024/1024/256/128x17) @@ -245,14 +251,16 @@ struct clip_layer { ggml_tensor * norm_conv_b = nullptr; ggml_tensor * linear_pos_w = nullptr; - ggml_tensor * conv_norm_w = nullptr; - ggml_tensor * conv_norm_b = nullptr; - ggml_tensor * conv_dw_w = nullptr; - ggml_tensor * conv_dw_b = nullptr; - ggml_tensor * conv_pw1_w = nullptr; - ggml_tensor * conv_pw1_b = nullptr; - ggml_tensor * conv_pw2_w = nullptr; - ggml_tensor * conv_pw2_b = nullptr; + ggml_tensor * conv_norm_w = nullptr; + ggml_tensor * conv_norm_b = nullptr; + ggml_tensor * conv_norm_mean = nullptr; // parakeet + ggml_tensor * conv_norm_var = nullptr; // parakeet + ggml_tensor * conv_dw_w = nullptr; + ggml_tensor * conv_dw_b = nullptr; + ggml_tensor * conv_pw1_w = nullptr; + ggml_tensor * conv_pw1_b = nullptr; + ggml_tensor * conv_pw2_w = nullptr; + ggml_tensor * conv_pw2_b = nullptr; // gemma4 audio conformer per-layer ggml_tensor * attn_pre_norm_w = nullptr; diff --git a/tools/mtmd/clip.cpp b/tools/mtmd/clip.cpp index 04614b93bd27..11f9820edeb8 100644 --- a/tools/mtmd/clip.cpp +++ b/tools/mtmd/clip.cpp @@ -1033,6 +1033,10 @@ static std::unique_ptr clip_get_graph_builder(clip_ctx * ctx, const { builder = std::make_unique(ctx, img); } break; + case PROJECTOR_TYPE_PARAKEET: + { + builder = std::make_unique(ctx, img); + } break; case PROJECTOR_TYPE_GRANITE4_VISION: { builder = std::make_unique(ctx, img); @@ -1356,6 +1360,20 @@ struct clip_model_loader { { get_u32(KEY_PROJ_SCALE_FACTOR, hparams.n_merge, false); } break; + case PROJECTOR_TYPE_PARAKEET: + { + get_u32(KEY_AUDIO_SUBSAMPLING_FACTOR, hparams.subsampling_factor); + GGML_ASSERT(hparams.subsampling_factor == 8 && + "subsampling_factor must match the conv strides in clip_graph_parakeet::build()"); + get_u32(KEY_A_CONV_KERNEL_SIZE, hparams.audio_conv_kernel_size); + GGML_ASSERT(hparams.audio_conv_kernel_size > 0 && hparams.audio_conv_kernel_size % 2 == 1 && + "audio_conv_kernel_size must be a positive odd integer"); + hparams.audio_chunk_len = 0; + hparams.audio_sample_rate = 16000; + hparams.audio_n_fft = 512; + hparams.audio_window_len = 400; + hparams.audio_hop_len = 160; + } break; case PROJECTOR_TYPE_IDEFICS3: { // use default llava-uhd preprocessing params @@ -1893,16 +1911,46 @@ struct clip_model_loader { return cur; }; - auto get_scalar = [&](const std::string & name, float default_val) { + auto get_vector = [&](const std::string & name) { + std::vector result; auto it = tensor_offset.find(name); if (it == tensor_offset.end()) { + return result; + } + + const int64_t idx = gguf_find_tensor(ctx_gguf.get(), name.c_str()); + if (idx < 0) { + throw std::runtime_error(string_format("%s: failed to find tensor %s\n", __func__, name.c_str())); + } + + if (const auto type = gguf_get_tensor_type(ctx_gguf.get(), idx); type != GGML_TYPE_F32) { + throw std::runtime_error(string_format("%s: %s must be %s, was %s\n", __func__, + name.c_str(), ggml_type_name(GGML_TYPE_F32), ggml_type_name(type))); + } + + const size_t n_bytes = gguf_get_tensor_size(ctx_gguf.get(), idx); + if (n_bytes == 0) { + throw std::runtime_error(string_format("%s: tensor %s is empty\n", __func__, name.c_str())); + } + + const size_t n_elems = n_bytes / sizeof(float); + result.resize(n_elems); + fin.seekg(it->second, std::ios::beg); + fin.read(reinterpret_cast(result.data()), n_bytes); + return result; + }; + + auto get_scalar = [&](const std::string & name, float default_val) { + auto v = get_vector(name); + if (v.empty()) { return default_val; } - size_t offset = it->second; - fin.seekg(offset, std::ios::beg); - float value; - fin.read(reinterpret_cast(&value), sizeof(float)); - return value; + if (v.size() != 1) { + throw std::runtime_error(string_format("%s: expected scalar tensor '%s' but got %d elements\n", + __func__, name.c_str(), (int) v.size())); + } + + return v[0]; }; model.class_embedding = get_tensor(TN_CLASS_EMBD, false); @@ -2800,6 +2848,68 @@ struct clip_model_loader { layer.conv_pw2_b = get_tensor(string_format(TN_CONV_PW2, prefix, il, "bias")); } } break; + case PROJECTOR_TYPE_PARAKEET: + { + + hparams.mel_filters = get_vector(TN_MEL_FILTERS); + hparams.window = get_vector(TN_WINDOW); + + // Subsampling layers (conv1d) + for (int i : {0, 2, 3, 5, 6}) { + model.pre_encode_conv_X_w[i] = get_tensor(string_format(TN_CONV1D, i, "weight")); + model.pre_encode_conv_X_b[i] = get_tensor(string_format(TN_CONV1D, i, "bias")); + } + model.pre_encode_out_w = get_tensor(string_format(TN_PRE_ENCODE_OUT, "weight")); + model.pre_encode_out_b = get_tensor(string_format(TN_PRE_ENCODE_OUT, "bias")); + + // Projection layers + model.mm_norm_pre_w = get_tensor(string_format(TN_MM_NORM_PRE, "weight"), false); + model.mm_0_w = get_tensor(string_format(TN_MM_AUDIO_MLP, 1, "weight"), false); + model.mm_1_w = get_tensor(string_format(TN_MM_AUDIO_MLP, 2, "weight"), false); + + // Encoder layers + for (int il = 0; il < hparams.n_layer; ++il) { + auto & layer = model.layers[il]; + + // Attention (from shared above) + + // Relative position encoding + layer.linear_pos_w = get_tensor(string_format(TN_LINEAR_POS, prefix, il, "weight")); + layer.pos_bias_u = get_tensor(string_format(TN_POS_BIAS_U, prefix, il)); + layer.pos_bias_v = get_tensor(string_format(TN_POS_BIAS_V, prefix, il)); + + // Convolution module + layer.conv_pw1_w = get_tensor(string_format(TN_CONV_PW1, prefix, il, "weight")); + layer.conv_pw1_b = get_tensor(string_format(TN_CONV_PW1, prefix, il, "bias"), false); + layer.conv_dw_w = get_tensor(string_format(TN_CONV_DW, prefix, il, "weight")); + layer.conv_dw_b = get_tensor(string_format(TN_CONV_DW, prefix, il, "bias"), false); + layer.conv_norm_w = get_tensor(string_format(TN_CONV_NORM, prefix, il, "weight")); + layer.conv_norm_b = get_tensor(string_format(TN_CONV_NORM, prefix, il, "bias")); + layer.conv_norm_mean = get_tensor(string_format(TN_CONV_NORM_MEAN, prefix, il)); + layer.conv_norm_var = get_tensor(string_format(TN_CONV_NORM_VAR, prefix, il)); + layer.conv_pw2_w = get_tensor(string_format(TN_CONV_PW2, prefix, il, "weight")); + layer.conv_pw2_b = get_tensor(string_format(TN_CONV_PW2, prefix, il, "bias"), false); + + // Feed-forward networks + layer.ff_norm_w = get_tensor(string_format(TN_FFN_NORM, prefix, il, "weight")); + layer.ff_norm_b = get_tensor(string_format(TN_FFN_NORM, prefix, il, "bias")); + + layer.ff_norm_1_w = get_tensor(string_format(TN_FFN_NORM_1, prefix, il, "weight")); + layer.ff_norm_1_b = get_tensor(string_format(TN_FFN_NORM_1, prefix, il, "bias")); + layer.ff_up_1_w = get_tensor(string_format(TN_FFN_UP_1, prefix, il, "weight")); + layer.ff_up_1_b = get_tensor(string_format(TN_FFN_UP_1, prefix, il, "bias"), false); + layer.ff_down_1_w = get_tensor(string_format(TN_FFN_DOWN_1, prefix, il, "weight")); + layer.ff_down_1_b = get_tensor(string_format(TN_FFN_DOWN_1, prefix, il, "bias"), false); + + // Layer norms + layer.norm_conv_w = get_tensor(string_format(TN_NORM_CONV, prefix, il, "weight")); + layer.norm_conv_b = get_tensor(string_format(TN_NORM_CONV, prefix, il, "bias")); + } + + model.mm_model_mlp_1_w = get_tensor(string_format(TN_MVLM_PROJ_MLP, 0, "weight")); + model.mm_model_mlp_2_w = get_tensor(string_format(TN_MVLM_PROJ_MLP, 1, "weight")); + model.mm_model_mlp_3_w = get_tensor(string_format(TN_MVLM_PROJ_MLP, 3, "weight")); + } break; case PROJECTOR_TYPE_GRANITE_SPEECH: { model.inp_proj_w = get_tensor(string_format(TN_INP_PROJ, "weight")); @@ -3645,6 +3755,10 @@ int clip_n_output_tokens(const clip_ctx * ctx, const clip_image_f32 * img) { } n_patches = n; } break; + case PROJECTOR_TYPE_PARAKEET: + { + n_patches = (img->nx() + (params.subsampling_factor - 1)) / params.subsampling_factor; + } break; case PROJECTOR_TYPE_GEMMA4UA: { n_patches = img->nx(); // no downsampling: one token per raw waveform frame @@ -4558,6 +4672,88 @@ bool clip_image_batch_encode(clip_ctx * ctx, int n_threads, const clip_image_f32 } set_input_f32("pos_emb", pos_emb); } break; + case PROJECTOR_TYPE_PARAKEET: + { + GGML_ASSERT(imgs.entries.size() == 1); + struct ggml_tensor * attn_mask = ggml_graph_get_tensor(gf, "attn_mask"); + const int n_q = attn_mask->ne[1]; + const int n_k = attn_mask->ne[0]; + const int n_frames = imgs.entries.front().nx(); + const int n_tokens_real = (n_frames + hparams.subsampling_factor-1) / hparams.subsampling_factor; + const float mask_value = -1e30f; + + std::vector mask_data(n_q * n_k); + if (n_k == n_q) { + // full attention: mask keys that are padding + for (int q = 0; q < n_q; ++q) { + for (int k = 0; k < n_k; ++k) { + mask_data[q * n_k + k] = (k >= n_tokens_real) ? mask_value : 0.0f; + } + } + } else { + // local attention: mask keys outside the valid window + const int att_left = n_k / 2; + for (int q = 0; q < n_q; ++q) { + for (int k = 0; k < n_k; ++k) { + const int key = q - att_left + k; + mask_data[q * n_k + k] = (key >= 0 && key < n_tokens_real) ? 0.0f : mask_value; + } + } + } + set_input_f32(attn_mask->name, mask_data); + + // local attention skew mask: zeroes out the probs that were + // computed for keys outside the valid sliding window. + if (struct ggml_tensor * local_mask = ggml_graph_get_tensor(gf, "local_mask")) { + const int lm_k = local_mask->ne[0]; + const int lm_q = local_mask->ne[1]; + const int window_size = lm_k - lm_q + 1; + std::vector lm_data(lm_q * lm_k); + for (int q = 0; q < lm_q; ++q) { + for (int k = 0; k < lm_k; ++k) { + const int rel = k - q; + lm_data[q * lm_k + k] = (rel >= 0 && rel < window_size) ? 1.0f : 0.0f; + } + } + set_input_f32(local_mask->name, lm_data); + } + + // Generate rotation frequencies for relative positional encoding. + { + const int n_state = hparams.n_embd; + const int d_half = n_state / 2; + const float log_10000 = logf(10000.0f); + std::vector freqs(d_half); + for (int k = 0; k < d_half; ++k) { + freqs[k] = expf(-(float(k * 2) * log_10000 / float(n_state))); + } + set_input_f32("pos_freqs", freqs); + } + + // Generate relative positional distance values which scaled by + // the frequency to produce the angles for sin/cos. + { + // window_size is only known after graph construction since it depends on + // n_time from the conv output, so we read it back from the graph tensor. + struct ggml_tensor * rel_pos = ggml_graph_get_tensor(gf, "rel_positions"); + const int window_size = rel_pos->ne[1]; + std::vector pos(window_size); + // local attention: window is fixed at [att_left, att_right] + // full attention: window covers the full sequence, centered + if (ggml_graph_get_tensor(gf, "local_mask")) { + const int att_left = window_size / 2; + for (int t = 0; t < window_size; ++t) { + pos[t] = float(att_left - t); + } + } else { + const int n_time = (window_size + 1) / 2; + for (int t = 0; t < window_size; ++t) { + pos[t] = float(n_time - 1 - t); + } + } + set_input_f32(rel_pos->name, pos); + } + } break; case PROJECTOR_TYPE_GRANITE_SPEECH: { const int context_size = ctx->model.hparams.audio_chunk_size; @@ -4841,6 +5037,8 @@ int clip_n_mmproj_embd(const struct clip_ctx * ctx) { return ctx->model.mm_ffn_down_w->ne[1]; case PROJECTOR_TYPE_MIMO_AUDIO: return ctx->model.mm_2_w->ne[1]; + case PROJECTOR_TYPE_PARAKEET: + return ctx->model.mm_1_w->ne[1]; default: GGML_ABORT("Unknown projector type"); } diff --git a/tools/mtmd/models/models.h b/tools/mtmd/models/models.h index caed438ec513..e54366a086f4 100644 --- a/tools/mtmd/models/models.h +++ b/tools/mtmd/models/models.h @@ -222,6 +222,11 @@ struct clip_graph_kimik25 : clip_graph { ggml_tensor * resize_position_embeddings_3d(uint32_t interpolation_mode); }; +struct clip_graph_parakeet : clip_graph { + clip_graph_parakeet(clip_ctx * ctx, const clip_image_f32 & img) : clip_graph(ctx, img) {} + ggml_cgraph * build() override; +}; + struct clip_graph_exaone4_5 : clip_graph { clip_graph_exaone4_5(clip_ctx * ctx, const clip_image_f32 & img) : clip_graph(ctx, img) {} ggml_cgraph * build() override; diff --git a/tools/mtmd/models/parakeet.cpp b/tools/mtmd/models/parakeet.cpp new file mode 100644 index 000000000000..8be141d93b37 --- /dev/null +++ b/tools/mtmd/models/parakeet.cpp @@ -0,0 +1,421 @@ +#include "models.h" + +static constexpr int PARAKEET_LOCAL_ATTN_THRESHOLD = 8192; +static constexpr int PARAKEET_LOCAL_ATTN_WINDOW = 128; + +// conv subsampling + conformer encoder +ggml_cgraph * clip_graph_parakeet::build() { + + // Conv subsampling + ggml_tensor * inp = build_inp_raw(1); + inp = ggml_cont(ctx0, ggml_transpose(ctx0, inp)); + + // [freq, time, channels, batch] + ggml_tensor * cur = ggml_conv_2d(ctx0, model.pre_encode_conv_X_w[0], inp, 2, 2, 1, 1, 1, 1); + cur = ggml_add(ctx0, cur, model.pre_encode_conv_X_b[0]); + cb(cur, "pre_conv_0", -1); + + cur = ggml_relu(ctx0, cur); + cb(cur, "pre_conv_0_relu", -1); + + // [freq, time, channels, batch] + cur = ggml_conv_2d_dw_direct(ctx0, model.pre_encode_conv_X_w[2], cur, 2, 2, 1, 1, 1, 1); + cur = ggml_add(ctx0, cur, model.pre_encode_conv_X_b[2]); + cb(cur, "pre_conv_2", -1); + + // [freq, time, channels, batch] + cur = ggml_conv_2d(ctx0, model.pre_encode_conv_X_w[3], cur, 1, 1, 0, 0, 1, 1); + cur = ggml_add(ctx0, cur, model.pre_encode_conv_X_b[3]); + cb(cur, "pre_conv_3", -1); + + cur = ggml_relu(ctx0, cur); + cb(cur, "pre_conv_3_relu", -1); + + // [freq, time, channels, batch] + cur = ggml_conv_2d_dw_direct(ctx0, model.pre_encode_conv_X_w[5], cur, 2, 2, 1, 1, 1, 1); + cb(cur, "pre_conv_5_direct", -1); + cur = ggml_add(ctx0, cur, model.pre_encode_conv_X_b[5]); + cb(cur, "pre_conv_5", -1); + + // [freq, time, channels, batch] + cur = ggml_conv_2d(ctx0, model.pre_encode_conv_X_w[6], cur, 1, 1, 0, 0, 1, 1); + cur = ggml_add(ctx0, cur, model.pre_encode_conv_X_b[6]); + cb(cur, "pre_conv_6", -1); + + cur = ggml_relu(ctx0, cur); + cb(cur, "pre_conv_6_relu", -1); + + // [freq, time, chan] + cur = ggml_permute(ctx0, cur, 0, 2, 1, 3); + // [freq, chan, time] + cur = ggml_cont(ctx0, cur); + + const int n_freq = cur->ne[0]; + const int n_chan = cur->ne[1]; + const int n_frames = cur->ne[2]; + + // [freq, time, chan, batch] -> [(freq * chan), time] + cur = ggml_reshape_2d(ctx0, cur, n_freq * n_chan, n_frames); + + cur = build_mm(model.pre_encode_out_w, cur); + cur = ggml_add(ctx0, cur, model.pre_encode_out_b); + + ggml_set_name(cur, "pre_enc_out"); + + // Encoder + + const auto & hparams = model.hparams; + const int n_layer = hparams.n_layer; + const int n_state = hparams.n_embd; + const float fc_factor = 0.5f; + + const int n_time = cur->ne[1]; + const bool local_attn = n_time > PARAKEET_LOCAL_ATTN_THRESHOLD; + const int att_left = local_attn ? PARAKEET_LOCAL_ATTN_WINDOW : n_time - 1; + const int att_right = local_attn ? PARAKEET_LOCAL_ATTN_WINDOW : n_time - 1; + const int window_size = local_attn ? att_left + att_right + 1 : 2 * n_time - 1; + const int d_half = n_state / 2; + const int mask_dim = local_attn ? window_size : n_time; + + // mask [key, n_time] + struct ggml_tensor * attn_mask = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, mask_dim, n_time); + ggml_set_name(attn_mask, "attn_mask"); + ggml_set_input(attn_mask); + + struct ggml_tensor * local_mask = nullptr; + if (local_attn) { + const int chunk = att_left + att_right; + local_mask = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, chunk + window_size - 1, chunk); + ggml_set_name(local_mask, "local_mask"); + ggml_set_input(local_mask); + } + + struct ggml_tensor * pos_freqs = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, d_half); + ggml_set_name(pos_freqs, "pos_freqs"); + ggml_set_input(pos_freqs); + + struct ggml_tensor * rel_positions = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, 1, window_size); + ggml_set_name(rel_positions, "rel_positions"); + ggml_set_input(rel_positions); + + struct ggml_tensor * freqs = ggml_repeat_4d(ctx0, pos_freqs, d_half, window_size, 1, 1); + struct ggml_tensor * theta = ggml_mul(ctx0, freqs, rel_positions); + + struct ggml_tensor * sin = ggml_reshape_3d(ctx0, ggml_sin(ctx0, theta), 1, d_half, window_size); + struct ggml_tensor * cos = ggml_reshape_3d(ctx0, ggml_cos(ctx0, theta), 1, d_half, window_size); + struct ggml_tensor * pos_emb = ggml_reshape_2d(ctx0, ggml_cont(ctx0, ggml_concat(ctx0, sin, cos, 0)), n_state, window_size); + ggml_set_name(pos_emb, "pos_emb"); + + for (int il = 0; il < n_layer; ++il) { + const auto & layer = model.layers[il]; + // FFN1 + { + struct ggml_tensor * residual = cur; + ggml_format_name(cur, "enc_%d_res", il); + + // norm + cur = ggml_norm(ctx0, cur, hparams.eps); + cur = ggml_add(ctx0, ggml_mul(ctx0, cur, layer.ff_norm_w), layer.ff_norm_b); + ggml_format_name(cur, "enc_%d_ffn_norm_1", il); + + cur = build_ffn(cur, layer.ff_up_w, nullptr, nullptr, nullptr, layer.ff_down_w, nullptr, FFN_SILU, il); + ggml_format_name(cur, "enc_%d_ffn_1", il); + + cur = ggml_add(ctx0, residual, ggml_scale(ctx0, cur, fc_factor)); + ggml_format_name(cur, "enc_%d_res_ffn", il); + } + + // self attention block using relative positional encoding from model.position_embedding. + { + // [feat, time_frames, 1, 1] + struct ggml_tensor * residual = cur; + + cur = ggml_norm(ctx0, cur, hparams.eps); + cur = ggml_add(ctx0, ggml_mul(ctx0, cur, layer.ln_1_w), layer.ln_1_b); + ggml_format_name(cur, "enc_%d_attn_norm", il); + + const int n_head = hparams.n_head; + const int d_head = n_state / n_head; + + // [feat, time_frames, 1, 1] + struct ggml_tensor * Q_cur = build_mm(layer.q_w, cur); + struct ggml_tensor * K_cur = build_mm(layer.k_w, cur); + struct ggml_tensor * V_cur = build_mm(layer.v_w, cur); + + // [d_head, n_heads, n_time, 1] + Q_cur = ggml_reshape_3d(ctx0, Q_cur, d_head, n_head, n_time); + K_cur = ggml_reshape_3d(ctx0, K_cur, d_head, n_head, n_time); + V_cur = ggml_reshape_3d(ctx0, V_cur, d_head, n_head, n_time); + + // [n_state, window_size] + struct ggml_tensor * pos = build_mm(layer.linear_pos_w, pos_emb); + // [feat, head, window_size, 1] + pos = ggml_reshape_3d(ctx0, pos, d_head, n_head, pos_emb->ne[1]); + // [feat, window_size, head, 1] + pos = ggml_cont(ctx0, ggml_permute(ctx0, pos, 0, 2, 1, 3)); + ggml_format_name(pos, "enc_%d_attn_pos", il); + + if (local_attn) { + const int chunk = att_left + att_right; + const int n_group = (n_time + chunk - 1) / chunk; + const int n_time_padded = n_group * chunk; + const int n_kv_chunk = chunk + window_size - 1; + const int n_kv_dense = n_kv_chunk * n_group; + const bool need_padding = n_time_padded > n_time; + + Q_cur = ggml_cont(ctx0, ggml_permute(ctx0, Q_cur, 0, 2, 1, 3)); + K_cur = ggml_cont(ctx0, ggml_permute(ctx0, K_cur, 0, 2, 1, 3)); + V_cur = ggml_cont(ctx0, ggml_permute(ctx0, V_cur, 0, 2, 1, 3)); + + // content bias + struct ggml_tensor * bias_u = ggml_reshape_3d(ctx0, layer.pos_bias_u, d_head, 1, n_head); + struct ggml_tensor * Q_u = ggml_add(ctx0, Q_cur, bias_u); + + // position bias + struct ggml_tensor * bias_v = ggml_reshape_3d(ctx0, layer.pos_bias_v, d_head, 1, n_head); + struct ggml_tensor * Q_v = ggml_add(ctx0, Q_cur, bias_v); + + // right pad the time dimension + struct ggml_tensor * Q_u_padded = need_padding ? + ggml_pad_ext(ctx0, Q_u, 0, 0, 0, n_time_padded - n_time, 0, 0, 0, 0) : Q_u; + Q_u_padded = ggml_reshape_4d(ctx0, Q_u_padded, d_head, chunk, n_group, n_head); + + // pad front and back for the first and last time frames + struct ggml_tensor * K_padded = ggml_pad_ext(ctx0, K_cur, 0, 0, att_left, att_right, 0, 0, 0, 0); + if (n_kv_dense > K_padded->ne[1]) { + K_padded = ggml_pad_ext(ctx0, K_padded, 0, 0, 0, n_kv_dense - K_padded->ne[1], 0, 0, 0, 0); + } + + // sliding window view: each group spans n_kv_chunk keys but steps by chunk + struct ggml_tensor * K_chunk = ggml_view_4d(ctx0, K_padded, + d_head, n_kv_chunk, n_group, n_head, + K_padded->nb[1], + (size_t) chunk * K_padded->nb[1], + K_padded->nb[2], + 0); + K_chunk = ggml_cont(ctx0, K_chunk); + + struct ggml_tensor * content_scores = ggml_mul_mat(ctx0, K_chunk, Q_u_padded); + + // trim the dense output down to window_size scores per query + content_scores = ggml_view_4d(ctx0, content_scores, + window_size, chunk, n_group, n_head, + (size_t) (chunk + window_size) * content_scores->nb[0], + content_scores->nb[2], + content_scores->nb[3], + 0); + content_scores = ggml_cont(ctx0, content_scores); + + // ungroup: [window_size, n_time_padded, n_head] + content_scores = ggml_reshape_3d(ctx0, content_scores, window_size, n_time_padded, n_head); + if (need_padding) { + content_scores = ggml_view_3d(ctx0, content_scores, + window_size, n_time, n_head, + content_scores->nb[1], + content_scores->nb[2], + 0); + } + + // Q_v: [d_head, time, head] + Q_v = ggml_cont(ctx0, ggml_permute(ctx0, Q_v, 0, 2, 1, 3)); + struct ggml_tensor * rel_pos_scores = ggml_mul_mat(ctx0, pos, Q_v); + + struct ggml_tensor * attn_scores = ggml_add(ctx0, content_scores, rel_pos_scores); + attn_scores = ggml_soft_max_ext(ctx0, attn_scores, attn_mask, 1.0f / std::sqrt(d_head), 0.0f); + ggml_format_name(attn_scores, "enc_%d_attn_probs", il); + + // expand probs back to n_kv_chunk width for the V matmul + struct ggml_tensor * probs_padded = need_padding ? + ggml_pad_ext(ctx0, attn_scores, 0, 0, 0, n_time_padded - n_time, 0, 0, 0, 0) : attn_scores; + + probs_padded = ggml_reshape_4d(ctx0, probs_padded, window_size, chunk, n_group, n_head); + probs_padded = ggml_pad_ext(ctx0, probs_padded, 0, chunk, 0, 0, 0, 0, 0, 0); + probs_padded = ggml_view_4d(ctx0, probs_padded, + n_kv_chunk, chunk, n_group, n_head, + (size_t) n_kv_chunk * probs_padded->nb[0], + probs_padded->nb[2], + probs_padded->nb[3], + 0); + probs_padded = ggml_cont(ctx0, probs_padded); + probs_padded = ggml_mul(ctx0, probs_padded, local_mask); + + struct ggml_tensor * V_padded = ggml_pad_ext(ctx0, V_cur, 0, 0, att_left, att_right, 0, 0, 0, 0); + if (n_kv_dense > V_padded->ne[1]) { + V_padded = ggml_pad_ext(ctx0, V_padded, 0, 0, 0, n_kv_dense - V_padded->ne[1], 0, 0, 0, 0); + } + V_padded = ggml_cont(ctx0, ggml_transpose(ctx0, V_padded)); + + struct ggml_tensor * V_chunk = ggml_view_4d(ctx0, V_padded, + n_kv_chunk, d_head, n_group, n_head, + V_padded->nb[1], + (size_t) chunk * V_padded->nb[0], + V_padded->nb[2], + 0); + V_chunk = ggml_cont(ctx0, V_chunk); + + cur = ggml_mul_mat(ctx0, V_chunk, probs_padded); + cur = ggml_reshape_3d(ctx0, cur, d_head, n_time_padded, n_head); + if (need_padding) { + cur = ggml_view_3d(ctx0, cur, d_head, n_time, n_head, cur->nb[1], cur->nb[2], 0); + } + cur = ggml_cont(ctx0, ggml_permute(ctx0, cur, 0, 2, 1, 3)); + cur = ggml_reshape_2d(ctx0, cur, n_state, n_time); + cur = build_mm(layer.o_w, cur); + } else { + // full attention + struct ggml_tensor * Q_u = ggml_add(ctx0, Q_cur, layer.pos_bias_u); + ggml_format_name(Q_u, "enc_%d_attn_q_u", il); + + struct ggml_tensor * K_prep = ggml_permute(ctx0, K_cur, 0, 2, 1, 3); + struct ggml_tensor * Q_prep = ggml_permute(ctx0, Q_u, 0, 2, 1, 3); + struct ggml_tensor * content_scores = ggml_mul_mat(ctx0, K_prep, Q_prep); + ggml_format_name(content_scores, "enc_%d_attn_content_scores", il); + + struct ggml_tensor * Q_v = ggml_add(ctx0, Q_cur, layer.pos_bias_v); + ggml_format_name(Q_v, "enc_%d_attn_q_v", il); + + Q_v = ggml_permute(ctx0, Q_v, 0, 2, 1, 3); + Q_v = ggml_cont(ctx0, Q_v); + ggml_format_name(Q_v, "enc_%d_attn_q_v_perm", il); + + struct ggml_tensor * rel_pos_scores = ggml_mul_mat(ctx0, pos, Q_v); + ggml_format_name(rel_pos_scores, "enc_%d_attn_rel_pos", il); + + // Relative positional shift + { + const auto pos_window = rel_pos_scores->ne[0]; + const auto n_frame = rel_pos_scores->ne[1]; + const auto n_head = rel_pos_scores->ne[2]; + + rel_pos_scores = ggml_pad(ctx0, rel_pos_scores, 1, 0, 0, 0); + rel_pos_scores = ggml_roll(ctx0, rel_pos_scores, 1, 0, 0, 0); + + rel_pos_scores = ggml_reshape_3d(ctx0, rel_pos_scores, n_frame, pos_window + 1, n_head); + rel_pos_scores = ggml_cont(ctx0, rel_pos_scores); + ggml_format_name(rel_pos_scores, "enc_%d_attn_rel_pos_reshaped", il); + + int center = pos_window / 2; + size_t offset = rel_pos_scores->nb[0] * (center+1); + + rel_pos_scores = ggml_view_3d(ctx0, rel_pos_scores, + n_frame, pos_window, n_head, + (pos_window) * 4, + rel_pos_scores->nb[2], + offset); + rel_pos_scores = ggml_cont(ctx0, rel_pos_scores); + ggml_format_name(rel_pos_scores, "enc_%d_attn_rel_pos_shifted", il); + + rel_pos_scores = ggml_view_3d(ctx0, rel_pos_scores, + content_scores->ne[0], + content_scores->ne[1], + rel_pos_scores->ne[2], + rel_pos_scores->nb[1], + rel_pos_scores->nb[2], + 0); + rel_pos_scores = ggml_cont(ctx0, rel_pos_scores); + ggml_format_name(rel_pos_scores, "enc_%d_attn_rel_pos_shifted_view", il); + } + + struct ggml_tensor * attn_scores = ggml_add(ctx0, content_scores, rel_pos_scores); + ggml_format_name(attn_scores, "enc_%d_attn_scores", il); + attn_scores = ggml_scale(ctx0, attn_scores, 1.0f / std::sqrt(d_head)); + attn_scores = ggml_add(ctx0, attn_scores, attn_mask); + ggml_format_name(attn_scores, "enc_%d_attn_scores_scaled", il); + + struct ggml_tensor * probs = ggml_soft_max(ctx0, attn_scores); + ggml_format_name(probs, "enc_%d_attn_probs", il); + + V_cur = ggml_cont(ctx0, ggml_permute(ctx0, V_cur, 1, 2, 0, 3)); + ggml_format_name(V_cur, "enc_%d_attn_v_cur", il); + cur = ggml_mul_mat(ctx0, probs, V_cur); + ggml_format_name(cur, "enc_%d_attn_inp", il); + + cur = ggml_permute(ctx0, cur, 2, 0, 1, 3); + cur = ggml_cont_2d(ctx0, cur, n_state, n_time); + cur = build_mm(layer.o_w, cur); + } + ggml_format_name(cur, "enc_%d_attn_out", il); + + cur = ggml_add(ctx0, residual, cur); + ggml_format_name(cur, "enc_%d_attn_res", il); + } + + // Convolution + { + struct ggml_tensor * residual = cur; + ggml_format_name(cur, "enc_%d_residual_conv", il); + + cur = ggml_norm(ctx0, cur, hparams.eps); + cur = ggml_add(ctx0, ggml_mul(ctx0, cur, layer.norm_conv_w), layer.norm_conv_b); + ggml_format_name(cur, "enc_%d_norm_conv", il); + + // pointwise 1d convolution: + cur = build_mm(layer.conv_pw1_w, cur); + ggml_format_name(cur, "enc_%d_conv_pw1", il); + + { + int64_t d = cur->ne[0] / 2; + struct ggml_tensor * signal = ggml_view_2d(ctx0, cur, d, cur->ne[1], cur->nb[1], 0); + struct ggml_tensor * gate = ggml_view_2d(ctx0, cur, d, cur->ne[1], cur->nb[1], d * cur->nb[0]); + + cur = ggml_mul(ctx0, signal, ggml_sigmoid(ctx0, gate)); + ggml_format_name(cur, "enc_%d_conv_glu", il); + } + + cur = ggml_cont(ctx0, ggml_transpose(ctx0, cur)); + + // use ggml_ssm_conv for f32 precision + const int dw_pad = (hparams.audio_conv_kernel_size - 1) / 2; + cur = ggml_pad(ctx0, cur, dw_pad, 0, 0, 0); + cur = ggml_roll(ctx0, cur, dw_pad, 0, 0, 0); + cur = ggml_pad(ctx0, cur, dw_pad, 0, 0, 0); + ggml_format_name(cur, "enc_%d_conv_dw_pad", il); + + cur = ggml_ssm_conv(ctx0, cur, layer.conv_dw_w); + ggml_format_name(cur, "enc_%d_conv_1d_dw", il); + + cur = ggml_sub(ctx0, cur, layer.conv_norm_mean); + struct ggml_tensor * std = ggml_sqrt(ctx0, layer.conv_norm_var); + cur = ggml_div(ctx0, cur, std); + cur = ggml_add(ctx0, ggml_mul(ctx0, cur, layer.conv_norm_w), layer.conv_norm_b); + ggml_format_name(cur, "enc_%d_conv_bn", il); + + cur = ggml_silu(ctx0, cur); + ggml_format_name(cur, "enc_%d_conv_silu", il); + + cur = build_mm(layer.conv_pw2_w, cur); + ggml_format_name(cur, "enc_%d_conv_pw2", il); + + cur = ggml_add(ctx0, residual, cur); + ggml_format_name(cur, "enc_%d_conv_res", il); + } + + // FFN2 + { + struct ggml_tensor * residual = cur; + cur = ggml_norm(ctx0, cur, hparams.eps); + cur = ggml_add(ctx0, ggml_mul(ctx0, cur, layer.ff_norm_1_w), layer.ff_norm_1_b); + ggml_format_name(cur, "enc_%d_ffn_norm_2", il); + + cur = build_ffn(cur, layer.ff_up_1_w, nullptr, nullptr, nullptr, layer.ff_down_1_w, nullptr, FFN_SILU, il); + cur = ggml_add(ctx0, residual, ggml_scale(ctx0, cur, 0.5)); + ggml_format_name(cur, "enc_%d_ffn_res", il); + } + + cur = ggml_norm(ctx0, cur, hparams.eps); + cur = ggml_add(ctx0, ggml_mul(ctx0, cur, layer.ln_2_w), layer.ln_2_b); + } + + cb(cur, "encoder_out", -1); + + cur = ggml_rms_norm(ctx0, cur, 1e-6); + cur = ggml_mul(ctx0, cur, model.mm_norm_pre_w); + cb(cur, "sound_projection.norm", -1); + + cur = build_ffn(cur, model.mm_0_w, model.mm_0_b, nullptr, nullptr, model.mm_1_w, model.mm_1_b, FFN_RELU_SQR, -1); + cb(cur, "projected", -1); + + ggml_build_forward_expand(gf, cur); + + return gf; +} diff --git a/tools/mtmd/mtmd-audio.cpp b/tools/mtmd/mtmd-audio.cpp index ed68951c0151..fea03557d05c 100644 --- a/tools/mtmd/mtmd-audio.cpp +++ b/tools/mtmd/mtmd-audio.cpp @@ -1022,6 +1022,209 @@ bool mtmd_audio_preprocessor_gemma4a::preprocess(const float * s } // +// mtmd_audio_preprocessor_parakeet implementation +// + +void mtmd_audio_preprocessor_parakeet::worker_thread( + int ith, + const float * window_func, + int window_size, + const std::vector & samples, + int n_samples, + int frame_size, + int frame_step, + int n_threads, + int n_fft_bins, + const mtmd_audio_cache & cache, + mtmd_audio_mel & mel) { + std::vector fft_in(frame_size * 2, 0.0); + std::vector fft_out(frame_size * 2 * 2 * 2); + + int n_fb = n_fft_bins; + int i = ith; + + GGML_ASSERT(n_fb == 1 + (frame_size / 2)); + + const double eps = 5.960464477539063e-08; + + for (; i < std::min(n_samples / frame_step + 1, (int) mel.n_len); i += n_threads) { + const int offset = i * frame_step; + const int window_pad_left = (frame_size - window_size) / 2; + + // Zero-pad left. + std::fill(fft_in.begin(), fft_in.begin() + window_pad_left, 0.0f); + + // Apply windowed samples in the center. + const int n_to_process = std::min({window_size, n_samples - offset}); + for (int j = 0; j < n_to_process; j++) { + fft_in[window_pad_left + j] = window_func[j] * samples[offset + window_pad_left + j]; + } + + // Zero-pad right. + std::fill(fft_in.begin() + window_pad_left + n_to_process, fft_in.begin() + frame_size, 0.0f); + + // FFT. + fft(cache, fft_in.data(), frame_size, fft_out.data()); + + // Calculate modulus^2 of complex numbers. + for (int j = 0; j < n_fb; j++) { + fft_out[j] = (fft_out[2 * j + 0] * fft_out[2 * j + 0] + fft_out[2 * j + 1] * fft_out[2 * j + 1]); + } + + // mel spectrogram. + for (int j = 0; j < mel.n_mel; j++) { + double sum = 0.0; + int k = 0; + for (k = 0; k < n_fb - 3; k += 4) { + sum += + fft_out[k + 0] * cache.filters.data[j * n_fb + k + 0] + + fft_out[k + 1] * cache.filters.data[j * n_fb + k + 1] + + fft_out[k + 2] * cache.filters.data[j * n_fb + k + 2] + + fft_out[k + 3] * cache.filters.data[j * n_fb + k + 3]; + } + for (; k < n_fb; k++) { + sum += fft_out[k] * cache.filters.data[j * n_fb + k]; + } + mel.data[j * mel.n_len + i] = std::log(sum + eps); + } + } + + // Otherwise fft_out are all zero. + const double empty_sum = std::log(eps); + for (; i < mel.n_len; i += n_threads) { + for (int j = 0; j < mel.n_mel; j++) { + mel.data[j * mel.n_len + i] = empty_sum; + } + } +} + +void mtmd_audio_preprocessor_parakeet::initialize() { + cache.fill_sin_cos_table(hparams.audio_n_fft); + + const size_t n_fft = hparams.audio_n_fft / 2 + 1; + GGML_ASSERT(hparams.mel_filters.size() == (size_t)hparams.n_mel_bins * n_fft); + cache.filters.n_mel = hparams.n_mel_bins; + cache.filters.n_fft = n_fft; + cache.filters.data = hparams.mel_filters; + + GGML_ASSERT(hparams.window.size() == (size_t)hparams.audio_window_len); + GGML_ASSERT(hparams.window.size() <= (size_t) hparams.audio_n_fft); + cache.hann_window = hparams.window; +} + +bool mtmd_audio_preprocessor_parakeet::preprocess(const float * samples, + size_t n_samples_in, + std::vector & output) { + if (n_samples_in == 0) { + return false; + } + + filter_params params; + params.n_mel = hparams.n_mel_bins; + params.n_fft_bins = 1 + (hparams.audio_n_fft / 2); + params.hann_window_size = hparams.audio_window_len; + params.hop_length = hparams.audio_hop_len; + params.sample_rate = hparams.audio_sample_rate; + + GGML_ASSERT(!cache.sin_vals.empty()); + GGML_ASSERT(!cache.cos_vals.empty()); + GGML_ASSERT(!cache.filters.data.empty()); + + const float * window_func = cache.hann_window.data(); + const int window_size = params.hann_window_size; + const int frame_size = (params.n_fft_bins - 1) * 2; + const int frame_step = params.hop_length; + + // Apply preemphasis filter (high-pass): x[i] = x[i] - 0.97 * x[i-1] + std::vector samples_preprocessed(samples, samples + n_samples_in); + { + const float preemph = 0.97f; + for (int i = n_samples_in - 1; i > 0; i--) { + samples_preprocessed[i] = samples_preprocessed[i] - preemph * samples_preprocessed[i - 1]; + } + } + + // Parakeet uses centered constant padding + const size_t pad = (size_t)(frame_size / 2); + std::vector samples_padded(n_samples_in + 2 * pad, 0.0f); + std::copy(samples_preprocessed.begin(), samples_preprocessed.end(), samples_padded.begin() + pad); + + mtmd_audio_mel out_full; + out_full.n_mel = params.n_mel; + out_full.n_len = (samples_padded.size() - frame_size) / frame_step + 1; + out_full.n_len_org = out_full.n_len; + out_full.data.resize(out_full.n_mel * out_full.n_len); + + const int n_threads = 4; + std::vector workers(n_threads - 1); + for (int iw = 0; iw < n_threads - 1; ++iw) { + workers[iw] = std::thread( + worker_thread, iw + 1, + window_func, + window_size, + std::cref(samples_padded), + samples_padded.size(), + frame_size, + frame_step, + n_threads, + params.n_fft_bins, + std::cref(cache), + std::ref(out_full) + ); + } + + worker_thread(0, + window_func, + window_size, + samples_padded, + samples_padded.size(), + frame_size, + frame_step, + n_threads, + params.n_fft_bins, + cache, + out_full); + + for (int iw = 0; iw < n_threads - 1; ++iw) { + workers[iw].join(); + } + + // Per-feature normalization (only on valid frames) + { + const double eps = 1e-5; + int valid_frames = n_samples_in / frame_step; + + for (int j = 0; j < out_full.n_mel; j++) { + double sum = 0.0; + double sq_diff_sum = 0.0; + + // Calculate Mean ONLY on valid audio frames + for (int i = 0; i < valid_frames; i++) { + sum += (double)out_full.data[j * out_full.n_len + i]; + } + double mean = sum / valid_frames; + + // Calculate Variance ONLY on valid audio frames + for (int i = 0; i < valid_frames; i++) { + double diff = (double)out_full.data[j * out_full.n_len + i] - mean; + sq_diff_sum += diff * diff; + } + + double std_dev = std::sqrt(sq_diff_sum / (valid_frames - 1.0)); + double denominator = std_dev + eps; + + // Apply to ALL frames (including the padded ones) + for (int i = 0; i < out_full.n_len; i++) { + out_full.data[j * out_full.n_len + i] = (float)((out_full.data[j * out_full.n_len + i] - mean) / denominator); + } + } + } + + output.push_back(std::move(out_full)); + return true; +} + + // mtmd_audio_preprocessor_gemma4ua // diff --git a/tools/mtmd/mtmd-audio.h b/tools/mtmd/mtmd-audio.h index d8ec72b9d54e..f65f282d96e2 100644 --- a/tools/mtmd/mtmd-audio.h +++ b/tools/mtmd/mtmd-audio.h @@ -120,6 +120,21 @@ struct mtmd_audio_preprocessor_mimo_audio : mtmd_audio_preprocessor { mtmd_audio_cache cache; }; +struct mtmd_audio_preprocessor_parakeet : mtmd_audio_preprocessor { + mtmd_audio_preprocessor_parakeet(clip_ctx * ctx) : mtmd_audio_preprocessor(ctx) { } + void initialize() override; + bool preprocess(const float * samples, size_t n_samples, std::vector & output) override; + + private: + mtmd_audio_cache cache; + + static void worker_thread(int ith, const float * window_func, int window_size, + const std::vector & samples, int n_samples, + int frame_size, int frame_step, int n_threads, + int n_fft_bins, + const mtmd_audio_cache & cache, mtmd_audio_mel & mel); +}; + // // streaming ISTFT - converts spectrogram frames back to audio one frame at a time // diff --git a/tools/mtmd/mtmd.cpp b/tools/mtmd/mtmd.cpp index 6e61cf3e520b..93ca8cbcf8ae 100644 --- a/tools/mtmd/mtmd.cpp +++ b/tools/mtmd/mtmd.cpp @@ -724,6 +724,10 @@ struct mtmd_context { aud_end = ""; audio_preproc = std::make_unique(ctx_a); } break; + case PROJECTOR_TYPE_PARAKEET: + { + audio_preproc = std::make_unique(ctx_a); + } break; case PROJECTOR_TYPE_GEMMA4UA: { aud_beg = "<|audio>"; From 8190848bb36c7df4251db4352bd81bc07d0a4385 Mon Sep 17 00:00:00 2001 From: Hongqiang Wang Date: Tue, 28 Jul 2026 11:04:42 -0700 Subject: [PATCH 043/190] opencl: skip the Adreno KQ/KQV image kernels for multi-stream batches (#26189) The Adreno KQ/KQV image1d kernels (ggml_cl_mul_mat_kq_kqv_adreno) ignore dim 3 entirely: the sub-buffer covers only nb02*ne02 bytes and the kernel receives no ne03/ne13/nb03/nb13 arguments. With the unified KV cache, multi-sequence batches (e.g. llama-perplexity with its default -b 2048, n_seq=4, or a multi-slot llama-server) present KQ/KQV as 4D tensors with ne3 = n_stream, so every stream past the first reads the first stream's K/V and produces garbage. Flash attention masks the bug where it is enabled; devices where FA is declined (e.g. Adreno 740) hit it with default settings. Route ne03/ne13 > 1 to the general path, which handles dim 3, and honor view_offs when creating the sub-buffers (currently always 0 for tensors reaching this function, but the function would silently misread any future view). Llama-3.2-1B-Instruct Q4_0, wiki.test.raw, 8 chunks, -ngl 99: - Adreno 740, default: PPL 1817.64 -> 15.61 - Adreno 740, -fa 0: PPL 1941.64 -> 15.61 - Adreno 840, -fa 0: PPL 1943.90 -> 15.50 - single-stream (-b 512) results unchanged (15.6090) - test-backend-ops -o MUL_MAT on 740: identical before/after (909 OK, 12 pre-existing q6_K failures) --- ggml/src/ggml-opencl/ggml-opencl.cpp | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/ggml/src/ggml-opencl/ggml-opencl.cpp b/ggml/src/ggml-opencl/ggml-opencl.cpp index a05d18ee30af..d07b8fe41a31 100644 --- a/ggml/src/ggml-opencl/ggml-opencl.cpp +++ b/ggml/src/ggml-opencl/ggml-opencl.cpp @@ -15675,7 +15675,7 @@ static void ggml_cl_mul_mat_kq_kqv_adreno(ggml_backend_t backend, const ggml_ten // <--------------------------------------------> // extra0 = src0->view_src ? (ggml_tensor_extra_cl *)src0->view_src->extra : (ggml_tensor_extra_cl *)src0->extra; - region.origin = (extra0->offset); + region.origin = (extra0->offset + src0->view_offs); if (nb01 > nb02) { // KQ region.size = nb01 * ne01; @@ -15691,7 +15691,7 @@ static void ggml_cl_mul_mat_kq_kqv_adreno(ggml_backend_t backend, const ggml_ten // create sub-buffer for B // <--------------------------------------------> // - region.origin = (extra1->offset); + region.origin = (extra1->offset + src1->view_offs); region.size = nb10 * ne10 * ne11 * ne12; B_sub_buffer = clCreateSubBuffer((extra1->data_device), 0, CL_BUFFER_CREATE_TYPE_REGION, ®ion, &status); CL_CHECK(status); @@ -15712,7 +15712,7 @@ static void ggml_cl_mul_mat_kq_kqv_adreno(ggml_backend_t backend, const ggml_ten // create sub-buffer for output C // <--------------------------------------------> // - region.origin = (extrad->offset); + region.origin = (extrad->offset + dst->view_offs); region.size = ne0 * ne1 * dst->ne[2] * dst->nb[0]; // size of C in bytes D_sub_buffer = clCreateSubBuffer((extrad->data_device), 0, CL_BUFFER_CREATE_TYPE_REGION, ®ion, &status); CL_CHECK(status); @@ -18591,6 +18591,8 @@ static void ggml_cl_mul_mat(ggml_backend_t backend, const ggml_tensor * src0, co #ifdef GGML_OPENCL_USE_ADRENO_KERNELS if(src0t == GGML_TYPE_F16 && src1t == GGML_TYPE_F32){ if (ne01 >= 64 && ne1 >= 32 && ne00 >= 16 && (ne12 % ne02) == 0 && + // the KQ/KQV image kernels do not handle dim 3 (multi-stream batches) + ne03 == 1 && ne13 == 1 && // dst is wrapped with image1d_buffer, the size limit applies, also src0 (ne0 * ne1 * dst->ne[2] * dst->nb[0] / 4 <= backend_ctx->image_max_buffer_size)) { // For KQ From bc71c24c9da1e7ba6f5993da0e61476155720b06 Mon Sep 17 00:00:00 2001 From: Reese Levine Date: Tue, 28 Jul 2026 11:13:06 -0700 Subject: [PATCH 044/190] ggml-webgpu: Fix some binding alias issues to support all archs, fix recurrent-state-rollback test (#25931) * Add overlap glu variant to support all archs, fix recurrent-state-rollback test * format * Fix all arch overlapped ranges * format * diagnose bus error on apple ci * More testing * more testing * more targeted testing * Fix bug in alignment for > 4gb buffer offsets * Fix bug in view offsets * Try avoiding multi_buffers * not fixed yet, more logging :( * Handle edge case in set_rows * Try looking at view source * Skip deepseek32 for now and clean up trace infrastructure * simplify skipping * last cleanup * actually final cleanup * update handling of overlap * format * try skipping other failing model --- .../ggml-webgpu/ggml-webgpu-shader-lib.hpp | 83 ++++--- ggml/src/ggml-webgpu/ggml-webgpu.cpp | 230 ++++++++++++++---- ggml/src/ggml-webgpu/wgsl-shaders/glu.wgsl | 17 +- .../ggml-webgpu/wgsl-shaders/ssm_scan.wgsl | 68 +++++- tests/test-llama-archs.cpp | 7 +- 5 files changed, 306 insertions(+), 99 deletions(-) diff --git a/ggml/src/ggml-webgpu/ggml-webgpu-shader-lib.hpp b/ggml/src/ggml-webgpu/ggml-webgpu-shader-lib.hpp index bed9265b8ab7..99d775c57624 100644 --- a/ggml/src/ggml-webgpu/ggml-webgpu-shader-lib.hpp +++ b/ggml/src/ggml-webgpu/ggml-webgpu-shader-lib.hpp @@ -73,11 +73,6 @@ inline bool ggml_webgpu_tensor_equal(const ggml_tensor * a, const ggml_tensor * return a->buffer == b->buffer && ggml_webgpu_tensor_addr(a) == ggml_webgpu_tensor_addr(b); } -inline bool ggml_webgpu_tensor_overlap(const ggml_tensor * a, const ggml_tensor * b) { - return a->buffer == b->buffer && ggml_webgpu_tensor_addr(a) < ggml_webgpu_tensor_addr(b) + ggml_nbytes(b) && - ggml_webgpu_tensor_addr(b) < ggml_webgpu_tensor_addr(a) + ggml_nbytes(a); -} - struct ggml_webgpu_shader_lib_context { ggml_tensor * src0; ggml_tensor * src1; @@ -118,6 +113,11 @@ struct ggml_webgpu_binary_shader_decisions { bool src_overlap = false; }; +struct ggml_webgpu_glu_shader_decisions { + uint32_t wg_size = 0; + bool src_overlap = false; +}; + struct ggml_webgpu_processed_shader { std::string wgsl; std::string variant; @@ -133,9 +133,12 @@ struct ggml_webgpu_ssm_scan_pipeline_key { int type; int d_state; bool xbc_overlap; + bool a_overlap; + bool ids_overlap; bool operator==(const ggml_webgpu_ssm_scan_pipeline_key & other) const { - return type == other.type && d_state == other.d_state && xbc_overlap == other.xbc_overlap; + return type == other.type && d_state == other.d_state && xbc_overlap == other.xbc_overlap && + a_overlap == other.a_overlap && ids_overlap == other.ids_overlap; } }; @@ -145,6 +148,8 @@ struct ggml_webgpu_ssm_scan_pipeline_key_hash { ggml_webgpu_hash_combine(seed, key.type); ggml_webgpu_hash_combine(seed, key.d_state); ggml_webgpu_hash_combine(seed, key.xbc_overlap); + ggml_webgpu_hash_combine(seed, key.a_overlap); + ggml_webgpu_hash_combine(seed, key.ids_overlap); return seed; } }; @@ -153,6 +158,8 @@ struct ggml_webgpu_ssm_scan_shader_decisions { uint32_t wg_size; uint32_t tokens_per_tile; bool xbc_overlap = false; + bool a_overlap = false; + bool ids_overlap = false; }; /** Argsort **/ @@ -264,7 +271,7 @@ struct ggml_webgpu_row_norm_pipeline_key_hash { struct ggml_webgpu_rms_norm_mul_pipeline_key { bool inplace; // rn_src == dst bool overlap; // mul_src == dst - bool src_overlap; // rn_src == mul_src + bool src_overlap; // rn_src binding overlaps mul_src binding bool operator==(const ggml_webgpu_rms_norm_mul_pipeline_key & other) const { return inplace == other.inplace && overlap == other.overlap && src_overlap == other.src_overlap; @@ -690,7 +697,8 @@ inline bool ggml_webgpu_flash_attn_kv_direct(const ggml_tensor * Q, inline ggml_webgpu_flash_attn_common_pipeline_key ggml_webgpu_flash_attn_make_common_pipeline_key( const ggml_webgpu_shader_lib_context & context, - uint32_t kv_direct_align) { + uint32_t kv_direct_align, + bool kv_overlap) { ggml_webgpu_flash_attn_common_pipeline_key key = {}; key.q_type = context.src0->type; key.k_type = context.src1->type; @@ -699,7 +707,7 @@ inline ggml_webgpu_flash_attn_common_pipeline_key ggml_webgpu_flash_attn_make_co key.head_dim_qk = (uint32_t) context.src0->ne[0]; key.head_dim_v = (uint32_t) context.src2->ne[0]; key.kv_direct = ggml_webgpu_flash_attn_kv_direct(context.src0, context.src1, context.src2, kv_direct_align); - key.kv_overlap = ggml_webgpu_tensor_overlap(context.src1, context.src2); + key.kv_overlap = kv_overlap; key.has_mask = context.src3 != nullptr; key.has_sinks = context.src4 != nullptr; key.uses_logit_softcap = ggml_get_op_params_f32(context.dst, 2) != 0.0f; @@ -1066,9 +1074,10 @@ struct ggml_webgpu_glu_pipeline_key { ggml_glu_op glu_op; ggml_type type; bool split; + bool src_overlap; bool operator==(const ggml_webgpu_glu_pipeline_key & other) const { - return glu_op == other.glu_op && type == other.type && split == other.split; + return glu_op == other.glu_op && type == other.type && split == other.split && src_overlap == other.src_overlap; } }; @@ -1078,6 +1087,7 @@ struct ggml_webgpu_glu_pipeline_key_hash { ggml_webgpu_hash_combine(seed, key.glu_op); ggml_webgpu_hash_combine(seed, key.type); ggml_webgpu_hash_combine(seed, key.split); + ggml_webgpu_hash_combine(seed, key.src_overlap); return seed; } }; @@ -1758,12 +1768,16 @@ class ggml_webgpu_shader_lib { return ssm_conv_pipelines[key]; } - webgpu_pipeline get_ssm_scan_pipeline(const ggml_webgpu_shader_lib_context & context) { + webgpu_pipeline get_ssm_scan_pipeline(const ggml_webgpu_shader_lib_context & context, + bool xbc_overlap, + bool a_overlap, + bool ids_overlap) { ggml_webgpu_ssm_scan_pipeline_key key = {}; key.type = context.dst->type; key.d_state = (int) context.src0->ne[0]; - key.xbc_overlap = ggml_webgpu_tensor_overlap(context.src1, context.src4) && - ggml_webgpu_tensor_overlap(context.src1, context.src5); + key.xbc_overlap = xbc_overlap; + key.a_overlap = a_overlap; + key.ids_overlap = ids_overlap; auto it = ssm_scan_pipelines.find(key); if (it != ssm_scan_pipelines.end()) { @@ -1798,7 +1812,12 @@ class ggml_webgpu_shader_lib { if (key.xbc_overlap) { defines.push_back("XBC_OVERLAP"); } - + if (key.a_overlap) { + defines.push_back("A_OVERLAP"); + } + if (key.ids_overlap) { + defines.push_back("IDS_OVERLAP"); + } variant += "_d" + std::to_string(key.d_state); auto processed = preprocessor.preprocess(wgsl_ssm_scan, defines); @@ -1806,6 +1825,8 @@ class ggml_webgpu_shader_lib { decisions->wg_size = wg_size; decisions->tokens_per_tile = tokens_per_tile; decisions->xbc_overlap = key.xbc_overlap; + decisions->a_overlap = key.a_overlap; + decisions->ids_overlap = key.ids_overlap; webgpu_pipeline pipeline = ggml_webgpu_create_pipeline(device, processed, variant); pipeline.context = decisions; ssm_scan_pipelines[key] = pipeline; @@ -2549,11 +2570,11 @@ class ggml_webgpu_shader_lib { return unary_pipelines[key]; } - webgpu_pipeline get_rms_norm_mul_pipeline(const ggml_webgpu_shader_lib_context & context) { + webgpu_pipeline get_rms_norm_mul_pipeline(const ggml_webgpu_shader_lib_context & context, bool src_overlap) { ggml_webgpu_rms_norm_mul_pipeline_key key = {}; key.inplace = ggml_webgpu_tensor_equal(context.src0, context.dst); key.overlap = ggml_webgpu_tensor_equal(context.src1, context.dst); - key.src_overlap = ggml_webgpu_tensor_overlap(context.src0, context.src1); + key.src_overlap = src_overlap; auto it = rms_norm_mul_pipelines.find(key); if (it != rms_norm_mul_pipelines.end()) { @@ -2589,13 +2610,13 @@ class ggml_webgpu_shader_lib { return rms_norm_mul_pipelines[key]; } - webgpu_pipeline get_binary_pipeline(const ggml_webgpu_shader_lib_context & context) { + webgpu_pipeline get_binary_pipeline(const ggml_webgpu_shader_lib_context & context, bool src_overlap) { ggml_webgpu_binary_pipeline_key key = {}; key.type = context.dst->type; key.op = context.dst->op; key.inplace = ggml_webgpu_tensor_equal(context.src0, context.dst); key.overlap = ggml_webgpu_tensor_equal(context.src1, context.dst); - key.src_overlap = ggml_webgpu_tensor_overlap(context.src0, context.src1); + key.src_overlap = src_overlap; auto it = binary_pipelines.find(key); if (it != binary_pipelines.end()) { @@ -2678,10 +2699,10 @@ class ggml_webgpu_shader_lib { return pipeline; } - webgpu_pipeline get_concat_pipeline(const ggml_webgpu_shader_lib_context & context) { + webgpu_pipeline get_concat_pipeline(const ggml_webgpu_shader_lib_context & context, bool src_overlap) { ggml_webgpu_concat_pipeline_key key = {}; key.type = context.dst->type; - key.src_overlap = ggml_webgpu_tensor_overlap(context.src0, context.src1); + key.src_overlap = src_overlap; auto it = concat_pipelines.find(key); if (it != concat_pipelines.end()) { @@ -2761,7 +2782,7 @@ class ggml_webgpu_shader_lib { return repeat_pipelines[key]; } - webgpu_pipeline get_flash_attn_pipeline(const ggml_webgpu_shader_lib_context & context) { + webgpu_pipeline get_flash_attn_pipeline(const ggml_webgpu_shader_lib_context & context, bool kv_overlap) { const bool can_use_subgroup_matrix = ggml_webgpu_flash_attn_can_use_subgroup_matrix_path( context.supports_subgroup_matrix, context.sg_mat_k, context.sg_mat_n, context.src0, context.src2); ggml_webgpu_flash_attn_decisions decisions = {}; @@ -2769,8 +2790,8 @@ class ggml_webgpu_shader_lib { decisions.q_tile = decisions.use_sg_matrix ? context.sg_mat_m : GGML_WEBGPU_FLASH_ATTN_TILE_Q_TILE; ggml_webgpu_flash_attn_pipeline_key key = {}; - key.common = - ggml_webgpu_flash_attn_make_common_pipeline_key(context, decisions.use_sg_matrix ? context.sg_mat_k : 1u); + key.common = ggml_webgpu_flash_attn_make_common_pipeline_key( + context, decisions.use_sg_matrix ? context.sg_mat_k : 1u, kv_overlap); key.common.kv_direct = decisions.use_sg_matrix && key.common.kv_direct; key.use_sg_matrix = decisions.use_sg_matrix; @@ -2824,9 +2845,10 @@ class ggml_webgpu_shader_lib { return flash_attn_pipelines[key]; } - webgpu_pipeline get_flash_attn_vec_pipeline(const ggml_webgpu_shader_lib_context & context) { + webgpu_pipeline get_flash_attn_vec_pipeline(const ggml_webgpu_shader_lib_context & context, bool kv_overlap) { ggml_webgpu_flash_attn_vec_pipeline_key key = {}; - key.common = ggml_webgpu_flash_attn_make_common_pipeline_key(context, GGML_WEBGPU_FLASH_ATTN_TILE_KV_VEC_WIDTH); + key.common = ggml_webgpu_flash_attn_make_common_pipeline_key(context, GGML_WEBGPU_FLASH_ATTN_TILE_KV_VEC_WIDTH, + kv_overlap); auto it = flash_attn_vec_pipelines.find(key); if (it != flash_attn_vec_pipelines.end()) { @@ -2984,11 +3006,12 @@ class ggml_webgpu_shader_lib { return cpy_pipelines[key]; } - webgpu_pipeline get_glu_pipeline(const ggml_webgpu_shader_lib_context & context) { + webgpu_pipeline get_glu_pipeline(const ggml_webgpu_shader_lib_context & context, bool src_overlap) { ggml_webgpu_glu_pipeline_key key = {}; key.glu_op = ggml_get_glu_op(context.dst); key.type = context.dst->type; key.split = (context.src1 != nullptr); + key.src_overlap = src_overlap; auto it = glu_pipelines.find(key); if (it != glu_pipelines.end()) { @@ -3039,7 +3062,10 @@ class ggml_webgpu_shader_lib { GGML_ABORT("Unsupported type for GLU shader"); } - if (key.split) { + if (key.src_overlap) { + defines.push_back("SRC_OVERLAP"); + variant += "_src_overlap"; + } else if (key.split) { variant += "_split"; } else { defines.push_back("NO_SPLIT"); @@ -3048,8 +3074,9 @@ class ggml_webgpu_shader_lib { defines.push_back(std::string("WG_SIZE=") + std::to_string(context.max_wg_size)); auto processed = preprocessor.preprocess(wgsl_glu, defines); - auto decisions = std::make_shared(); + auto decisions = std::make_shared(); decisions->wg_size = context.max_wg_size; + decisions->src_overlap = key.src_overlap; webgpu_pipeline pipeline = ggml_webgpu_create_pipeline(device, processed, variant); pipeline.context = decisions; glu_pipelines[key] = pipeline; diff --git a/ggml/src/ggml-webgpu/ggml-webgpu.cpp b/ggml/src/ggml-webgpu/ggml-webgpu.cpp index 75286ec7313c..2add5da0b493 100644 --- a/ggml/src/ggml-webgpu/ggml-webgpu.cpp +++ b/ggml/src/ggml-webgpu/ggml-webgpu.cpp @@ -374,18 +374,59 @@ static wgpu::Buffer ggml_webgpu_tensor_buf(const ggml_tensor * tensor) { return ctx->buffer; } +static size_t ggml_webgpu_tensor_misalignment(const ggml_tensor * t, size_t alignment) { + size_t offset = ggml_webgpu_tensor_offset(t); + return offset & (alignment - 1); +} + static size_t ggml_webgpu_tensor_misalignment(webgpu_context & ctx, const ggml_tensor * t) { + return ggml_webgpu_tensor_misalignment(t, ctx->global_ctx->capabilities.limits.minStorageBufferOffsetAlignment); +} + +static size_t ggml_webgpu_tensor_align_offset(const ggml_tensor * t, size_t alignment) { size_t offset = ggml_webgpu_tensor_offset(t); - return offset & (ctx->global_ctx->capabilities.limits.minStorageBufferOffsetAlignment - 1); + return offset & ~(alignment - 1); } static size_t ggml_webgpu_tensor_align_offset(webgpu_context & ctx, const ggml_tensor * t) { - size_t offset = ggml_webgpu_tensor_offset(t); - return offset & ~(ctx->global_ctx->capabilities.limits.minStorageBufferOffsetAlignment - 1); + return ggml_webgpu_tensor_align_offset(t, ctx->global_ctx->capabilities.limits.minStorageBufferOffsetAlignment); } -static size_t ggml_webgpu_tensor_binding_size(webgpu_context & ctx, ggml_tensor * t) { - return ROUNDUP_POW2(ggml_nbytes(t) + ggml_webgpu_tensor_misalignment(ctx, t), WEBGPU_STORAGE_BUF_BINDING_MULT); +static size_t ggml_webgpu_tensor_binding_size(const ggml_tensor * t, size_t alignment) { + return ROUNDUP_POW2(ggml_nbytes(t) + ggml_webgpu_tensor_misalignment(t, alignment), + WEBGPU_STORAGE_BUF_BINDING_MULT); +} + +static size_t ggml_webgpu_tensor_binding_size(webgpu_context & ctx, const ggml_tensor * t) { + return ggml_webgpu_tensor_binding_size(t, ctx->global_ctx->capabilities.limits.minStorageBufferOffsetAlignment); +} + +static bool ggml_webgpu_tensor_binding_overlap(const webgpu_global_context & global_ctx, + const ggml_tensor * a, + const ggml_tensor * b) { + if (a->buffer != b->buffer) { + return false; + } + + const size_t alignment = global_ctx->capabilities.limits.minStorageBufferOffsetAlignment; + const size_t a_offset = ggml_webgpu_tensor_align_offset(a, alignment); + const size_t b_offset = ggml_webgpu_tensor_align_offset(b, alignment); + return a_offset < b_offset + ggml_webgpu_tensor_binding_size(b, alignment) && + b_offset < a_offset + ggml_webgpu_tensor_binding_size(a, alignment); +} + +static bool ggml_webgpu_tensor_binding_overlap_range(const webgpu_global_context & global_ctx, + ggml_tensor * tensor, + ggml_backend_buffer_t buffer, + size_t offset, + size_t size) { + if (tensor->buffer != buffer) { + return false; + } + + const size_t alignment = global_ctx->capabilities.limits.minStorageBufferOffsetAlignment; + const size_t tensor_offset = ggml_webgpu_tensor_align_offset(tensor, alignment); + return tensor_offset < offset + size && offset < tensor_offset + ggml_webgpu_tensor_binding_size(tensor, alignment); } struct ggml_webgpu_merged_binding_range { @@ -1188,39 +1229,76 @@ static webgpu_encoded_op ggml_webgpu_ssm_scan(webgpu_context & ctx, ggml_webgpu_shader_lib_context shader_lib_ctx = {}; shader_lib_ctx.src0 = src0; shader_lib_ctx.src1 = src1; + shader_lib_ctx.src2 = src2; + shader_lib_ctx.src3 = src3; shader_lib_ctx.src4 = src4; shader_lib_ctx.src5 = src5; shader_lib_ctx.dst = dst; shader_lib_ctx.max_wg_size = ctx->global_ctx->capabilities.limits.maxComputeInvocationsPerWorkgroup; shader_lib_ctx.supports_subgroups = ctx->global_ctx->capabilities.supports_subgroups; + bool xbc_overlap = ggml_webgpu_tensor_binding_overlap(ctx->global_ctx, src1, src2) || + ggml_webgpu_tensor_binding_overlap(ctx->global_ctx, src1, src4) || + ggml_webgpu_tensor_binding_overlap(ctx->global_ctx, src1, src5) || + ggml_webgpu_tensor_binding_overlap(ctx->global_ctx, src2, src4) || + ggml_webgpu_tensor_binding_overlap(ctx->global_ctx, src2, src5) || + ggml_webgpu_tensor_binding_overlap(ctx->global_ctx, src4, src5); + bool a_overlap = false; + bool ids_overlap = false; + ggml_webgpu_merged_binding_range xbc_merged_range = {}; + if (xbc_overlap) { + xbc_merged_range = ggml_webgpu_tensor_merged_binding_range(ctx, { src1, src2, src4, src5 }); + a_overlap = ggml_webgpu_tensor_binding_overlap_range(ctx->global_ctx, src3, src1->buffer, + xbc_merged_range.offset, xbc_merged_range.size); + if (a_overlap) { + xbc_merged_range = ggml_webgpu_tensor_merged_binding_range(ctx, { src1, src2, src3, src4, src5 }); + } + ids_overlap = ggml_webgpu_tensor_binding_overlap_range(ctx->global_ctx, src6, src1->buffer, + xbc_merged_range.offset, xbc_merged_range.size); + if (ids_overlap) { + xbc_merged_range = + a_overlap ? ggml_webgpu_tensor_merged_binding_range(ctx, { src1, src2, src3, src4, src5, src6 }) : + ggml_webgpu_tensor_merged_binding_range(ctx, { src1, src2, src4, src5, src6 }); + } + } - webgpu_pipeline pipeline = ctx->shader_lib->get_ssm_scan_pipeline(shader_lib_ctx); - auto * decisions = static_cast(pipeline.context.get()); - const bool xbc_overlap = decisions->xbc_overlap; + webgpu_pipeline pipeline = + ctx->shader_lib->get_ssm_scan_pipeline(shader_lib_ctx, xbc_overlap, a_overlap, ids_overlap); + auto * decisions = static_cast(pipeline.context.get()); + xbc_overlap = decisions->xbc_overlap; + a_overlap = decisions->a_overlap; + ids_overlap = decisions->ids_overlap; uint32_t offset_x = (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, src1) / ggml_type_size(src1->type)); + uint32_t offset_dt = (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, src2) / ggml_type_size(src2->type)); + uint32_t offset_A = (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, src3) / ggml_type_size(src3->type)); uint32_t offset_B = (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, src4) / ggml_type_size(src4->type)); uint32_t offset_C = (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, src5) / ggml_type_size(src5->type)); + uint32_t offset_ids = (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, src6) / ggml_type_size(src6->type)); size_t xbc_bind_offset = 0; size_t xbc_bind_size = 0; if (xbc_overlap) { - const ggml_webgpu_merged_binding_range merged_range = - ggml_webgpu_tensor_merged_binding_range(ctx, { src1, src4, src5 }); - xbc_bind_offset = merged_range.offset; - xbc_bind_size = merged_range.size; - offset_x = ggml_webgpu_tensor_merged_element_offset(src1, merged_range); - offset_B = ggml_webgpu_tensor_merged_element_offset(src4, merged_range); - offset_C = ggml_webgpu_tensor_merged_element_offset(src5, merged_range); + xbc_bind_offset = xbc_merged_range.offset; + xbc_bind_size = xbc_merged_range.size; + offset_x = ggml_webgpu_tensor_merged_element_offset(src1, xbc_merged_range); + offset_dt = ggml_webgpu_tensor_merged_element_offset(src2, xbc_merged_range); + if (a_overlap) { + offset_A = ggml_webgpu_tensor_merged_element_offset(src3, xbc_merged_range); + } + offset_B = ggml_webgpu_tensor_merged_element_offset(src4, xbc_merged_range); + offset_C = ggml_webgpu_tensor_merged_element_offset(src5, xbc_merged_range); + if (ids_overlap) { + offset_ids = ggml_webgpu_tensor_merged_element_offset(src6, xbc_merged_range); + } } std::vector params = { (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, src0) / ggml_type_size(src0->type)), offset_x, - (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, src2) / ggml_type_size(src2->type)), - (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, src3) / ggml_type_size(src3->type)), + offset_dt, + offset_A, offset_B, offset_C, - (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, src6) / ggml_type_size(src6->type)), + offset_ids, (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, dst) / ggml_type_size(dst->type)), (uint32_t) (src0->nb[1] / ggml_type_size(src0->type)), @@ -1260,10 +1338,19 @@ static webgpu_encoded_op ggml_webgpu_ssm_scan(webgpu_context & ctx, if (xbc_overlap) { entries.push_back( ggml_webgpu_make_bind_group_entry(1, ggml_webgpu_tensor_buf(src1), xbc_bind_offset, xbc_bind_size)); - entries.push_back(ggml_webgpu_make_tensor_bind_group_entry(ctx, 2, src2)); - entries.push_back(ggml_webgpu_make_tensor_bind_group_entry(ctx, 3, src3)); - entries.push_back(ggml_webgpu_make_tensor_bind_group_entry(ctx, 4, src6)); - entries.push_back(ggml_webgpu_make_tensor_bind_group_entry(ctx, 5, dst)); + if (ids_overlap) { + if (!a_overlap) { + entries.push_back(ggml_webgpu_make_tensor_bind_group_entry(ctx, 2, src3)); + } + entries.push_back(ggml_webgpu_make_tensor_bind_group_entry(ctx, a_overlap ? 2 : 3, dst)); + } else if (a_overlap) { + entries.push_back(ggml_webgpu_make_tensor_bind_group_entry(ctx, 2, src6)); + entries.push_back(ggml_webgpu_make_tensor_bind_group_entry(ctx, 3, dst)); + } else { + entries.push_back(ggml_webgpu_make_tensor_bind_group_entry(ctx, 2, src3)); + entries.push_back(ggml_webgpu_make_tensor_bind_group_entry(ctx, 3, src6)); + entries.push_back(ggml_webgpu_make_tensor_bind_group_entry(ctx, 4, dst)); + } } else { entries.push_back(ggml_webgpu_make_tensor_bind_group_entry(ctx, 1, src1)); entries.push_back(ggml_webgpu_make_tensor_bind_group_entry(ctx, 2, src2)); @@ -1381,11 +1468,10 @@ static std::optional ggml_webgpu_set_rows(webgpu_context & ct (uint32_t) (idx->ne[1]), (uint32_t) (idx->ne[2]) }; - std::vector entries = { - ggml_webgpu_make_tensor_bind_group_entry(ctx, 0, src), - ggml_webgpu_make_tensor_bind_group_entry(ctx, 1, idx), - ggml_webgpu_make_tensor_bind_group_entry(ctx, 2, dst), - }; + std::vector entries; + entries.push_back(ggml_webgpu_make_tensor_bind_group_entry(ctx, 0, src)); + entries.push_back(ggml_webgpu_make_tensor_bind_group_entry(ctx, 1, idx)); + entries.push_back(ggml_webgpu_make_tensor_bind_group_entry(ctx, 2, dst)); if (decisions->i64_idx) { entries.push_back(ggml_webgpu_make_bind_group_entry(3, ctx->set_rows_dev_error_buf, 0, @@ -1892,7 +1978,7 @@ static ggml_webgpu_flash_attn_op ggml_webgpu_flash_attn_prepare(webgpu_context & op.has_mask = mask != nullptr; op.has_sinks = sinks != nullptr; - op.kv_overlap = ggml_webgpu_tensor_overlap(K, V); + op.kv_overlap = ggml_webgpu_tensor_binding_overlap(ctx->global_ctx, K, V); uint32_t offset_k = (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, K) / ggml_type_size(K->type)); uint32_t offset_v = (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, V) / ggml_type_size(V->type)); @@ -1964,7 +2050,7 @@ static uint32_t ggml_webgpu_flash_attn_vec_nwg(uint32_t vec_nwg_cap, uint32_t kv } static webgpu_encoded_op ggml_webgpu_flash_attn_direct(webgpu_context & ctx, const ggml_webgpu_flash_attn_op & op) { - webgpu_pipeline pipeline = ctx->shader_lib->get_flash_attn_pipeline(op.shader_lib_ctx); + webgpu_pipeline pipeline = ctx->shader_lib->get_flash_attn_pipeline(op.shader_lib_ctx, op.kv_overlap); auto * decisions = static_cast(pipeline.context.get()); uint32_t wg_per_head = CEIL_DIV(op.shader_lib_ctx.src0->ne[1], decisions->q_tile); uint32_t wg_x = wg_per_head * op.shader_lib_ctx.src0->ne[2] * op.shader_lib_ctx.src0->ne[3]; @@ -1979,7 +2065,7 @@ static webgpu_encoded_op ggml_webgpu_flash_attn_vec(webgpu_context & ct ggml_tensor * sinks, ggml_tensor * dst, ggml_webgpu_flash_attn_op op) { - webgpu_pipeline pipeline = ctx->shader_lib->get_flash_attn_vec_pipeline(op.shader_lib_ctx); + webgpu_pipeline pipeline = ctx->shader_lib->get_flash_attn_vec_pipeline(op.shader_lib_ctx, op.kv_overlap); auto * decisions = static_cast(pipeline.context.get()); wgpu::Buffer blk_buf = {}; @@ -2249,8 +2335,9 @@ static webgpu_encoded_op ggml_webgpu_binary_op(webgpu_context & ctx, shader_lib_ctx.dst = dst; shader_lib_ctx.max_wg_size = ctx->global_ctx->capabilities.limits.maxComputeInvocationsPerWorkgroup; - webgpu_pipeline pipeline = ctx->shader_lib->get_binary_pipeline(shader_lib_ctx); - auto * decisions = static_cast(pipeline.context.get()); + const bool src_overlap = ggml_webgpu_tensor_binding_overlap(ctx->global_ctx, src0, src1); + webgpu_pipeline pipeline = ctx->shader_lib->get_binary_pipeline(shader_lib_ctx, src_overlap); + auto * decisions = static_cast(pipeline.context.get()); uint32_t ne = (uint32_t) ggml_nelements(dst); @@ -2372,6 +2459,9 @@ static webgpu_encoded_op ggml_webgpu_concat(webgpu_context & ctx, ggml_tensor * dst) { uint32_t ne = (uint32_t) ggml_nelements(dst); uint32_t dim = (uint32_t) dst->op_params[0]; + if (ggml_nbytes(src0) == 0 && ggml_nbytes(src1) == 0) { + return {}; + } ggml_webgpu_shader_lib_context shader_lib_ctx = {}; shader_lib_ctx.src0 = src0; @@ -2379,20 +2469,34 @@ static webgpu_encoded_op ggml_webgpu_concat(webgpu_context & ctx, shader_lib_ctx.dst = dst; shader_lib_ctx.max_wg_size = ctx->global_ctx->capabilities.limits.maxComputeInvocationsPerWorkgroup; - webgpu_pipeline pipeline = ctx->shader_lib->get_concat_pipeline(shader_lib_ctx); - auto * decisions = static_cast(pipeline.context.get()); + const bool src_overlap = ggml_webgpu_tensor_binding_overlap(ctx->global_ctx, src0, src1) || + ggml_nbytes(src0) == 0 || ggml_nbytes(src1) == 0; + webgpu_pipeline pipeline = ctx->shader_lib->get_concat_pipeline(shader_lib_ctx, src_overlap); + auto * decisions = static_cast(pipeline.context.get()); uint32_t offset_src0 = (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, src0) / ggml_type_size(src0->type)); uint32_t offset_src1 = (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, src1) / ggml_type_size(src1->type)); size_t merged_offset = 0; size_t merged_size = 0; if (decisions->src_overlap) { - const ggml_webgpu_merged_binding_range merged_range = - ggml_webgpu_tensor_merged_binding_range(ctx, { src0, src1 }); - merged_offset = merged_range.offset; - merged_size = merged_range.size; - offset_src0 = ggml_webgpu_tensor_merged_element_offset(src0, merged_range); - offset_src1 = ggml_webgpu_tensor_merged_element_offset(src1, merged_range); + if (ggml_nbytes(src0) == 0) { + merged_offset = ggml_webgpu_tensor_align_offset(ctx, src1); + merged_size = ggml_webgpu_tensor_binding_size(ctx, src1); + offset_src0 = 0; + offset_src1 = (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, src1) / ggml_type_size(src1->type)); + } else if (ggml_nbytes(src1) == 0) { + merged_offset = ggml_webgpu_tensor_align_offset(ctx, src0); + merged_size = ggml_webgpu_tensor_binding_size(ctx, src0); + offset_src0 = (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, src0) / ggml_type_size(src0->type)); + offset_src1 = 0; + } else { + const ggml_webgpu_merged_binding_range merged_range = + ggml_webgpu_tensor_merged_binding_range(ctx, { src0, src1 }); + merged_offset = merged_range.offset; + merged_size = merged_range.size; + offset_src0 = ggml_webgpu_tensor_merged_element_offset(src0, merged_range); + offset_src1 = ggml_webgpu_tensor_merged_element_offset(src1, merged_range); + } } std::vector params = { ne, @@ -2518,8 +2622,9 @@ static std::optional ggml_webgpu_rms_norm_mul(webgpu_context shader_lib_ctx.dst = dst; shader_lib_ctx.max_wg_size = ctx->global_ctx->capabilities.limits.maxComputeInvocationsPerWorkgroup; - webgpu_pipeline pipeline = ctx->shader_lib->get_rms_norm_mul_pipeline(shader_lib_ctx); - auto * decisions = static_cast(pipeline.context.get()); + const bool src_overlap = ggml_webgpu_tensor_binding_overlap(ctx->global_ctx, rn_src, mul_src); + webgpu_pipeline pipeline = ctx->shader_lib->get_rms_norm_mul_pipeline(shader_lib_ctx, src_overlap); + auto * decisions = static_cast(pipeline.context.get()); if (decisions->src_overlap) { const ggml_webgpu_merged_binding_range merged_range = @@ -2678,15 +2783,30 @@ static webgpu_encoded_op ggml_webgpu_glu(webgpu_context & ctx, shader_lib_ctx.dst = dst; shader_lib_ctx.max_wg_size = ctx->global_ctx->capabilities.limits.maxComputeInvocationsPerWorkgroup; - webgpu_pipeline pipeline = ctx->shader_lib->get_glu_pipeline(shader_lib_ctx); + const bool src_overlap = src1 != nullptr && ggml_webgpu_tensor_binding_overlap(ctx->global_ctx, src0, src1); + webgpu_pipeline pipeline = ctx->shader_lib->get_glu_pipeline(shader_lib_ctx, src_overlap); - auto * decisions = static_cast(pipeline.context.get()); + auto * decisions = static_cast(pipeline.context.get()); const int split = (src1 != nullptr); + uint32_t offset_src0 = (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, src0) / ggml_type_size(src0->type)); + uint32_t offset_src1 = + src1 != nullptr ? (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, src1) / ggml_type_size(src1->type)) : 0; + size_t merged_offset = 0; + size_t merged_size = 0; + if (decisions->src_overlap) { + const ggml_webgpu_merged_binding_range merged_range = + ggml_webgpu_tensor_merged_binding_range(ctx, { src0, src1 }); + merged_offset = merged_range.offset; + merged_size = merged_range.size; + offset_src0 = ggml_webgpu_tensor_merged_element_offset(src0, merged_range); + offset_src1 = ggml_webgpu_tensor_merged_element_offset(src1, merged_range); + } + std::vector params = { - (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, src0) / ggml_type_size(src0->type)), - src1 != nullptr ? (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, src1) / ggml_type_size(src1->type)) : 0, + offset_src0, + offset_src1, (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, dst) / ggml_type_size(dst->type)), (uint32_t) (src0->nb[1] / ggml_type_size(src0->type)), (uint32_t) (src0->nb[2] / ggml_type_size(src0->type)), @@ -2709,11 +2829,15 @@ static webgpu_encoded_op ggml_webgpu_glu(webgpu_context & ctx, ggml_webgpu_u32_from_f32(ggml_get_op_params_f32(dst, 3)), // limit, for swiglu_oai }; - std::vector entries = { - ggml_webgpu_make_tensor_bind_group_entry(ctx, 0, src0), - }; - uint32_t dst_binding = 1; - if (split) { + std::vector entries; + uint32_t dst_binding = 1; + if (decisions->src_overlap) { + entries.push_back( + ggml_webgpu_make_bind_group_entry(0, ggml_webgpu_tensor_buf(src0), merged_offset, merged_size)); + } else { + entries.push_back(ggml_webgpu_make_tensor_bind_group_entry(ctx, 0, src0)); + } + if (split && !decisions->src_overlap) { dst_binding = 2; entries.push_back(ggml_webgpu_make_tensor_bind_group_entry(ctx, 1, src1)); } @@ -4285,8 +4409,8 @@ static bool ggml_backend_webgpu_device_supports_op(ggml_backend_dev_t dev, const if (!supports_op) { break; } - if (ggml_webgpu_tensor_overlap(src1, src2) && src1->type != src2->type && - !ggml_is_quantized(src1->type) && !ggml_is_quantized(src2->type)) { + if (ggml_webgpu_tensor_binding_overlap(ctx->webgpu_global_ctx, src1, src2) && + src1->type != src2->type && !ggml_is_quantized(src1->type) && !ggml_is_quantized(src2->type)) { supports_op = false; break; } diff --git a/ggml/src/ggml-webgpu/wgsl-shaders/glu.wgsl b/ggml/src/ggml-webgpu/wgsl-shaders/glu.wgsl index e6d7608cec5d..d03f1c207d98 100644 --- a/ggml/src/ggml-webgpu/wgsl-shaders/glu.wgsl +++ b/ggml/src/ggml-webgpu/wgsl-shaders/glu.wgsl @@ -96,7 +96,22 @@ struct Params { @group(0) @binding(0) var src0: array; -#ifdef NO_SPLIT +#ifdef SRC_OVERLAP +@group(0) @binding(1) +var dst: array; + +@group(0) @binding(2) +var params: Params; + +fn a_value(base: u32) -> DataType { + return src0[base]; +} + +fn b_value(base: u32) -> DataType { + return src0[base]; +} + +#elif defined(NO_SPLIT) @group(0) @binding(1) var dst: array; diff --git a/ggml/src/ggml-webgpu/wgsl-shaders/ssm_scan.wgsl b/ggml/src/ggml-webgpu/wgsl-shaders/ssm_scan.wgsl index 05761dec353a..66bfdd64015c 100644 --- a/ggml/src/ggml-webgpu/wgsl-shaders/ssm_scan.wgsl +++ b/ggml/src/ggml-webgpu/wgsl-shaders/ssm_scan.wgsl @@ -46,12 +46,29 @@ struct Params { @group(0) @binding(0) var s_in: array; #ifdef XBC_OVERLAP -@group(0) @binding(1) var x_B_C_merged: array; -@group(0) @binding(2) var dt: array; -@group(0) @binding(3) var A: array; -@group(0) @binding(4) var ids: array; -@group(0) @binding(5) var dst: array; -@group(0) @binding(6) var params: Params; +#ifdef IDS_OVERLAP +@group(0) @binding(1) var x_dt_B_C_ids_merged: array; +#ifdef A_OVERLAP +@group(0) @binding(2) var dst: array; +@group(0) @binding(3) var params: Params; +#else +@group(0) @binding(2) var A: array; +@group(0) @binding(3) var dst: array; +@group(0) @binding(4) var params: Params; +#endif +#else +@group(0) @binding(1) var x_dt_B_C_merged: array; +#ifdef A_OVERLAP +@group(0) @binding(2) var ids: array; +@group(0) @binding(3) var dst: array; +@group(0) @binding(4) var params: Params; +#else +@group(0) @binding(2) var A: array; +@group(0) @binding(3) var ids: array; +@group(0) @binding(4) var dst: array; +@group(0) @binding(5) var params: Params; +#endif +#endif #else @group(0) @binding(1) var x: array; @group(0) @binding(2) var dt: array; @@ -71,6 +88,24 @@ fn reduce_base(token_in_tile: u32) -> u32 { return token_in_tile * WG_SIZE; } +#ifdef XBC_OVERLAP +fn read_merged_f32(idx: u32) -> f32 { +#ifdef IDS_OVERLAP + return bitcast(x_dt_B_C_ids_merged[idx]); +#else + return x_dt_B_C_merged[idx]; +#endif +} +#endif + +fn read_state_slot(i3: u32) -> u32 { +#ifdef IDS_OVERLAP + return x_dt_B_C_ids_merged[params.offset_ids + i3]; +#else + return u32(ids[params.offset_ids + i3]); +#endif +} + @compute @workgroup_size(WG_SIZE) fn main( @builtin(local_invocation_id) local_id: vec3, @@ -90,13 +125,18 @@ fn main( let ir = head_seq % params.n_head; let i3 = head_seq / params.n_head; - let state_slot = u32(ids[params.offset_ids + i3]); + let state_slot = read_state_slot(i3); let g = ir / (params.n_head / params.n_group); let s_idx = params.offset_s + tid + i1 * params.stride_s1 + ir * params.stride_s2 + state_slot * params.stride_s3; var s_prev = s_in[s_idx]; - let A0 = A[params.offset_A + (tid % params.a_ne0) + ir * params.stride_A1]; + let a_idx = params.offset_A + (tid % params.a_ne0) + ir * params.stride_A1; +#ifdef A_OVERLAP + let A0 = read_merged_f32(a_idx); +#else + let A0 = A[a_idx]; +#endif for (var token_base = 0u; token_base < params.n_seq_tokens; token_base += TOKENS_PER_TILE) { if (tid < TOKENS_PER_TILE) { @@ -104,11 +144,15 @@ fn main( if (token < params.n_seq_tokens) { let x_idx = params.offset_x + i1 + ir * params.stride_x1 + token * params.stride_x2 + i3 * params.stride_x3; let dt_idx = params.offset_dt + ir + token * params.stride_dt1 + i3 * params.stride_dt2; +#ifdef XBC_OVERLAP + let dt0 = read_merged_f32(dt_idx); +#else let dt0 = dt[dt_idx]; +#endif let dtsp = select(log(1.0 + exp(dt0)), dt0, dt0 > 20.0); shared_dtsp[tid] = dtsp; #ifdef XBC_OVERLAP - shared_x_dt[tid] = x_B_C_merged[x_idx] * dtsp; + shared_x_dt[tid] = read_merged_f32(x_idx) * dtsp; #else shared_x_dt[tid] = x[x_idx] * dtsp; #endif @@ -130,7 +174,7 @@ fn main( let b_idx = params.offset_B + tid + g * params.stride_B1 + token * params.stride_B2 + i3 * params.stride_B3; let c_idx = params.offset_C + tid + g * params.stride_C1 + token * params.stride_C2 + i3 * params.stride_C3; #ifdef XBC_OVERLAP - let s = s_prev * dA + x_B_C_merged[b_idx] * x_dt; + let s = s_prev * dA + read_merged_f32(b_idx) * x_dt; #else let s = s_prev * dA + B[b_idx] * x_dt; #endif @@ -138,7 +182,7 @@ fn main( #ifdef USE_SUBGROUP_REDUCTION #ifdef XBC_OVERLAP - let subgroup_partial = subgroupAdd(s * x_B_C_merged[c_idx]); + let subgroup_partial = subgroupAdd(s * read_merged_f32(c_idx)); #else let subgroup_partial = subgroupAdd(s * C[c_idx]); #endif @@ -147,7 +191,7 @@ fn main( } #else #ifdef XBC_OVERLAP - shared_reduce[reduce_idx] = s * x_B_C_merged[c_idx]; + shared_reduce[reduce_idx] = s * read_merged_f32(c_idx); #else shared_reduce[reduce_idx] = s * C[c_idx]; #endif diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index d02e65c9ead0..a1ed2a76f879 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -428,9 +428,9 @@ static bool arch_supported(const llm_arch arch) { return false; } - // FIXME some models are segfaulting with WebGPU: + // FIXME: these hit scheduler/view-backed-output issues with WebGPU on CI. #ifdef GGML_USE_WEBGPU - if (arch == LLM_ARCH_QWEN3NEXT || arch == LLM_ARCH_QWEN35 || arch == LLM_ARCH_QWEN35MOE || arch == LLM_ARCH_KIMI_LINEAR) { + if (arch == LLM_ARCH_DEEPSEEK32 || arch == LLM_ARCH_GLM_DSA) { return false; } #endif // GGML_USE_WEBGPU @@ -600,9 +600,6 @@ static int test_backends(const llm_arch target_arch, const size_t seed, const gg std::string status_roundtrip = "\033[1;33mSKIP\033[0m"; char nmse_str[12] = {0}; bool skip = !arch_supported(arch) || (dc.split_mode == LLAMA_SPLIT_MODE_TENSOR && dc.devs.empty()); -#if defined(GGML_USE_WEBGPU) - skip = true; // FIXME -#endif // GGML_USE_WEBGPU if (!skip) { if (logits_cpu.empty()) { model_and_ctx_cpu = get_model_and_ctx(gguf_ctx.get(), nullptr, seed, {}, LLAMA_SPLIT_MODE_LAYER, encode); From e9fa0781f1c25fc4fe8c86be1edc6970661ad6f0 Mon Sep 17 00:00:00 2001 From: Guido Imperiale Date: Tue, 28 Jul 2026 20:02:33 +0100 Subject: [PATCH 045/190] model: Add Laguna-S-2.1 LLM_TYPE (#26233) --- src/llama-model.cpp | 1 + src/llama-model.h | 1 + src/models/laguna.cpp | 1 + 3 files changed, 3 insertions(+) diff --git a/src/llama-model.cpp b/src/llama-model.cpp index be0a0df55d62..7a70585fa4d8 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -818,6 +818,7 @@ const char * llm_type_name(llm_type type) { case LLM_TYPE_100B_A6B: return "100B.A6B"; case LLM_TYPE_102B_A12B: return "102B.A12B"; case LLM_TYPE_106B_A12B: return "106B.A12B"; + case LLM_TYPE_118B_A8B: return "118B.A8B"; case LLM_TYPE_120B_A12B: return "120B.A12B"; case LLM_TYPE_122B_A10B: return "122B.A10B"; case LLM_TYPE_196B_A11B: return "196B.A11B"; diff --git a/src/llama-model.h b/src/llama-model.h index d6a40fa30204..056a6efa59e8 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -130,6 +130,7 @@ enum llm_type { LLM_TYPE_100B_A6B, LLM_TYPE_102B_A12B, // Solar-Open LLM_TYPE_106B_A12B, // GLM-4.5-Air + LLM_TYPE_118B_A8B, // Laguna-S-2 LLM_TYPE_120B_A12B, // Nemotron 3 Super LLM_TYPE_122B_A10B, // Qwen3.5 LLM_TYPE_196B_A11B, // Step3.5-Flash diff --git a/src/models/laguna.cpp b/src/models/laguna.cpp index fb55ec12f934..82c9a9538cd4 100644 --- a/src/models/laguna.cpp +++ b/src/models/laguna.cpp @@ -58,6 +58,7 @@ void llama_model_laguna::load_arch_hparams(llama_model_loader & ml) { switch (hparams.n_layer()) { case 40: type = LLM_TYPE_30B_A3B; break; // Laguna-XS.2 + case 48: type = LLM_TYPE_118B_A8B; break; // Laguna-S.2 case 70: type = LLM_TYPE_230B_A10B; break; // Laguna-M.1 default: type = LLM_TYPE_UNKNOWN; } From 7be2c65dc9adee9bae784478be0f656e3b683431 Mon Sep 17 00:00:00 2001 From: Satinder Grewal Date: Wed, 29 Jul 2026 18:02:31 +1200 Subject: [PATCH 046/190] model: add NextN/MTP speculative decoding support for GLM_DSA (GLM-5.2) (#25980) * model: add NextN/MTP speculative decoding support for GLM_DSA (GLM-5.2) Adds GLM-5.2 NextN/MTP as a --spec-type draft-mtp target: nextn tensor loading via the qwen35moe/step35-style presence probe, a graph_mtp builder (enorm/hnorm/eh_proj + dense MLA + sigmoid-gated MoE with shared expert + shared head with fallbacks, _s scale tensors passed for NVFP4), t_h_nextn extraction in the trunk graph, and MTP-context KV setup: the draft head runs dense MLA, so the MTP context uses a plain attention KV cache holding only the nextn layer(s) (same pattern as the hybrid Qwen3.5 MTP context) while the main context keeps the DSA cache, now filtered to trunk layers only. Co-Authored-By: Claude Fable 5 * convert : support --mtp/--no-mtp export for GlmMoeDsaForCausalLM (GLM-5.2) Opt GLM-5.2 into the supports_mtp_export contract (post-#25641 shape, mirroring HYV3Model/Step35Model): --no-mtp drops the appended NextN block (blk.78) and its nextn_predict_layers KV; --mtp keeps only the NextN block plus shared embeddings/norm/lm_head. Default (bundled) output is unchanged. Co-Authored-By: Claude Fable 5 --------- Co-authored-by: Claude Fable 5 --- conversion/glm.py | 43 ++++++- src/llama-model.cpp | 53 +++++++- src/models/glm-dsa.cpp | 281 +++++++++++++++++++++++++++++++++++++++-- src/models/models.h | 4 + 4 files changed, 366 insertions(+), 15 deletions(-) diff --git a/conversion/glm.py b/conversion/glm.py index d85268a62149..cc34cddbf843 100644 --- a/conversion/glm.py +++ b/conversion/glm.py @@ -1,6 +1,8 @@ from __future__ import annotations -from typing import Iterable, TYPE_CHECKING +import re + +from typing import Callable, Iterable, TYPE_CHECKING import torch @@ -213,12 +215,47 @@ def set_vocab(self): class GlmMoeDsaModel(DeepseekV2Model): model_arch = gguf.MODEL_ARCH.GLM_DSA skip_mtp = False + supports_mtp_export = True + + # Trunk layer count, stashed before indexing so the classmethod + # filter_tensors can identify the appended NextN/MTP block (mirrors + # HYV3Model / Step35Model). + _n_main_layers: int | None = None def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) - self.block_count = self.hparams["num_hidden_layers"] + self.hparams.get("num_nextn_predict_layers", 0) + self.block_count = self.hparams["num_hidden_layers"] + if not self.no_mtp: + self.block_count += self.hparams.get("num_nextn_predict_layers", 0) self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count) + def index_tensors(self, remote_hf_model_id: str | None = None): + type(self)._n_main_layers = self.hparams["num_hidden_layers"] + return super().index_tensors(remote_hf_model_id=remote_hf_model_id) + + @classmethod + def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None: + if (titem := super().filter_tensors(item)) is None: + return None + name, gen = titem + + # GLM-5.2 appends the NextN/MTP block past num_hidden_layers + # (model.layers.78 -> blk.78 in the 79-block file). + assert cls._n_main_layers is not None + is_mtp = (m := re.match(r"model\.layers\.(\d+)\.", name)) is not None and int(m.group(1)) >= cls._n_main_layers + + # --no-mtp: drop the appended NextN block entirely. + if is_mtp and cls.no_mtp: + return None + # --mtp: keep ONLY NextN-block tensors plus the shared embeddings/ + # norm/lm_head (so the resulting GGUF carries just the draft head). + if cls.mtp_only and not is_mtp and name not in ( + "model.embed_tokens.weight", "model.norm.weight", "lm_head.weight", + ): + return None + + return name, gen + def set_vocab(self): return self._set_vocab_glm() @@ -230,7 +267,7 @@ def set_gguf_parameters(self): self.gguf_writer.add_rope_dimension_count(int(rope_dim * partial_rotary_factor)) # NextN/MTP prediction layers - if (num_nextn_predict_layers := self.hparams.get("num_nextn_predict_layers")) is not None: + if not self.no_mtp and (num_nextn_predict_layers := self.hparams.get("num_nextn_predict_layers")) is not None: self.gguf_writer.add_nextn_predict_layers(num_nextn_predict_layers) # DSA indexer parameters diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 7a70585fa4d8..a8422ff6052b 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -2072,7 +2072,6 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, res = nullptr; } break; case LLM_ARCH_DEEPSEEK32: - case LLM_ARCH_GLM_DSA: { res = new llama_kv_cache_dsa( *this, @@ -2089,6 +2088,56 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, nullptr, nullptr); } break; + case LLM_ARCH_GLM_DSA: + { + if (params.ctx_type == LLAMA_CONTEXT_TYPE_MTP && hparams.n_layer_nextn > 0) { + // The NextN/MTP draft head runs dense MLA (no DSA indexer), so the + // MTP context uses a plain attention KV cache holding only the + // nextn layer(s) - same pattern as the hybrid Qwen3.5 MTP context. + llama_kv_cache::layer_filter_cb filter = + [&](uint32_t il) { return il >= hparams.n_layer(); }; + + res = new llama_kv_cache( + *this, + hparams, + params.type_k, + params.type_v, + !cparams.flash_attn, + cparams.offload_kqv, + cparams.kv_unified, + cparams.n_ctx_seq, + cparams.n_seq_max, + 1, + hparams.n_swa, + hparams.swa_type, + nullptr, + filter, + nullptr, + nullptr); + } else { + // Main context: DSA cache for the trunk layers only - the nextn + // layer(s) are never attended by the trunk graph. + llama_kv_cache::layer_filter_cb filter = nullptr; + if (hparams.n_layer_nextn > 0) { + filter = [&](uint32_t il) { return il < hparams.n_layer(); }; + } + + res = new llama_kv_cache_dsa( + *this, + params.type_k, + params.type_v, + !cparams.flash_attn, + cparams.offload_kqv, + cparams.kv_unified, + cparams.n_ctx_seq, + cparams.n_seq_max, + 1, + hparams.n_swa, + hparams.swa_type, + filter, + nullptr); + } + } break; // Models that need standard caching should rely on recurrent/hybrid // checks default: @@ -2194,7 +2243,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, filter = [&](uint32_t il) { return il >= hparams.n_layer(); }; } - if ((arch == LLM_ARCH_STEP35 || arch == LLM_ARCH_HY_V3) && hparams.n_layer_nextn > 0) { + if ((arch == LLM_ARCH_STEP35 || arch == LLM_ARCH_HY_V3 || arch == LLM_ARCH_GLM_DSA) && hparams.n_layer_nextn > 0) { if (params.ctx_type == LLAMA_CONTEXT_TYPE_MTP) { filter = [&](uint32_t il) { return il >= hparams.n_layer(); }; } else { diff --git a/src/models/glm-dsa.cpp b/src/models/glm-dsa.cpp index df190e1f634b..bd1c4df21281 100644 --- a/src/models/glm-dsa.cpp +++ b/src/models/glm-dsa.cpp @@ -72,15 +72,27 @@ void llama_model_glm_dsa::load_arch_hparams(llama_model_loader & ml) { ml.get_key_or_arr(LLM_KV_ATTENTION_INDEXER_TYPES, hparams.is_indexer_full_impl, hparams.n_layer(), false); switch (hparams.n_layer()) { - case 78: type = LLM_TYPE_744B_A40B; break; + case 78: // GGUF with NextN/MTP metadata: n_layer() excludes the nextn layer + case 79: + type = LLM_TYPE_744B_A40B; break; default: type = LLM_TYPE_UNKNOWN; } } -void llama_model_glm_dsa::load_arch_tensors(llama_model_loader &) { +void llama_model_glm_dsa::load_arch_tensors(llama_model_loader & ml) { LLAMA_LOAD_LOCALS; const int64_t n_expert_shared = hparams.n_expert_shared; + // MTP-only: the GGUF carries only the NextN/MTP block(s) (user split target/draft). + const bool mtp_only = (hparams.n_layer_nextn > 0) && (ml.get_weight("blk.0.attn_norm.weight") == nullptr); + // Trunk-only: the GGUF declares MTP layers in metadata but the actual MTP + // tensors live in a separate file (or were stripped at conversion). Mark + // MTP tensors NOT_REQUIRED so the trunk loads cleanly. + const std::string mtp_probe = "blk." + std::to_string(n_layer) + ".nextn.eh_proj.weight"; + const bool trunk_only = (hparams.n_layer_nextn > 0) && (ml.get_weight(mtp_probe.c_str()) == nullptr); + const int trunk_flags = mtp_only ? TENSOR_NOT_REQUIRED : 0; + const int mtp_flags = trunk_only ? TENSOR_NOT_REQUIRED : 0; + const bool is_mla = hparams.is_mla(); if (!is_mla) { throw std::runtime_error("GLM_DSA architecture requires MLA"); @@ -109,12 +121,9 @@ void llama_model_glm_dsa::load_arch_tensors(llama_model_loader &) { } for (int i = 0; i < n_layer_all; ++i) { - int flags = 0; - if (i >= n_layer) { - // skip all tensors in the NextN layers - // TODO @ngxson : TENSOR_NOT_REQUIRED was a hack, need to remove it later - flags |= TENSOR_SKIP | TENSOR_NOT_REQUIRED; - } + // NextN/MTP layers (i >= n_layer) are full decoder blocks used by the + // LLM_GRAPH_TYPE_DECODER_MTP draft head; load them like qwen35moe/step35/hy_v3. + const int flags = (i >= n_layer) ? mtp_flags : trunk_flags; auto & layer = layers[i]; @@ -167,7 +176,7 @@ void llama_model_glm_dsa::load_arch_tensors(llama_model_loader &) { layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", i), {n_embd, n_ff_exp * n_expert_shared}, flags); } - // NextN/MTP tensors (preserved but unused) - conditionally load for last n_layer_nextn + // NextN/MTP tensors - the NextN-specific wiring around the extra decoder block if (i >= n_layer) { layer.nextn.eh_proj = create_tensor(tn(LLM_TENSOR_NEXTN_EH_PROJ, "weight", i), { 2 * n_embd, n_embd }, flags); layer.nextn.enorm = create_tensor(tn(LLM_TENSOR_NEXTN_ENORM, "weight", i), { n_embd }, flags); @@ -182,6 +191,9 @@ void llama_model_glm_dsa::load_arch_tensors(llama_model_loader &) { } std::unique_ptr llama_model_glm_dsa::build_arch_graph(const llm_graph_params & params) const { + if (params.gtype == LLM_GRAPH_TYPE_DECODER_MTP) { + return std::make_unique(*this, params); + } return std::make_unique(*this, params); } @@ -469,7 +481,9 @@ llama_model_glm_dsa::graph::graph(const llama_model & model, const llm_graph_par Qcur, Kcur, Vcur, nullptr, nullptr, model.layers[il].wv_b, top_k, kq_scale, il); } } - if (il == n_layer - 1 && inp_out_ids) { + // when unmasked nextn embeddings are requested, t_h_nextn must keep all rows, + // so the early output masking has to be skipped (it is applied after the final norm instead) + if (il == n_layer - 1 && inp_out_ids && (!cparams.embeddings_nextn || cparams.embeddings_nextn_masked)) { cur = ggml_get_rows(ctx0, cur, inp_out_ids); inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids); } @@ -532,6 +546,14 @@ llama_model_glm_dsa::graph::graph(const llama_model & model, const llm_graph_par cur = build_norm(cur, model.output_norm, NULL, LLM_NORM_RMS, -1); + // post-norm hidden state feeds the NextN/MTP draft head + cb(cur, "h_nextn", -1); + res->t_h_nextn = cur; + + if (cparams.embeddings_nextn && !cparams.embeddings_nextn_masked && inp_out_ids) { + cur = ggml_get_rows(ctx0, cur, inp_out_ids); + } + cb(cur, "result_norm", -1); res->t_embd = cur; @@ -543,3 +565,242 @@ llama_model_glm_dsa::graph::graph(const llama_model & model, const llm_graph_par ggml_build_forward_expand(gf, cur); } + +// LLM_GRAPH_TYPE_DECODER_MTP draft head for GLM-5.2 (GLM_DSA). +// Semantics mirror the deepseek-family NextN/MTP layer: +// enorm(embed) + hnorm(prev_hidden) -> concat(e, h) -> eh_proj -> +// full glm_dsa decoder block (dense MLA attention + sigmoid-gated MoE FFN +// with shared expert, exactly as the trunk deepseek2 graph builds it) -> +// shared_head_norm (fallback output_norm) -> shared LM head. +// The DSA indexer is not used at runtime (same as the trunk graph). +llama_model_glm_dsa::graph_mtp::graph_mtp(const llama_model & model, const llm_graph_params & params) + : llm_graph_context(params) { + GGML_ASSERT(hparams.n_layer_nextn > 0 && "GLM_DSA MTP requires n_layer_nextn > 0"); + GGML_ASSERT(hparams.n_layer_nextn == 1 && "GLM_DSA MTP currently only supports a single MTP block"); + GGML_ASSERT(hparams.is_mla() && "GLM_DSA MTP requires MLA"); + + const int il = hparams.n_layer() + cparams.nextn_layer_offset; + GGML_ASSERT(cparams.nextn_layer_offset >= 0 && + cparams.nextn_layer_offset < (int) hparams.n_layer_nextn && + "nextn_layer_offset out of range [0, n_layer_nextn)"); + const auto & layer = model.layers[il]; + + GGML_ASSERT(layer.nextn.eh_proj && "MTP block missing nextn.eh_proj"); + GGML_ASSERT(layer.nextn.enorm && "MTP block missing nextn.enorm"); + GGML_ASSERT(layer.nextn.hnorm && "MTP block missing nextn.hnorm"); + GGML_ASSERT(layer.ffn_gate_inp && "MTP block missing ffn_gate_inp"); + + // note: these are the actual head sizes you get when treating as MHA or after "decompression" using wv_b for MLA + const int64_t n_embd_head_k = hparams.n_embd_head_k_mla(); + + const int64_t n_embd_head_qk_rope = hparams.n_rot(); + const int64_t n_embd_head_qk_nope = n_embd_head_k - n_embd_head_qk_rope; + + const uint32_t kv_lora_rank = hparams.n_lora_kv; + + // We have to pre-scale kq_scale and attn_factor to make the YaRN RoPE work correctly. + // See the deepseek2 trunk graph for the detailed explanation - this must match it EXACTLY. + GGML_ASSERT(ext_factor >= 0.0f); + const float attn_factor_org = attn_factor * (1.0f + 0.1f * logf(1.0f / freq_scale)); + + const float mscale = attn_factor_org * (1.0f + 0.1f * hparams.rope_yarn_log_mul * logf(1.0f / freq_scale)); + const float kq_scale = 1.0f * mscale * mscale / sqrtf(float(n_embd_head_k)); + + // TODO: extract in a common llm_graph_context::build_inp_embd_h() + auto inp = std::make_unique(hparams.n_embd); + + inp->tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens); + ggml_set_input(inp->tokens); + + inp->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd_inp(), n_tokens); + ggml_set_input(inp->embd); + + ggml_tensor * tok_embd; + if (ubatch.token) { + ggml_tensor * tok_embd_w = layer.nextn.embed_tokens ? layer.nextn.embed_tokens : model.tok_embd; + + tok_embd = ggml_get_rows(ctx0, tok_embd_w, inp->tokens); + } else { + tok_embd = inp->embd; + } + cb(tok_embd, "mtp_tok_embd", il); + + inp->h = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd, n_tokens); + ggml_set_input(inp->h); + ggml_set_name(inp->h, "mtp_h_input"); + + ggml_tensor * h_embd = inp->h; + + res->add_input(std::move(inp)); + + ggml_tensor * inp_pos = build_inp_pos(); + ggml_tensor * inp_out_ids = build_inp_out_ids(); + + // MLA with the absorption optimization uses a K-only cache (V is a view of K) + auto * inp_attn = build_attn_inp_k(); + + ggml_tensor * h_norm = build_norm(h_embd, layer.nextn.hnorm, nullptr, LLM_NORM_RMS, il); + cb(h_norm, "mtp_hnorm", il); + + ggml_tensor * e_norm = build_norm(tok_embd, layer.nextn.enorm, nullptr, LLM_NORM_RMS, il); + cb(e_norm, "mtp_enorm", il); + + ggml_tensor * concat = ggml_concat(ctx0, e_norm, h_norm, /*dim=*/ 0); + cb(concat, "mtp_concat", il); + + ggml_tensor * cur = build_lora_mm(layer.nextn.eh_proj, concat, layer.nextn.eh_proj_s); + cb(cur, "mtp_eh_proj", il); + + ggml_tensor * inpSA = cur; + + cur = build_norm(cur, layer.attn_norm, nullptr, LLM_NORM_RMS, il); + cb(cur, "mtp_attn_norm", il); + + // self-attention: dense MLA, same construction as the deepseek2 trunk graph + { + ggml_tensor * q = ggml_mul_mat(ctx0, layer.wq_a, cur); + cb(q, "mtp_q", il); + + q = build_norm(q, layer.attn_q_a_norm, nullptr, LLM_NORM_RMS, il); + cb(q, "mtp_q", il); + + q = ggml_mul_mat(ctx0, layer.wq_b, q); + cb(q, "mtp_q", il); + + // split into {n_embd_head_qk_nope, n_head, n_tokens} + ggml_tensor * q_nope = + ggml_view_3d(ctx0, q, n_embd_head_qk_nope, n_head, n_tokens, ggml_row_size(q->type, n_embd_head_k), + ggml_row_size(q->type, n_embd_head_k) * n_head, 0); + cb(q_nope, "mtp_q_nope", il); + + // and {n_embd_head_qk_rope, n_head, n_tokens} + ggml_tensor * q_pe = ggml_view_3d( + ctx0, q, n_embd_head_qk_rope, n_head, n_tokens, ggml_row_size(q->type, n_embd_head_k), + ggml_row_size(q->type, n_embd_head_k) * n_head, ggml_row_size(q->type, n_embd_head_qk_nope)); + cb(q_pe, "mtp_q_pe", il); + + ggml_tensor * kv_cmpr_pe = ggml_mul_mat(ctx0, layer.wkv_a_mqa, cur); + cb(kv_cmpr_pe, "mtp_kv_cmpr_pe", il); + + // split into {kv_lora_rank, n_tokens} + ggml_tensor * kv_cmpr = + ggml_view_2d(ctx0, kv_cmpr_pe, kv_lora_rank, n_tokens, + ggml_row_size(kv_cmpr_pe->type, kv_lora_rank + n_embd_head_qk_rope), 0); + cb(kv_cmpr, "mtp_kv_cmpr", il); + + // and {n_embd_head_qk_rope, 1, n_tokens} + ggml_tensor * k_pe = ggml_view_3d(ctx0, kv_cmpr_pe, n_embd_head_qk_rope, 1, n_tokens, + ggml_row_size(kv_cmpr_pe->type, kv_lora_rank + n_embd_head_qk_rope), + ggml_row_size(kv_cmpr_pe->type, kv_lora_rank + n_embd_head_qk_rope), + ggml_row_size(kv_cmpr_pe->type, kv_lora_rank)); + cb(k_pe, "mtp_k_pe", il); + + q_pe = ggml_rope_ext(ctx0, q_pe, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow); + cb(q_pe, "mtp_q_pe", il); + + k_pe = ggml_rope_ext(ctx0, k_pe, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow); + cb(k_pe, "mtp_k_pe", il); + + kv_cmpr = build_norm(kv_cmpr, layer.attn_kv_a_norm, nullptr, LLM_NORM_RMS, il); + cb(kv_cmpr, "mtp_kv_cmpr", il); + + // {n_embd_head_qk_nope, n_tokens, n_head} + q_nope = ggml_permute(ctx0, q_nope, 0, 2, 1, 3); + cb(q_nope, "mtp_q_nope_perm", il); + + // {n_embd_head_qk_nope, kv_lora_rank, n_head} x {n_embd_head_qk_nope, n_tokens, n_head} + ggml_tensor * q_nope_absorbed = ggml_mul_mat(ctx0, layer.wk_b, q_nope); + cb(q_nope_absorbed, "mtp_q_nope_absorbed", il); + + // {kv_lora_rank, n_head, n_tokens} + q_nope_absorbed = ggml_permute(ctx0, q_nope_absorbed, 0, 2, 1, 3); + cb(q_nope_absorbed, "mtp_q_nope_absorbed_perm", il); + + // {n_embd_head_qk_rope + kv_lora_rank, n_head, n_tokens} + // note: rope must go first for in-place context shifting in build_rope_shift() + ggml_tensor * Qcur = ggml_concat(ctx0, q_nope_absorbed, q_pe, 0); + cb(Qcur, "mtp_Qcur", il); + + kv_cmpr = ggml_reshape_3d(ctx0, kv_cmpr, kv_lora_rank, 1, n_tokens); + cb(kv_cmpr, "mtp_kv_cmpr_reshape", il); + + // {n_embd_head_qk_rope + kv_lora_rank, 1, n_tokens} + ggml_tensor * Kcur = ggml_concat(ctx0, kv_cmpr, k_pe, 0); + cb(Kcur, "mtp_Kcur", il); + + // {kv_lora_rank, 1, n_tokens} + ggml_tensor * Vcur = kv_cmpr; + cb(Vcur, "mtp_Vcur", il); + + // note: MLA with the absorption optimization converts into MQA (ie: GQA with 1 group) + cur = build_attn(inp_attn, + layer.wo, NULL, layer.wo_s, + Qcur, Kcur, Vcur, nullptr, nullptr, layer.wv_b, kq_scale, il); + cb(cur, "mtp_attn_out", il); + } + + ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpSA); + cb(ffn_inp, "mtp_ffn_inp", il); + + cur = build_norm(ffn_inp, layer.ffn_norm, NULL, LLM_NORM_RMS, il); + cb(cur, "mtp_ffn_norm", il); + + // MoE FFN with shared expert - same construction as the deepseek2 trunk graph + ggml_tensor * moe_out = build_moe_ffn(cur, + layer.ffn_gate_inp, + layer.ffn_up_exps, + layer.ffn_gate_exps, + layer.ffn_down_exps, + layer.ffn_exp_probs_b, + n_expert, n_expert_used, + LLM_FFN_SILU, hparams.expert_weights_norm, + hparams.expert_weights_scale, + (llama_expert_gating_func_type) hparams.expert_gating_func, + il, + nullptr, + layer.ffn_gate_up_exps, + layer.ffn_up_exps_s, + layer.ffn_gate_exps_s, + layer.ffn_down_exps_s); + cb(moe_out, "mtp_ffn_moe_out", il); + + // FFN shared expert + ggml_tensor * ffn_shexp = + build_ffn(cur, + layer.ffn_up_shexp, NULL, layer.ffn_up_shexp_s, + layer.ffn_gate_shexp, NULL, layer.ffn_gate_shexp_s, + layer.ffn_down_shexp, NULL, layer.ffn_down_shexp_s, + NULL, LLM_FFN_SILU, LLM_FFN_PAR, il); + cb(ffn_shexp, "mtp_ffn_shexp", il); + + cur = ggml_add(ctx0, moe_out, ffn_shexp); + cb(cur, "mtp_ffn_out", il); + + cur = ggml_add(ctx0, cur, ffn_inp); + cb(cur, "mtp_post_ffn", il); + + // shared_head_norm applied after the decoder block, before the shared LM head. + // The post-norm hidden state seeds the next MTP step. + ggml_tensor * head_norm_w = layer.nextn.shared_head_norm + ? layer.nextn.shared_head_norm + : model.output_norm; + GGML_ASSERT(head_norm_w && "GLM_DSA MTP: missing both nextn.shared_head_norm and output_norm"); + cur = build_norm(cur, head_norm_w, nullptr, LLM_NORM_RMS, -1); + + cb(cur, "h_nextn", -1); + res->t_h_nextn = cur; + + cur = ggml_get_rows(ctx0, cur, inp_out_ids); + cb(cur, "mtp_shared_head_norm", -1); + + ggml_tensor * head_w = layer.nextn.shared_head_head ? layer.nextn.shared_head_head : model.output; + ggml_tensor * head_s = layer.nextn.shared_head_head ? layer.nextn.shared_head_head_s : model.output_s; + GGML_ASSERT(head_w && "GLM_DSA MTP: missing LM head (nextn.shared_head_head or model.output)"); + cur = build_lora_mm(head_w, cur, head_s); + cb(cur, "result_output", -1); + + res->t_logits = cur; + ggml_build_forward_expand(gf, cur); +} diff --git a/src/models/models.h b/src/models/models.h index 92ebfafa1e29..c73136f3bdcc 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -1237,6 +1237,10 @@ struct llama_model_glm_dsa : public llama_model_base { graph(const llama_model & model, const llm_graph_params & params); }; + struct graph_mtp : public llm_graph_context { + graph_mtp(const llama_model & model, const llm_graph_params & params); + }; + std::unique_ptr build_arch_graph(const llm_graph_params & params) const override; }; From 60bccc3763395e01b039aa1ddeacc8cc0ea69f70 Mon Sep 17 00:00:00 2001 From: Geramy Loveless Date: Tue, 28 Jul 2026 23:43:45 -0700 Subject: [PATCH 047/190] add rdna3.5, and 3 to mmq configs so they can be tuned independently. (#26199) --- ggml/src/ggml-cuda/mmq-config-rdna3-5.cuh | 278 ++++++++++++++++++++ ggml/src/ggml-cuda/mmq-config-rdna3.cuh | 278 ++++++++++++++++++++ ggml/src/ggml-cuda/mmq-config-rdna4.cuh | 298 +++++++++++----------- ggml/src/ggml-cuda/mmq.cuh | 16 +- 4 files changed, 717 insertions(+), 153 deletions(-) create mode 100644 ggml/src/ggml-cuda/mmq-config-rdna3-5.cuh create mode 100644 ggml/src/ggml-cuda/mmq-config-rdna3.cuh diff --git a/ggml/src/ggml-cuda/mmq-config-rdna3-5.cuh b/ggml/src/ggml-cuda/mmq-config-rdna3-5.cuh new file mode 100644 index 000000000000..e420a32f0c41 --- /dev/null +++ b/ggml/src/ggml-cuda/mmq-config-rdna3-5.cuh @@ -0,0 +1,278 @@ +static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config_rdna3_5(ggml_type type, int J, bool fallback) { + CASE(GGML_TYPE_Q1_0, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q1_0, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q1_0, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q1_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q1_0, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q1_0, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q1_0, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q1_0, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q1_0, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q1_0, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q1_0, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q1_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_Q4_0, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_0, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_0, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_0, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_0, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_0, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_0, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_0, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_0, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_0, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_Q4_1, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_1, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_1, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_1, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_1, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_1, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_1, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_1, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_1, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_1, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_1, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_1, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_Q5_0, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_0, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_0, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_0, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_0, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_0, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_0, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_0, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_0, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_0, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_Q5_1, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_1, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_1, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_1, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_1, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_1, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_1, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_1, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_1, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_1, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_1, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_1, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_Q8_0, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q8_0, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q8_0, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q8_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q8_0, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q8_0, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q8_0, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q8_0, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q8_0, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q8_0, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q8_0, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q8_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + +// --------------------------------------------------------------------------------------------- + + CASE(GGML_TYPE_Q2_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q2_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q2_K, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q2_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_K, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_K, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_K, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_Q3_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q3_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q3_K, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q3_K, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q3_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q3_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q3_K, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q3_K, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q3_K, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q3_K, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q3_K, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q3_K, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_Q4_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_K, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_K, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_K, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_K, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_K, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_K, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_K, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_K, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_Q5_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_K, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_K, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_K, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_K, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_K, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_K, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_K, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_K, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_Q6_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q6_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q6_K, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q6_K, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q6_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q6_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q6_K, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q6_K, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q6_K, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q6_K, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q6_K, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q6_K, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); + +// --------------------------------------------------------------------------------------------- + + CASE(GGML_TYPE_IQ1_S, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ1_S, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ1_S, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ1_S, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ1_S, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ1_S, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ1_S, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ1_S, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ1_S, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ1_S, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ1_S, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ1_S, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_IQ2_XXS, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_XXS, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_XXS, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XXS, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_IQ2_XS, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_XS, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_XS, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_XS, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_XS, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XS, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XS, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XS, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XS, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XS, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XS, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XS, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_IQ2_S, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_S, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_S, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_S, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_S, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_S, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_S, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_S, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_S, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_S, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_S, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_S, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_IQ3_XXS, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ3_XXS, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ3_XXS, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_XXS, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_IQ3_S, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ3_S, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ3_S, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ3_S, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ3_S, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_S, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_S, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_S, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_S, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_S, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_S, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_S, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_IQ4_XS, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ4_XS, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ4_XS, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ4_XS, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ4_XS, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_XS, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_XS, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_XS, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_XS, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_XS, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_XS, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_XS, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_IQ4_NL, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ4_NL, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ4_NL, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ4_NL, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ4_NL, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_NL, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_NL, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_NL, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_NL, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_NL, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_NL, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_NL, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + +// --------------------------------------------------------------------------------------------- + + CASE(GGML_TYPE_MXFP4, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_MXFP4, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_MXFP4, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_MXFP4, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_MXFP4, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_MXFP4, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_MXFP4, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_MXFP4, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_MXFP4, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_MXFP4, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_MXFP4, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_MXFP4, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_NVFP4, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_NVFP4, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_NVFP4, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_NVFP4, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_NVFP4, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_NVFP4, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_NVFP4, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_NVFP4, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_NVFP4, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_NVFP4, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_NVFP4, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_NVFP4, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false); + + return ggml_cuda_mmq_config(GGML_TYPE_COUNT, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, 256, false, true); +} diff --git a/ggml/src/ggml-cuda/mmq-config-rdna3.cuh b/ggml/src/ggml-cuda/mmq-config-rdna3.cuh new file mode 100644 index 000000000000..12262306090e --- /dev/null +++ b/ggml/src/ggml-cuda/mmq-config-rdna3.cuh @@ -0,0 +1,278 @@ +static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config_rdna3(ggml_type type, int J, bool fallback) { + CASE(GGML_TYPE_Q1_0, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q1_0, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q1_0, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q1_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q1_0, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q1_0, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q1_0, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q1_0, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q1_0, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q1_0, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q1_0, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q1_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_Q4_0, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_0, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_0, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_0, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_0, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_0, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_0, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_0, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_0, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_0, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_Q4_1, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_1, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_1, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_1, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_1, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_1, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_1, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_1, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_1, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_1, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_1, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_1, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_Q5_0, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_0, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_0, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_0, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_0, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_0, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_0, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_0, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_0, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_0, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_Q5_1, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_1, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_1, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_1, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_1, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_1, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_1, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_1, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_1, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_1, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_1, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_1, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_Q8_0, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q8_0, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q8_0, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q8_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q8_0, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q8_0, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q8_0, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q8_0, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q8_0, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q8_0, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q8_0, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q8_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + +// --------------------------------------------------------------------------------------------- + + CASE(GGML_TYPE_Q2_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q2_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q2_K, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q2_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_K, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_K, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_K, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_Q3_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q3_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q3_K, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q3_K, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q3_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q3_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q3_K, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q3_K, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q3_K, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q3_K, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q3_K, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q3_K, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_Q4_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_K, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_K, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_K, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_K, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_K, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_K, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_K, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_K, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_Q5_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_K, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_K, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_K, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_K, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_K, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_K, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_K, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_K, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_Q6_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q6_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q6_K, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q6_K, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q6_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q6_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q6_K, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q6_K, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q6_K, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q6_K, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q6_K, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q6_K, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); + +// --------------------------------------------------------------------------------------------- + + CASE(GGML_TYPE_IQ1_S, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ1_S, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ1_S, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ1_S, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ1_S, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ1_S, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ1_S, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ1_S, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ1_S, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ1_S, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ1_S, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ1_S, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_IQ2_XXS, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_XXS, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_XXS, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XXS, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_IQ2_XS, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_XS, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_XS, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_XS, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_XS, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XS, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XS, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XS, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XS, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XS, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XS, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XS, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_IQ2_S, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_S, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_S, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_S, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_S, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_S, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_S, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_S, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_S, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_S, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_S, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_S, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_IQ3_XXS, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ3_XXS, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ3_XXS, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_XXS, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_IQ3_S, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ3_S, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ3_S, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ3_S, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ3_S, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_S, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_S, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_S, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_S, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_S, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_S, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_S, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_IQ4_XS, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ4_XS, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ4_XS, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ4_XS, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ4_XS, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_XS, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_XS, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_XS, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_XS, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_XS, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_XS, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_XS, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_IQ4_NL, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ4_NL, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ4_NL, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ4_NL, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ4_NL, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_NL, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_NL, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_NL, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_NL, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_NL, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_NL, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_NL, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + +// --------------------------------------------------------------------------------------------- + + CASE(GGML_TYPE_MXFP4, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_MXFP4, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_MXFP4, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_MXFP4, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_MXFP4, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_MXFP4, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_MXFP4, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_MXFP4, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_MXFP4, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_MXFP4, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_MXFP4, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_MXFP4, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_NVFP4, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_NVFP4, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_NVFP4, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_NVFP4, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_NVFP4, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_NVFP4, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_NVFP4, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_NVFP4, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_NVFP4, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_NVFP4, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_NVFP4, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_NVFP4, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false); + + return ggml_cuda_mmq_config(GGML_TYPE_COUNT, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, 256, false, true); +} diff --git a/ggml/src/ggml-cuda/mmq-config-rdna4.cuh b/ggml/src/ggml-cuda/mmq-config-rdna4.cuh index 6280e80ee4ce..a224ecafc37d 100644 --- a/ggml/src/ggml-cuda/mmq-config-rdna4.cuh +++ b/ggml/src/ggml-cuda/mmq-config-rdna4.cuh @@ -1,77 +1,77 @@ static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config_rdna4(ggml_type type, int J, bool fallback) { - CASE(GGML_TYPE_Q1_0, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); - CASE(GGML_TYPE_Q1_0, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); - CASE(GGML_TYPE_Q1_0, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q1_0, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q1_0, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q1_0, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); CASE(GGML_TYPE_Q1_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); - CASE(GGML_TYPE_Q1_0, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_Q1_0, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_Q1_0, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_Q1_0, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q1_0, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q1_0, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q1_0, 128, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q1_0, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); CASE(GGML_TYPE_Q1_0, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); CASE(GGML_TYPE_Q1_0, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); CASE(GGML_TYPE_Q1_0, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); CASE(GGML_TYPE_Q1_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_Q4_0, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); - CASE(GGML_TYPE_Q4_0, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); - CASE(GGML_TYPE_Q4_0, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_0, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_0, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_0, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); CASE(GGML_TYPE_Q4_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); - CASE(GGML_TYPE_Q4_0, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_Q4_0, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_Q4_0, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_Q4_0, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_0, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_0, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_0, 128, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_0, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); CASE(GGML_TYPE_Q4_0, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); CASE(GGML_TYPE_Q4_0, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); CASE(GGML_TYPE_Q4_0, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); CASE(GGML_TYPE_Q4_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_Q4_1, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); - CASE(GGML_TYPE_Q4_1, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); - CASE(GGML_TYPE_Q4_1, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_1, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_1, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_1, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); CASE(GGML_TYPE_Q4_1, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); - CASE(GGML_TYPE_Q4_1, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_Q4_1, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_Q4_1, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_Q4_1, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_1, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_1, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_1, 128, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_1, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); CASE(GGML_TYPE_Q4_1, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); CASE(GGML_TYPE_Q4_1, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); CASE(GGML_TYPE_Q4_1, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); CASE(GGML_TYPE_Q4_1, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_Q5_0, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); - CASE(GGML_TYPE_Q5_0, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); - CASE(GGML_TYPE_Q5_0, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_0, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_0, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_0, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); CASE(GGML_TYPE_Q5_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); - CASE(GGML_TYPE_Q5_0, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_Q5_0, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_Q5_0, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_Q5_0, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_0, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_0, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_0, 128, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_0, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); CASE(GGML_TYPE_Q5_0, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); CASE(GGML_TYPE_Q5_0, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); CASE(GGML_TYPE_Q5_0, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); CASE(GGML_TYPE_Q5_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_Q5_1, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); - CASE(GGML_TYPE_Q5_1, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); - CASE(GGML_TYPE_Q5_1, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_1, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_1, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_1, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); CASE(GGML_TYPE_Q5_1, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); - CASE(GGML_TYPE_Q5_1, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_Q5_1, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_Q5_1, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_Q5_1, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_1, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_1, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_1, 128, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_1, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); CASE(GGML_TYPE_Q5_1, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); CASE(GGML_TYPE_Q5_1, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); CASE(GGML_TYPE_Q5_1, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); CASE(GGML_TYPE_Q5_1, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_Q8_0, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); - CASE(GGML_TYPE_Q8_0, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); - CASE(GGML_TYPE_Q8_0, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q8_0, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q8_0, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q8_0, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); CASE(GGML_TYPE_Q8_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); - CASE(GGML_TYPE_Q8_0, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_Q8_0, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_Q8_0, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_Q8_0, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q8_0, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q8_0, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q8_0, 128, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q8_0, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); CASE(GGML_TYPE_Q8_0, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); CASE(GGML_TYPE_Q8_0, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); CASE(GGML_TYPE_Q8_0, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); @@ -79,66 +79,62 @@ static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_conf // --------------------------------------------------------------------------------------------- - CASE(GGML_TYPE_Q2_K, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true); - CASE(GGML_TYPE_Q2_K, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true); - CASE(GGML_TYPE_Q2_K, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true); - CASE(GGML_TYPE_Q2_K, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true); - CASE(GGML_TYPE_Q2_K, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_Q2_K, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_Q2_K, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_Q2_K, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q2_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q2_K, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q2_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_K, 128, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_K, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false); CASE(GGML_TYPE_Q2_K, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_Q2_K, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_Q2_K, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_Q2_K, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_Q3_K, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); - CASE(GGML_TYPE_Q3_K, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); - CASE(GGML_TYPE_Q3_K, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q3_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q3_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q3_K, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); CASE(GGML_TYPE_Q3_K, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); - CASE(GGML_TYPE_Q3_K, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_Q3_K, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_Q3_K, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_Q3_K, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q3_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q3_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q3_K, 128, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q3_K, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); CASE(GGML_TYPE_Q3_K, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); CASE(GGML_TYPE_Q3_K, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); CASE(GGML_TYPE_Q3_K, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); CASE(GGML_TYPE_Q3_K, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_Q4_K, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); - CASE(GGML_TYPE_Q4_K, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); - CASE(GGML_TYPE_Q4_K, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_K, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); CASE(GGML_TYPE_Q4_K, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); - CASE(GGML_TYPE_Q4_K, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_Q4_K, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_Q4_K, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_Q4_K, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_K, 128, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_K, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); CASE(GGML_TYPE_Q4_K, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); CASE(GGML_TYPE_Q4_K, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); CASE(GGML_TYPE_Q4_K, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); CASE(GGML_TYPE_Q4_K, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_Q5_K, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); - CASE(GGML_TYPE_Q5_K, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); - CASE(GGML_TYPE_Q5_K, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_K, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); CASE(GGML_TYPE_Q5_K, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); - CASE(GGML_TYPE_Q5_K, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_Q5_K, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_Q5_K, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_Q5_K, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_K, 128, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_K, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); CASE(GGML_TYPE_Q5_K, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); CASE(GGML_TYPE_Q5_K, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); CASE(GGML_TYPE_Q5_K, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); CASE(GGML_TYPE_Q5_K, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_Q6_K, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true); - CASE(GGML_TYPE_Q6_K, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true); - CASE(GGML_TYPE_Q6_K, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q6_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q6_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q6_K, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true); CASE(GGML_TYPE_Q6_K, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true); - CASE(GGML_TYPE_Q6_K, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_Q6_K, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_Q6_K, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_Q6_K, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q6_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q6_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q6_K, 128, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q6_K, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); CASE(GGML_TYPE_Q6_K, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); CASE(GGML_TYPE_Q6_K, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); CASE(GGML_TYPE_Q6_K, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); @@ -146,105 +142,105 @@ static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_conf // --------------------------------------------------------------------------------------------- - CASE(GGML_TYPE_IQ1_S, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); - CASE(GGML_TYPE_IQ1_S, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); - CASE(GGML_TYPE_IQ1_S, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ1_S, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ1_S, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ1_S, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); CASE(GGML_TYPE_IQ1_S, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); - CASE(GGML_TYPE_IQ1_S, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_IQ1_S, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_IQ1_S, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_IQ1_S, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ1_S, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ1_S, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ1_S, 128, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ1_S, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); CASE(GGML_TYPE_IQ1_S, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); CASE(GGML_TYPE_IQ1_S, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); CASE(GGML_TYPE_IQ1_S, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); CASE(GGML_TYPE_IQ1_S, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); - CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); - CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_XXS, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_XXS, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_XXS, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); - CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XXS, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XXS, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XXS, 128, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XXS, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_IQ2_XS, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); - CASE(GGML_TYPE_IQ2_XS, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); - CASE(GGML_TYPE_IQ2_XS, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_XS, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_XS, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_XS, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); CASE(GGML_TYPE_IQ2_XS, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); - CASE(GGML_TYPE_IQ2_XS, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_IQ2_XS, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_IQ2_XS, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_IQ2_XS, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XS, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XS, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XS, 128, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XS, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); CASE(GGML_TYPE_IQ2_XS, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); CASE(GGML_TYPE_IQ2_XS, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); CASE(GGML_TYPE_IQ2_XS, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); CASE(GGML_TYPE_IQ2_XS, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_IQ2_S, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); - CASE(GGML_TYPE_IQ2_S, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); - CASE(GGML_TYPE_IQ2_S, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_S, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_S, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_S, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); CASE(GGML_TYPE_IQ2_S, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); - CASE(GGML_TYPE_IQ2_S, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_IQ2_S, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_IQ2_S, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_IQ2_S, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_S, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_S, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_S, 128, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_S, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); CASE(GGML_TYPE_IQ2_S, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); CASE(GGML_TYPE_IQ2_S, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); CASE(GGML_TYPE_IQ2_S, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); CASE(GGML_TYPE_IQ2_S, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); - CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); - CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ3_XXS, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ3_XXS, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ3_XXS, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); - CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_XXS, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_XXS, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_XXS, 128, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_XXS, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_IQ3_S, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); - CASE(GGML_TYPE_IQ3_S, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); - CASE(GGML_TYPE_IQ3_S, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ3_S, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ3_S, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ3_S, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); CASE(GGML_TYPE_IQ3_S, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); - CASE(GGML_TYPE_IQ3_S, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_IQ3_S, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_IQ3_S, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_IQ3_S, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_S, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_S, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_S, 128, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_S, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); CASE(GGML_TYPE_IQ3_S, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); CASE(GGML_TYPE_IQ3_S, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); CASE(GGML_TYPE_IQ3_S, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); CASE(GGML_TYPE_IQ3_S, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_IQ4_XS, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); - CASE(GGML_TYPE_IQ4_XS, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); - CASE(GGML_TYPE_IQ4_XS, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ4_XS, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ4_XS, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ4_XS, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); CASE(GGML_TYPE_IQ4_XS, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); - CASE(GGML_TYPE_IQ4_XS, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_IQ4_XS, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_IQ4_XS, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_IQ4_XS, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_XS, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_XS, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_XS, 128, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_XS, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); CASE(GGML_TYPE_IQ4_XS, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); CASE(GGML_TYPE_IQ4_XS, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); CASE(GGML_TYPE_IQ4_XS, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); CASE(GGML_TYPE_IQ4_XS, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_IQ4_NL, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); - CASE(GGML_TYPE_IQ4_NL, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); - CASE(GGML_TYPE_IQ4_NL, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ4_NL, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ4_NL, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ4_NL, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); CASE(GGML_TYPE_IQ4_NL, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); - CASE(GGML_TYPE_IQ4_NL, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_IQ4_NL, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_IQ4_NL, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_IQ4_NL, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_NL, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_NL, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_NL, 128, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_NL, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); CASE(GGML_TYPE_IQ4_NL, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); CASE(GGML_TYPE_IQ4_NL, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); CASE(GGML_TYPE_IQ4_NL, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); @@ -252,27 +248,27 @@ static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_conf // --------------------------------------------------------------------------------------------- - CASE(GGML_TYPE_MXFP4, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); - CASE(GGML_TYPE_MXFP4, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); - CASE(GGML_TYPE_MXFP4, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_MXFP4, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_MXFP4, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_MXFP4, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); CASE(GGML_TYPE_MXFP4, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); - CASE(GGML_TYPE_MXFP4, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_MXFP4, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_MXFP4, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_MXFP4, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_MXFP4, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_MXFP4, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_MXFP4, 128, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_MXFP4, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); CASE(GGML_TYPE_MXFP4, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); CASE(GGML_TYPE_MXFP4, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); CASE(GGML_TYPE_MXFP4, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); CASE(GGML_TYPE_MXFP4, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_NVFP4, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true); - CASE(GGML_TYPE_NVFP4, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true); - CASE(GGML_TYPE_NVFP4, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_NVFP4, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_NVFP4, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_NVFP4, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true); CASE(GGML_TYPE_NVFP4, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true); - CASE(GGML_TYPE_NVFP4, 256, 2, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_NVFP4, 256, 2, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_NVFP4, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_NVFP4, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_NVFP4, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_NVFP4, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_NVFP4, 128, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_NVFP4, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false); CASE(GGML_TYPE_NVFP4, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false); CASE(GGML_TYPE_NVFP4, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false); CASE(GGML_TYPE_NVFP4, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false); diff --git a/ggml/src/ggml-cuda/mmq.cuh b/ggml/src/ggml-cuda/mmq.cuh index 71e3b2647a8e..5f30f5f6bc6f 100644 --- a/ggml/src/ggml-cuda/mmq.cuh +++ b/ggml/src/ggml-cuda/mmq.cuh @@ -218,6 +218,8 @@ struct ggml_cuda_mmq_config { #include "mmq-config-cdna.cuh" #include "mmq-config-rdna2.cuh" +#include "mmq-config-rdna3.cuh" +#include "mmq-config-rdna3-5.cuh" #include "mmq-config-rdna4.cuh" #undef CASE @@ -227,9 +229,15 @@ static __host__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config(const ggml_type ty if (GGML_CUDA_CC_IS_CDNA(cc)) { return ggml_cuda_mmq_get_config_cdna(type, J, fallback); } - if (amd_wmma_available(cc)) { + if (GGML_CUDA_CC_IS_RDNA4(cc)) { return ggml_cuda_mmq_get_config_rdna4(type, J, fallback); } + if (GGML_CUDA_CC_IS_RDNA3_5(cc)) { + return ggml_cuda_mmq_get_config_rdna3_5(type, J, fallback); + } + if (GGML_CUDA_CC_IS_RDNA3(cc)) { // covers RDNA 3.0 + return ggml_cuda_mmq_get_config_rdna3(type, J, fallback); + } return ggml_cuda_mmq_get_config_rdna2(type, J, fallback); } if (blackwell_mma_available(cc)) { @@ -245,8 +253,12 @@ static constexpr __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config(ggml_t #ifdef GGML_USE_HIP #ifdef CDNA return ggml_cuda_mmq_get_config_cdna(type, J, fallback); -#elif defined(AMD_WMMA_AVAILABLE) +#elif defined(RDNA4) return ggml_cuda_mmq_get_config_rdna4(type, J, fallback); +#elif defined(RDNA3_5) + return ggml_cuda_mmq_get_config_rdna3_5(type, J, fallback); +#elif defined(RDNA3) + return ggml_cuda_mmq_get_config_rdna3(type, J, fallback); #else return ggml_cuda_mmq_get_config_rdna2(type, J, fallback); #endif // CDNA From f5b9bd39b56c7a7839a9795a100b6a00b84ac961 Mon Sep 17 00:00:00 2001 From: Aman Gupta Date: Wed, 29 Jul 2026 15:04:30 +0800 Subject: [PATCH 048/190] RPC: add tensor_memset (#25912) --- ggml/include/ggml-rpc.h | 4 +- ggml/src/ggml-rpc/ggml-rpc.cpp | 83 +++++++++++++++++++++++++++++++++- 2 files changed, 84 insertions(+), 3 deletions(-) diff --git a/ggml/include/ggml-rpc.h b/ggml/include/ggml-rpc.h index 16ca33947a2e..276aea00ea1b 100644 --- a/ggml/include/ggml-rpc.h +++ b/ggml/include/ggml-rpc.h @@ -6,9 +6,9 @@ extern "C" { #endif -#define RPC_PROTO_MAJOR_VERSION 4 +#define RPC_PROTO_MAJOR_VERSION 5 #define RPC_PROTO_MINOR_VERSION 0 -#define RPC_PROTO_PATCH_VERSION 3 +#define RPC_PROTO_PATCH_VERSION 0 #ifdef __cplusplus static_assert(GGML_OP_COUNT == 101, "GGML_OP_COUNT has changed - update RPC_PROTO_PATCH_VERSION"); diff --git a/ggml/src/ggml-rpc/ggml-rpc.cpp b/ggml/src/ggml-rpc/ggml-rpc.cpp index d38057721834..17c53a5f049e 100644 --- a/ggml/src/ggml-rpc/ggml-rpc.cpp +++ b/ggml/src/ggml-rpc/ggml-rpc.cpp @@ -71,6 +71,7 @@ enum rpc_cmd { RPC_CMD_HELLO, RPC_CMD_DEVICE_COUNT, RPC_CMD_GRAPH_RECOMPUTE, + RPC_CMD_MEMSET_TENSOR, RPC_CMD_COUNT, }; @@ -152,6 +153,13 @@ struct rpc_msg_buffer_clear_req { uint8_t value; }; +struct rpc_msg_memset_tensor_req { + rpc_tensor tensor; + uint64_t offset; + uint64_t size; + uint8_t value; +}; + struct rpc_msg_set_tensor_hash_req { rpc_tensor tensor; uint64_t offset; @@ -462,6 +470,19 @@ static enum ggml_status ggml_backend_rpc_buffer_init_tensor(ggml_backend_buffer_ return GGML_STATUS_SUCCESS; } +static void ggml_backend_rpc_buffer_memset_tensor( + ggml_backend_buffer_t buffer, ggml_tensor * tensor, uint8_t value, size_t offset, size_t size) { + ggml_backend_rpc_buffer_context * ctx = (ggml_backend_rpc_buffer_context *)buffer->context; + rpc_msg_memset_tensor_req request = { + /* .tensor = */ serialize_tensor(tensor), + /* .offset = */ offset, + /* .size = */ size, + /* .value = */ value, + }; + bool status = send_rpc_cmd(ctx->sock, RPC_CMD_MEMSET_TENSOR, &request, sizeof(request), nullptr, 0); + RPC_STATUS_ASSERT(status); +} + static void ggml_backend_rpc_buffer_set_tensor(ggml_backend_buffer_t buffer, ggml_tensor * tensor, const void * data, size_t offset, size_t size) { ggml_backend_rpc_buffer_context * ctx = (ggml_backend_rpc_buffer_context *)buffer->context; rpc_tensor rpc_tensor = serialize_tensor(tensor); @@ -531,7 +552,7 @@ static ggml_backend_buffer_i ggml_backend_rpc_buffer_interface = { /* .free_buffer = */ ggml_backend_rpc_buffer_free_buffer, /* .get_base = */ ggml_backend_rpc_buffer_get_base, /* .init_tensor = */ ggml_backend_rpc_buffer_init_tensor, - /* .memset_tensor = */ NULL, + /* .memset_tensor = */ ggml_backend_rpc_buffer_memset_tensor, /* .set_tensor = */ ggml_backend_rpc_buffer_set_tensor, /* .get_tensor = */ ggml_backend_rpc_buffer_get_tensor, /* .set_tensor_2d = */ NULL, @@ -831,6 +852,7 @@ class rpc_server { bool buffer_get_base(const rpc_msg_buffer_get_base_req & request, rpc_msg_buffer_get_base_rsp & response); bool free_buffer(const rpc_msg_free_buffer_req & request); bool buffer_clear(const rpc_msg_buffer_clear_req & request); + bool memset_tensor(const rpc_msg_memset_tensor_req & request); bool set_tensor(const std::vector & input); bool set_tensor_hash(const rpc_msg_set_tensor_hash_req & request, rpc_msg_set_tensor_hash_rsp & response); bool get_tensor(const rpc_msg_get_tensor_req & request, std::vector & response); @@ -989,6 +1011,52 @@ bool rpc_server::buffer_clear(const rpc_msg_buffer_clear_req & request) { return true; } +bool rpc_server::memset_tensor(const rpc_msg_memset_tensor_req & request) { + struct ggml_init_params params { + /*.mem_size =*/ ggml_tensor_overhead(), + /*.mem_buffer =*/ NULL, + /*.no_alloc =*/ true, + }; + ggml_context_ptr ctx_ptr { ggml_init(params) }; + GGML_ASSERT(ctx_ptr != nullptr); + ggml_context * ctx = ctx_ptr.get(); + ggml_tensor * tensor = deserialize_tensor(ctx, &request.tensor); + if (tensor == nullptr || tensor->buffer == nullptr) { + GGML_LOG_ERROR("[%s] error deserializing tensor\n", __func__); + return false; + } + + const uint64_t tensor_size = ggml_nbytes(tensor); + if (request.offset > tensor_size || request.size > tensor_size - request.offset) { + GGML_LOG_ERROR("[%s] tensor region (offset=%" PRIu64 ", size=%" PRIu64 ") out of tensor bounds [0, %" PRIu64 ")\n", + __func__, request.offset, request.size, tensor_size); + return false; + } + + const uint64_t buffer_start = (uint64_t) ggml_backend_buffer_get_base(tensor->buffer); + const uint64_t buffer_size = ggml_backend_buffer_get_size(tensor->buffer); + if (request.tensor.data < buffer_start) { + GGML_LOG_ERROR("[%s] tensor data before buffer start\n", __func__); + return false; + } + const uint64_t data_offset = request.tensor.data - buffer_start; + if (data_offset > buffer_size || + request.offset > buffer_size - data_offset || + request.size > buffer_size - data_offset - request.offset) { + GGML_LOG_ERROR("[%s] tensor region out of buffer bounds\n", __func__); + return false; + } + if (tensor->buffer->iface.memset_tensor == nullptr) { + GGML_LOG_ERROR("[%s] memset not implemented by backend buffer\n", __func__); + return false; + } + + LOG_DBG("[%s] buffer: %p, data: %p, offset: %" PRIu64 ", size: %" PRIu64 ", value: %u\n", + __func__, (void *) tensor->buffer, tensor->data, request.offset, request.size, request.value); + ggml_backend_tensor_memset(tensor, request.value, request.offset, request.size); + return true; +} + ggml_tensor * rpc_server::deserialize_tensor(struct ggml_context * ctx, const rpc_tensor * tensor) { // Validate tensor type before using it if (tensor->type >= GGML_TYPE_COUNT) { @@ -1585,6 +1653,19 @@ static void rpc_serve_client(const std::vector & backends, const } break; } + case RPC_CMD_MEMSET_TENSOR: { + rpc_msg_memset_tensor_req request; + if (!recv_msg(sock, &request, sizeof(request))) { + return; + } + if (!server.memset_tensor(request)) { + return; + } + if (!send_msg(sock, nullptr, 0)) { + return; + } + break; + } case RPC_CMD_SET_TENSOR: { std::vector input; if (!recv_msg(sock, input)) { From e1af89a6815737a5db132eee23a94a8ee58553e0 Mon Sep 17 00:00:00 2001 From: Kaben Nanlohy Date: Wed, 29 Jul 2026 04:53:44 -0600 Subject: [PATCH 049/190] conversion: fix Qwen2.5-Omni mmproj conversion regression (#26262) --- conversion/qwenvl.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/conversion/qwenvl.py b/conversion/qwenvl.py index 7befd0c8d816..202a47961b3c 100644 --- a/conversion/qwenvl.py +++ b/conversion/qwenvl.py @@ -179,12 +179,12 @@ def set_gguf_parameters(self): def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None: name, gen = item - if not name.startswith("visual.") and not name.startswith("audio_tower."): - return None - if name.startswith("thinker."): name = name.replace("thinker.", "") + if not name.startswith("visual.") and not name.startswith("audio_tower."): + return None + if "audio_bos_eos_token" in name: # this tensor is left unused in transformers code # https://github.com/huggingface/transformers/blob/6e3063422c4b1c014aa60c32b9254fd2902f0f28/src/transformers/models/qwen2_5_omni/modular_qwen2_5_omni.py#L1809 From 992c325323f925cb82c86778e5e91a63de199063 Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Wed, 29 Jul 2026 14:59:44 +0300 Subject: [PATCH 050/190] server : add trace logging for slot similarity checking (#26271) Adds trace logging in server-context.cpp for slot similarity checking during prompt cache slot selection, including skip reasons and similarity calculation details. Assisted-by: llama.cpp:Qwen3.6-27B --- tools/server/server-context.cpp | 19 ++++++++++++------- tools/server/server-task.cpp | 14 +++++++------- 2 files changed, 19 insertions(+), 14 deletions(-) diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index dba15d426fd4..749bd9aac194 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -1537,7 +1537,7 @@ struct server_context_impl { // find the slot that has at least n% prompt similarity if (slot_prompt_similarity != 0.0f) { - float sim_best = 0; + float f_sim_best = 0; for (server_slot & slot : slots) { if (task.id_slot != -1 && slot.id != task.id_slot) { @@ -1546,6 +1546,7 @@ struct server_context_impl { // skip the slot if it is not available if (slot.is_processing()) { + SLT_TRC(slot, " - skipping, is_processing = %d\n", slot.is_processing()); continue; } @@ -1553,26 +1554,30 @@ struct server_context_impl { // skip the slot if it does not contains cached tokens if (tokens.empty()) { + SLT_TRC(slot, "%s", " - skipping, slot is empty\n"); continue; } // fraction of the Longest Common Prefix length with respect to the input prompt length - const float sim_cur = float(tokens.get_common_prefix(task.tokens)) / task.tokens.size(); + const size_t lcp_len = tokens.get_common_prefix(task.tokens); + const float f_sim_cur = float(lcp_len) / task.tokens.size(); + + SLT_TRC(slot, " - checking sim = %.3f (%zu/%zu) > %.3f\n", f_sim_cur, lcp_len, task.tokens.size(), slot_prompt_similarity); // select the current slot if the criteria match - if (sim_cur > sim_best && sim_cur > slot_prompt_similarity) { - sim_best = sim_cur; + if (f_sim_cur > f_sim_best && f_sim_cur > slot_prompt_similarity) { + f_sim_best = f_sim_cur; ret = &slot; } } if (ret != nullptr) { - const float f_keep = (sim_best*task.tokens.size()) / ret->prompt.tokens.size(); + const float f_keep = (f_sim_best*task.tokens.size()) / ret->prompt.tokens.size(); if (task.id_slot == -1) { - SLT_INF(*ret, "selected slot by LCP similarity, sim_best = %.3f (> %.3f thold), f_keep = %.3f\n", - sim_best, slot_prompt_similarity, f_keep); + SLT_INF(*ret, "selected slot by LCP similarity, f_sim_best = %.3f (> %.3f thold), f_keep = %.3f\n", + f_sim_best, slot_prompt_similarity, f_keep); } // if we are about to lose a large portion of the existing context - save it in the prompt cache diff --git a/tools/server/server-task.cpp b/tools/server/server-task.cpp index 99e63b05f544..070f1ade241c 100644 --- a/tools/server/server-task.cpp +++ b/tools/server/server-task.cpp @@ -1742,9 +1742,9 @@ bool server_prompt_cache::load(server_prompt & prompt, const server_tokens & tok const int lcp_best = prompt.tokens.get_common_prefix(tokens_new); float f_keep_best = prompt.tokens.size() > 0 ? float(lcp_best) / prompt.tokens.size() : -1.0f; // empty slot: any cache entry wins - float sim_best = float(lcp_best) / tokens_new.size(); + float f_sim_best = float(lcp_best) / tokens_new.size(); - SRV_TRC(" - looking for better prompt, base f_keep = %.3f, sim = %.3f\n", f_keep_best, sim_best); + SRV_TRC(" - looking for better prompt, base f_keep = %.3f, f_sim = %.3f\n", f_keep_best, f_sim_best); auto it_best = states.end(); @@ -1753,25 +1753,25 @@ bool server_prompt_cache::load(server_prompt & prompt, const server_tokens & tok const int lcp_cur = it->prompt.tokens.get_common_prefix(tokens_new); const float f_keep_cur = float(lcp_cur) / it->prompt.tokens.size(); - const float sim_cur = float(lcp_cur) / tokens_new.size(); + const float f_sim_cur = float(lcp_cur) / tokens_new.size(); - SRV_TRC(" - prompt with length %7zu, lcp = %7d, f_keep = %.3f, sim = %.3f\n", it->prompt.tokens.size(), lcp_cur, f_keep_cur, sim_cur); + SRV_TRC(" - prompt with length %7zu, lcp = %7d, f_keep = %.3f, f_sim = %.3f\n", it->prompt.tokens.size(), lcp_cur, f_keep_cur, f_sim_cur); // don't trash large prompts if (f_keep_cur < 0.25f) { continue; } - if (f_keep_best < f_keep_cur && sim_best < sim_cur) { + if (f_keep_best < f_keep_cur && f_sim_best < f_sim_cur) { f_keep_best = f_keep_cur; - sim_best = sim_cur; + f_sim_best = f_sim_cur; it_best = it; } } if (it_best != states.end()) { - SRV_TRC(" - found better prompt with f_keep = %.3f, sim = %.3f\n", f_keep_best, sim_best); + SRV_TRC(" - found better prompt with f_keep = %.3f, f_sim = %.3f\n", f_keep_best, f_sim_best); { auto & data = it_best->data.main; From e2f59ed71dacdbd665f406eabc43e8da51b30a90 Mon Sep 17 00:00:00 2001 From: "Alessandro de Oliveira Faria (A.K.A.CABELO)" Date: Wed, 29 Jul 2026 09:16:02 -0300 Subject: [PATCH 051/190] vendor: update BoringSSL to 0.20260728.0 (#26241) --- vendor/cpp-httplib/CMakeLists.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vendor/cpp-httplib/CMakeLists.txt b/vendor/cpp-httplib/CMakeLists.txt index 7460a9ce93ec..c6eb372b5f2c 100644 --- a/vendor/cpp-httplib/CMakeLists.txt +++ b/vendor/cpp-httplib/CMakeLists.txt @@ -41,7 +41,7 @@ if (LLAMA_BUILD_BORINGSSL) set(FIPS OFF CACHE BOOL "Enable FIPS (BoringSSL)") set(BORINGSSL_GIT "https://boringssl.googlesource.com/boringssl" CACHE STRING "BoringSSL git repository") - set(BORINGSSL_VERSION "0.20260713.0" CACHE STRING "BoringSSL version") + set(BORINGSSL_VERSION "0.20260728.0" CACHE STRING "BoringSSL version") message(STATUS "Fetching BoringSSL version ${BORINGSSL_VERSION}") From 11b068d06605288ce7917534b46d52b47823dc13 Mon Sep 17 00:00:00 2001 From: Titaniumtown Date: Wed, 29 Jul 2026 05:16:57 -0700 Subject: [PATCH 052/190] sycl: contiguous fast path + 32-bit index math for unary elementwise ops (#25946) * sycl: contiguous fast path + 32-bit index math for unary elementwise ops * sycl: use fastdiv for elementwise index math --- ggml/src/ggml-sycl/element_wise.cpp | 129 +++++++++++++++++++--------- 1 file changed, 87 insertions(+), 42 deletions(-) diff --git a/ggml/src/ggml-sycl/element_wise.cpp b/ggml/src/ggml-sycl/element_wise.cpp index b2406e11b5af..3cd055494ecf 100644 --- a/ggml/src/ggml-sycl/element_wise.cpp +++ b/ggml/src/ggml-sycl/element_wise.cpp @@ -306,29 +306,43 @@ static __dpct_inline__ T op_trunc(T x) { } } +template +static void unary_op_flat_kernel(const T * x, T * dst, const int k, const sycl::nd_item<1> & item_ct1, F func) { + SYCL_GLOBAL_ID_LOOP(k, item_ct1) { + dst[i] = func(x[i]); + } +} + template static void unary_op_generic_kernel( const T * x, T * dst, const int k, - const int64_t ne0, const int64_t ne1, const int64_t ne2, const int64_t ne3, + const sycl::uint3 ne0_fd, const sycl::uint3 ne1_fd, const sycl::uint3 ne2_fd, const size_t nb0, const size_t nb1, const size_t nb2, const size_t nb3, const size_t nbd0, const size_t nbd1, const size_t nbd2, const size_t nbd3, const sycl::nd_item<1> & item_ct1, F func) { - (void) ne3; + // 32-bit index math: k is int, so every logical index fits u32. 64-bit integer div/mod is + // emulated on Xe and dominates this kernel otherwise, and even the 32-bit divide is worth + // avoiding -- the divisors are launch-invariant, so the magic numbers are precomputed + // host-side and each division becomes a multiply-high plus a shift. + // Byte offsets are widened back to size_t only for the final address math. SYCL_GLOBAL_ID_LOOP(k, item_ct1) { - const int64_t i0 = i % ne0; - const int64_t i1 = (i / ne0) % ne1; - const int64_t i2 = (i / (ne0*ne1)) % ne2; - const int64_t i3 = i / (ne0*ne1*ne2); + sycl::uint2 dm = fast_div_modulo((uint32_t) i, ne0_fd); + const uint32_t i0 = dm.y(); + dm = fast_div_modulo(dm.x(), ne1_fd); + const uint32_t i1 = dm.y(); + dm = fast_div_modulo(dm.x(), ne2_fd); + const uint32_t i2 = dm.y(); + const uint32_t i3 = dm.x(); const char * src_base = (const char *) x; char * dst_base = (char *) dst; - const T * srcp = (const T *)(src_base + i0*nb0 + i1*nb1 + i2*nb2 + i3*nb3 ); - T * dstp = (T *)(dst_base + i0*nbd0 + i1*nbd1 + i2*nbd2 + i3*nbd3); + const T * srcp = (const T *)(src_base + (size_t) i0*nb0 + (size_t) i1*nb1 + (size_t) i2*nb2 + (size_t) i3*nb3 ); + T * dstp = (T *)(dst_base + (size_t) i0*nbd0 + (size_t) i1*nbd1 + (size_t) i2*nbd2 + (size_t) i3*nbd3); *dstp = func(*srcp); } @@ -407,46 +421,51 @@ static void clamp(const T * x, T * dst, const float min, const float max, const } template -static void gated_op_fused_geglu(const T * x, const T * g, T * dst, const uint64_t k, const uint64_t n, const uint64_t o0, const uint64_t o1, const sycl::nd_item<1> &item_ct1) { +static void gated_op_fused_geglu(const T * x, const T * g, T * dst, const uint64_t k, const sycl::uint3 n_fd, const uint64_t o0, const uint64_t o1, const sycl::nd_item<1> &item_ct1) { SYCL_GLOBAL_ID_LOOP(k, item_ct1) { - const int64_t j0 = (i / n) * o0 + (i % n); - const int64_t j1 = o0 == o1 ? j0 : (i / n) * o1 + (i % n); + const sycl::uint2 rc = fast_div_modulo((uint32_t) i, n_fd); + const int64_t j0 = rc.x() * o0 + rc.y(); + const int64_t j1 = o0 == o1 ? j0 : rc.x() * o1 + rc.y(); dst[i] = op_gelu(x[j0]) * g[j1]; } } template -static void gated_op_fused_reglu(const T * x, const T * g, T * dst, const uint64_t k, const uint64_t n, const uint64_t o0, const uint64_t o1, const sycl::nd_item<1> &item_ct1) { +static void gated_op_fused_reglu(const T * x, const T * g, T * dst, const uint64_t k, const sycl::uint3 n_fd, const uint64_t o0, const uint64_t o1, const sycl::nd_item<1> &item_ct1) { SYCL_GLOBAL_ID_LOOP(k, item_ct1) { - const int64_t j0 = (i / n) * o0 + (i % n); - const int64_t j1 = o0 == o1 ? j0 : (i / n) * o1 + (i % n); + const sycl::uint2 rc = fast_div_modulo((uint32_t) i, n_fd); + const int64_t j0 = rc.x() * o0 + rc.y(); + const int64_t j1 = o0 == o1 ? j0 : rc.x() * o1 + rc.y(); dst[i] = op_relu(x[j0]) * g[j1]; } } template -static void gated_op_fused_swiglu(const T * x, const T * g, T * dst, const uint64_t k, const uint64_t n, const uint64_t o0, const uint64_t o1, const sycl::nd_item<1> &item_ct1) { +static void gated_op_fused_swiglu(const T * x, const T * g, T * dst, const uint64_t k, const sycl::uint3 n_fd, const uint64_t o0, const uint64_t o1, const sycl::nd_item<1> &item_ct1) { SYCL_GLOBAL_ID_LOOP(k, item_ct1) { - const int64_t j0 = (i / n) * o0 + (i % n); - const int64_t j1 = o0 == o1 ? j0 : (i / n) * o1 + (i % n); + const sycl::uint2 rc = fast_div_modulo((uint32_t) i, n_fd); + const int64_t j0 = rc.x() * o0 + rc.y(); + const int64_t j1 = o0 == o1 ? j0 : rc.x() * o1 + rc.y(); dst[i] = op_silu(x[j0]) * g[j1]; } } template -static void gated_op_fused_geglu_erf(const T * x, const T * g, T * dst, const uint64_t k, const uint64_t n, const uint64_t o0, const uint64_t o1, const sycl::nd_item<1> &item_ct1) { +static void gated_op_fused_geglu_erf(const T * x, const T * g, T * dst, const uint64_t k, const sycl::uint3 n_fd, const uint64_t o0, const uint64_t o1, const sycl::nd_item<1> &item_ct1) { SYCL_GLOBAL_ID_LOOP(k, item_ct1) { - const int64_t j0 = (i / n) * o0 + (i % n); - const int64_t j1 = o0 == o1 ? j0 : (i / n) * o1 + (i % n); + const sycl::uint2 rc = fast_div_modulo((uint32_t) i, n_fd); + const int64_t j0 = rc.x() * o0 + rc.y(); + const int64_t j1 = o0 == o1 ? j0 : rc.x() * o1 + rc.y(); dst[i] = op_gelu_erf(x[j0]) * g[j1]; } } template -static void gated_op_fused_geglu_quick(const T * x, const T * g, T * dst, const uint64_t k, const uint64_t n, const uint64_t o0, const uint64_t o1, const sycl::nd_item<1> &item_ct1) { +static void gated_op_fused_geglu_quick(const T * x, const T * g, T * dst, const uint64_t k, const sycl::uint3 n_fd, const uint64_t o0, const uint64_t o1, const sycl::nd_item<1> &item_ct1) { SYCL_GLOBAL_ID_LOOP(k, item_ct1) { - const int64_t j0 = (i / n) * o0 + (i % n); - const int64_t j1 = o0 == o1 ? j0 : (i / n) * o1 + (i % n); + const sycl::uint2 rc = fast_div_modulo((uint32_t) i, n_fd); + const int64_t j0 = rc.x() * o0 + rc.y(); + const int64_t j1 = o0 == o1 ? j0 : rc.x() * o1 + rc.y(); dst[i] = op_gelu_quick(x[j0]) * g[j1]; } } @@ -529,6 +548,10 @@ static inline void dispatch_ggml_sycl_op_fused_glu(ggml_backend_sycl_context & c GGML_ASSERT(dst->ne[0] == nc); GGML_ASSERT(ggml_is_contiguous_1(dst->src[0])); GGML_ASSERT(ggml_is_contiguous(dst)); + // The fused GLU kernels index with 32-bit fastdiv, which is exact only for indices below + // 2^31. A dst that large is ~8 GB at f32, and the grid sizing already narrows to 32 bits, + // so assert the bound rather than carry a second code path for it. + GGML_ASSERT(ggml_nelements(dst) < ((int64_t) 1 << 31)); const int32_t swapped = ((const int32_t *) dst->op_params)[1]; void * src0_d = src0->data; void * src1_d = src1 ? src1->data : src0->data; @@ -597,7 +620,6 @@ static inline void ggml_sycl_op_unary( const int64_t ne0 = dst->ne[0]; const int64_t ne1 = dst->ne[1]; const int64_t ne2 = dst->ne[2]; - const int64_t ne3 = dst->ne[3]; const size_t nb0 = src0->nb[0]; const size_t nb1 = src0->nb[1]; @@ -609,24 +631,42 @@ static inline void ggml_sycl_op_unary( const size_t nbd2 = dst->nb[2]; const size_t nbd3 = dst->nb[3]; + // Hot unary ops (FFN/GDN silu, sigmoid, ...) run on contiguous tensors; + // skip the strided index math entirely for them. + const bool contiguous = ggml_is_contiguous(src0) && ggml_is_contiguous(dst); + ggml_sycl_detail::dispatch_ggml_sycl_op_unary(ctx, dst, [=](const auto* src, auto* dst_ptr, int k_elements, queue_ptr stream) { const int num_blocks = ceil_div(k_elements, 256); - stream->parallel_for( - sycl::nd_range<1>(sycl::range<1>(num_blocks) * sycl::range<1>(256), - sycl::range<1>(256)), - [=](sycl::nd_item<1> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]] { - unary_op_generic_kernel( - src, dst_ptr, k_elements, - ne0, ne1, ne2, ne3, - nb0, nb1, nb2, nb3, - nbd0, nbd1, nbd2, nbd3, - item_ct1, - func - ); - }); + if (contiguous) { + stream->parallel_for( + sycl::nd_range<1>(sycl::range<1>(num_blocks) * sycl::range<1>(256), + sycl::range<1>(256)), + [=](sycl::nd_item<1> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]] { + unary_op_flat_kernel(src, dst_ptr, k_elements, item_ct1, func); + }); + } else { + // Launch-invariant divisors: compute the magic numbers once on the host so the + // kernel never issues an integer divide. Only the strided path needs them. + const sycl::uint3 ne0_fd = init_fastdiv_values((uint32_t) ne0); + const sycl::uint3 ne1_fd = init_fastdiv_values((uint32_t) ne1); + const sycl::uint3 ne2_fd = init_fastdiv_values((uint32_t) ne2); + stream->parallel_for( + sycl::nd_range<1>(sycl::range<1>(num_blocks) * sycl::range<1>(256), + sycl::range<1>(256)), + [=](sycl::nd_item<1> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]] { + unary_op_generic_kernel( + src, dst_ptr, k_elements, + ne0_fd, ne1_fd, ne2_fd, + nb0, nb1, nb2, nb3, + nbd0, nbd1, nbd2, nbd3, + item_ct1, + func + ); + }); + } }); } @@ -930,10 +970,11 @@ static inline void ggml_sycl_op_geglu(ggml_backend_sycl_context & ctx, ggml_tens ggml_sycl_detail::dispatch_ggml_sycl_op_fused_glu(ctx, dst, [](const auto* x_ptr, const auto* g_ptr, auto* dst_ptr, uint64_t k, uint64_t n, uint64_t o0, uint64_t o1, queue_ptr main_stream) { const uint32_t num_blocks = ceil_div(k, SYCL_GELU_BLOCK_SIZE); + const sycl::uint3 n_fd = init_fastdiv_values((uint32_t) n); main_stream->parallel_for( sycl::nd_range<1>((num_blocks * sycl::range<1>(SYCL_GELU_BLOCK_SIZE)), sycl::range<1>(SYCL_GELU_BLOCK_SIZE)), [=](sycl::nd_item<1> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]] { - gated_op_fused_geglu(x_ptr, g_ptr, dst_ptr, k, n, o0, o1, item_ct1); + gated_op_fused_geglu(x_ptr, g_ptr, dst_ptr, k, n_fd, o0, o1, item_ct1); }); }); } @@ -942,10 +983,11 @@ static inline void ggml_sycl_op_reglu(ggml_backend_sycl_context & ctx, ggml_tens ggml_sycl_detail::dispatch_ggml_sycl_op_fused_glu(ctx, dst, [](const auto* x_ptr, const auto* g_ptr, auto* dst_ptr, uint64_t k, uint64_t n, uint64_t o0, uint64_t o1, queue_ptr main_stream) { const uint32_t num_blocks = ceil_div((uint32_t)k, SYCL_RELU_BLOCK_SIZE); // Using RELU block size for reglu + const sycl::uint3 n_fd = init_fastdiv_values((uint32_t) n); main_stream->parallel_for( sycl::nd_range<1>((num_blocks * sycl::range<1>(SYCL_RELU_BLOCK_SIZE)), sycl::range<1>(SYCL_RELU_BLOCK_SIZE)), [=](sycl::nd_item<1> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]] { - gated_op_fused_reglu(x_ptr, g_ptr, dst_ptr, k, n, o0, o1, item_ct1); + gated_op_fused_reglu(x_ptr, g_ptr, dst_ptr, k, n_fd, o0, o1, item_ct1); }); }); } @@ -954,10 +996,11 @@ static inline void ggml_sycl_op_swiglu(ggml_backend_sycl_context & ctx, ggml_ten ggml_sycl_detail::dispatch_ggml_sycl_op_fused_glu(ctx, dst, [](const auto* x_ptr, const auto* g_ptr, auto* dst_ptr, uint64_t k, uint64_t n, uint64_t o0, uint64_t o1, queue_ptr main_stream) { const uint32_t num_blocks = ceil_div((uint32_t)k, SYCL_SILU_BLOCK_SIZE); // Using SILU block size for swiglu + const sycl::uint3 n_fd = init_fastdiv_values((uint32_t) n); main_stream->parallel_for( sycl::nd_range<1>((num_blocks * sycl::range<1>(SYCL_SILU_BLOCK_SIZE)), sycl::range<1>(SYCL_SILU_BLOCK_SIZE)), [=](sycl::nd_item<1> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]] { - gated_op_fused_swiglu(x_ptr, g_ptr, dst_ptr, k, n, o0, o1, item_ct1); + gated_op_fused_swiglu(x_ptr, g_ptr, dst_ptr, k, n_fd, o0, o1, item_ct1); }); }); } @@ -1057,10 +1100,11 @@ static inline void ggml_sycl_op_geglu_erf(ggml_backend_sycl_context & ctx, ggml_ ggml_sycl_detail::dispatch_ggml_sycl_op_fused_glu(ctx, dst, [](const auto* x_ptr, const auto* g_ptr, auto* dst_ptr, uint64_t k, uint64_t n, uint64_t o0, uint64_t o1, queue_ptr main_stream) { const uint32_t num_blocks = ceil_div(k, SYCL_GELU_BLOCK_SIZE); + const sycl::uint3 n_fd = init_fastdiv_values((uint32_t) n); main_stream->parallel_for( sycl::nd_range<1>((num_blocks * sycl::range<1>(SYCL_GELU_BLOCK_SIZE)), sycl::range<1>(SYCL_GELU_BLOCK_SIZE)), [=](sycl::nd_item<1> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]] { - gated_op_fused_geglu_erf(x_ptr, g_ptr, dst_ptr, k, n, o0, o1, item_ct1); + gated_op_fused_geglu_erf(x_ptr, g_ptr, dst_ptr, k, n_fd, o0, o1, item_ct1); }); }); } @@ -1069,10 +1113,11 @@ static inline void ggml_sycl_op_geglu_quick(ggml_backend_sycl_context & ctx, ggm ggml_sycl_detail::dispatch_ggml_sycl_op_fused_glu(ctx, dst, [](const auto* x_ptr, const auto* g_ptr, auto* dst_ptr, uint64_t k, uint64_t n, uint64_t o0, uint64_t o1, queue_ptr main_stream) { const uint32_t num_blocks = ceil_div(k, SYCL_GELU_BLOCK_SIZE); + const sycl::uint3 n_fd = init_fastdiv_values((uint32_t) n); main_stream->parallel_for( sycl::nd_range<1>((num_blocks * sycl::range<1>(SYCL_GELU_BLOCK_SIZE)), sycl::range<1>(SYCL_GELU_BLOCK_SIZE)), [=](sycl::nd_item<1> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]] { - gated_op_fused_geglu_quick(x_ptr, g_ptr, dst_ptr, k, n, o0, o1, item_ct1); + gated_op_fused_geglu_quick(x_ptr, g_ptr, dst_ptr, k, n_fd, o0, o1, item_ct1); }); }); } From caa596ab3f0f8768ee326d6e3d5d39782194676c Mon Sep 17 00:00:00 2001 From: Kakaru <97896816+KakaruHayate@users.noreply.github.com> Date: Wed, 29 Jul 2026 20:27:35 +0800 Subject: [PATCH 053/190] ggml-cuda : disable MMQ on devices with less than 48 KiB shared memory (#26141) ggml_cuda_should_use_mmq() selects MMQ purely from the quantization type. The current MMQ configurations are designed and maintained against a minimum of 48 KiB per-block shared memory, the limit provided by NVIDIA Pascal GPUs and later. On devices that report less, no supported MMQ tile fits and mul_mat_q_switch_J() aborts when every tile size exceeds the device's per-block shared memory budget. Disable MMQ when smpbo < 48 KiB so the caller falls back to the BLAS path instead of hitting GGML_ABORT. Some current MUSA QY1 devices report only 28 KiB and are covered by this guard. Reproduced on a Moore Threads MTT S70 (arch mp_21, 28 KiB shared memory per block) with an RWKV-7 0.1B Q8_0 model: $ llama-bench -m rwkv7-g1d-0.1b-Q8_0.gguf -p 128 -n 0 J_best=0 ggml/src/ggml-cuda/template-instances/../mmq.cuh:1521: fatal error (core dumped) Only prefill (batch > 1) is affected; token generation is fine. After the fix the same device falls back to the BLAS path: Q8_0 pp128 1470.7 t/s, tg8 55.3 t/s (was: abort) FP16 unchanged Q4_K_M unchanged This matches a -DGGML_CUDA_FORCE_CUBLAS=ON build (pp128 1464.2 t/s), which confirms the fallback path is the one being taken. This is not MUSA-specific: any device with less than 48 KiB per-block shared memory is affected. Co-authored-by: KakaruHayate --- ggml/src/ggml-cuda/mmq.cu | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/ggml/src/ggml-cuda/mmq.cu b/ggml/src/ggml-cuda/mmq.cu index 8a0f4d3b5cbf..8cf1a3eb94f6 100644 --- a/ggml/src/ggml-cuda/mmq.cu +++ b/ggml/src/ggml-cuda/mmq.cu @@ -296,6 +296,15 @@ bool ggml_cuda_should_use_mmq(enum ggml_type type, int cc, int64_t ne11, int64_t return false; } + // MMQ tiles require at least 48 KiB per-block shared memory; fall back to BLAS otherwise. + { + const int id = ggml_cuda_get_device(); + const size_t smpbo = ggml_cuda_info().devices[id].smpbo; + if (smpbo < 48 * 1024) { + return false; + } + } + if (turing_mma_available(cc)) { return true; } From afeebe103bd99cda8f5dfaefcabadf890db7fda7 Mon Sep 17 00:00:00 2001 From: Xuan-Son Nguyen Date: Wed, 29 Jul 2026 18:02:30 +0200 Subject: [PATCH 054/190] llama: move suppress_tokens handling to common/sampling (#26276) * llama: move suppress_tokens handling to common/sampling * address security issues * rm has_logit_bias --- common/common.h | 4 ---- common/sampling.cpp | 15 +++++++++++++-- include/llama.h | 3 +++ src/llama-vocab.cpp | 17 ++++++++++++++++- src/models/gemma4.cpp | 37 ------------------------------------- 5 files changed, 32 insertions(+), 44 deletions(-) diff --git a/common/common.h b/common/common.h index e7c55ae925fc..14889193d9b5 100644 --- a/common/common.h +++ b/common/common.h @@ -294,10 +294,6 @@ struct common_params_sampling { bool backend_sampling = false; - bool has_logit_bias() const { - return !logit_bias.empty(); - } - // print the parameters into a string std::string print() const; }; diff --git a/common/sampling.cpp b/common/sampling.cpp index 7b241e34f77f..256ac161e20f 100644 --- a/common/sampling.cpp +++ b/common/sampling.cpp @@ -310,8 +310,19 @@ struct common_sampler * common_sampler_init(const struct llama_model * model, st } } - if (params.has_logit_bias()) { - samplers.push_back(llama_sampler_init_logit_bias(llama_vocab_n_tokens(vocab), params.logit_bias.size(), params.logit_bias.data())); + // logit bias: user biases + model suppress tokens (-INFINITY) + { + std::vector merged = params.logit_bias; + + int32_t n_suppress = 0; + const llama_token * suppress = llama_vocab_get_suppress_tokens(vocab, &n_suppress); + for (int32_t i = 0; i < n_suppress; ++i) { + merged.push_back({ suppress[i], -INFINITY }); + } + + if (!merged.empty()) { + samplers.push_back(llama_sampler_init_logit_bias(llama_vocab_n_tokens(vocab), merged.size(), merged.data())); + } } if (params.mirostat == 0) { diff --git a/include/llama.h b/include/llama.h index 3c6d22be8999..a3ceecf5ef47 100644 --- a/include/llama.h +++ b/include/llama.h @@ -1102,6 +1102,9 @@ extern "C" { LLAMA_API bool llama_vocab_get_add_eos(const struct llama_vocab * vocab); LLAMA_API bool llama_vocab_get_add_sep(const struct llama_vocab * vocab); + // model-specific suppress tokens (gguf key: tokenizer.ggml.suppress_tokens) + LLAMA_API const llama_token * llama_vocab_get_suppress_tokens(const struct llama_vocab * vocab, int32_t * n_suppress_tokens); + LLAMA_API llama_token llama_vocab_fim_pre(const struct llama_vocab * vocab); LLAMA_API llama_token llama_vocab_fim_suf(const struct llama_vocab * vocab); LLAMA_API llama_token llama_vocab_fim_mid(const struct llama_vocab * vocab); diff --git a/src/llama-vocab.cpp b/src/llama-vocab.cpp index 9164a4dd888d..443cd4640858 100644 --- a/src/llama-vocab.cpp +++ b/src/llama-vocab.cpp @@ -2578,7 +2578,14 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) { if (suppress_idx != -1) { const int n = gguf_get_arr_n(ctx, suppress_idx); const int32_t * data = (const int32_t *) gguf_get_arr_data(ctx, suppress_idx); - suppress_tokens.assign(data, data + n); + // drop out-of-range ids + suppress_tokens.reserve(n); + for (int i = 0; i < n; ++i) { + const int32_t id = data[i]; + if (id >= 0 && id < (int) id_to_token.size()) { + suppress_tokens.push_back(id); + } + } } } @@ -4205,6 +4212,14 @@ bool llama_vocab_get_add_sep(const struct llama_vocab * vocab) { return vocab->get_add_sep(); } +const llama_token * llama_vocab_get_suppress_tokens(const struct llama_vocab * vocab, int32_t * n_suppress_tokens) { + const std::vector & tokens = vocab->get_suppress_tokens(); + if (n_suppress_tokens) { + *n_suppress_tokens = (int32_t) tokens.size(); + } + return tokens.data(); +} + llama_token llama_vocab_fim_pre(const struct llama_vocab * vocab) { return vocab->token_fim_pre(); } diff --git a/src/models/gemma4.cpp b/src/models/gemma4.cpp index 6a96979cebde..e44f423bdbc5 100644 --- a/src/models/gemma4.cpp +++ b/src/models/gemma4.cpp @@ -142,33 +142,6 @@ static ggml_tensor * ggml_view_2d_slice(ggml_context * ctx0, ggml_tensor * x, in idx * x->ne[0] * x->ne[1] * ggml_element_size(x)); } -// TODO @ngxson : maybe improve this in the future -class llm_graph_input_logits_bias : public llm_graph_input_i { -public: - llm_graph_input_logits_bias(const llama_vocab & vocab) { - arr.resize(vocab.n_tokens(), 0.0f); - for (llama_token id : vocab.get_suppress_tokens()) { - if (0 <= id && id < (int32_t)vocab.n_tokens()) { - arr[id] = -INFINITY; - } - } - } - virtual ~llm_graph_input_logits_bias() = default; - - void set_input(const llama_ubatch * /*ubatch*/) override { - const int64_t n_vocab = arr.size(); - ggml_backend_tensor_set(logits_bias, arr.data(), 0, n_vocab*ggml_element_size(logits_bias)); - } - - bool can_reuse(const llm_graph_params & /*params*/) override { - return true; - } - - ggml_tensor * logits_bias = nullptr; // F32 [n_vocab] - - std::vector arr; -}; - llama_model_gemma4::graph::graph(const llama_model & model, const llm_graph_params & params) : llm_graph_context(params), model(model), @@ -429,16 +402,6 @@ llama_model_gemma4::graph::graph(const llama_model & model, const llm_graph_para cur = ggml_scale(ctx0, cur, hparams.f_final_logit_softcapping); } - // apply logits bias if needed (e.g. for gemma4_unified patch) - // this is to mirror the suppress_tokens patch on transformers, to avoid model from outputing and tokens (which is a known issue related to the checkpoint) - // TODO: maybe handle this inside the sampling system in the future - if (!model.vocab.get_suppress_tokens().empty()) { - auto inp_bias = std::make_unique(model.vocab); - inp_bias->logits_bias = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, inp_bias->arr.size()); - cur = ggml_add(ctx0, cur, inp_bias->logits_bias); - res->add_input(std::move(inp_bias)); - } - cb(cur, "result_output", -1); res->t_logits = cur; From 3018a11e79e489b657dbb77c95694889ccff92df Mon Sep 17 00:00:00 2001 From: Aleksander Grygier Date: Wed, 29 Jul 2026 19:25:13 +0200 Subject: [PATCH 055/190] fix: increase greeting spacing on md screens (#26287) --- .../components/app/chat/ChatScreen/ChatScreenGreeting.svelte | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tools/ui/src/lib/components/app/chat/ChatScreen/ChatScreenGreeting.svelte b/tools/ui/src/lib/components/app/chat/ChatScreen/ChatScreenGreeting.svelte index 018949aff87a..5b44bcf858b1 100644 --- a/tools/ui/src/lib/components/app/chat/ChatScreen/ChatScreenGreeting.svelte +++ b/tools/ui/src/lib/components/app/chat/ChatScreen/ChatScreenGreeting.svelte @@ -11,7 +11,7 @@