From a52603f2c31b3263a4502cba227a4e452282dc70 Mon Sep 17 00:00:00 2001 From: bri-prism <288398250+bri-prism@users.noreply.github.com> Date: Mon, 13 Jul 2026 19:05:30 +0100 Subject: [PATCH 1/2] metal: fix M5 device creation + add Q2_0 multi-column mul_mv kernels Two independent Metal changes: - The AGX_RELAX_CDM_CTXSTORE_TIMEOUT override (added for the long-context command-buffer timeout on M1/M2, ggml-org#20141) prevents MTLCreateSystemDefaultDevice() from returning a device on M5 / current macOS. Keep it on by default and disable it only when sysctl reports an M5 chip; GGML_METAL_RELAX_CDM_CTXSTORE_TIMEOUT=0/1 forces either way. - Add Q1_0-style multi-column mul_mv variants for Q2_0 (nr1 2/3/4): read the streamed weights once for nr1 output columns via a 2-bit weight expansion and an FMA inner loop, instead of re-reading them per column on the mul_mv_ext path. Opt-in via GGML_METAL_Q2_0_NR1 (default routing unchanged). Measured [4096,14336] on M5 Pro: nr1_2 93.2 us at ne11=2 vs 122 for the ext route. 41/41 test-backend-ops MUL_MAT q2_0 on both routings. --- ggml/src/ggml-metal/ggml-metal-device.cpp | 16 +++ ggml/src/ggml-metal/ggml-metal-ops.cpp | 6 +- ggml/src/ggml-metal/ggml-metal.cpp | 31 ++++- ggml/src/ggml-metal/ggml-metal.metal | 162 ++++++++++++++++++---- 4 files changed, 188 insertions(+), 27 deletions(-) diff --git a/ggml/src/ggml-metal/ggml-metal-device.cpp b/ggml/src/ggml-metal/ggml-metal-device.cpp index 24122e7b4136..86a60c060565 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.cpp +++ b/ggml/src/ggml-metal/ggml-metal-device.cpp @@ -862,6 +862,22 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv(ggml_meta { nsg = N_SG_Q2_0; nr0 = N_R0_Q2_0; + + // multi-column variants, same scheme as Q1_0 above: read the + // streamed q2_0 weights once per nr1 src1 columns. + // EXPERIMENTAL, opt-in via GGML_METAL_Q2_0_NR1 (0/absent keeps + // the default routing, i.e. the mul_mv_ext path for ne11 2..8). + // Measured (M5 Pro, [4096,14336]): nr1_2 = 93.2 us at ne11=2 + // vs 122 for the ext route (+31%); ne11=4 via 2 passes = 171 + // vs 183. But nr1_3 = 195 vs 152 ext at ne11=3 (occupancy + // cliff at tpb=16) -- routing is NOT settled yet, hence opt-in. + static const int nr1_max = getenv("GGML_METAL_Q2_0_NR1") ? atoi(getenv("GGML_METAL_Q2_0_NR1")) : 0; + + const int nr1_force = nr1_max >= 2 && nr1_max <= 4 ? nr1_max : 0; + if (nr1_force > 1 && ne11 >= 2) { + nr1 = std::min(nr1_force, 4); + suffix = nr1 == 2 ? "_nr1_2" : nr1 == 3 ? "_nr1_3" : "_nr1_4"; + } } break; case GGML_TYPE_Q4_0: { diff --git a/ggml/src/ggml-metal/ggml-metal-ops.cpp b/ggml/src/ggml-metal/ggml-metal-ops.cpp index 7a87c8dc04c9..879e826373a3 100644 --- a/ggml/src/ggml-metal/ggml-metal-ops.cpp +++ b/ggml/src/ggml-metal/ggml-metal-ops.cpp @@ -2068,6 +2068,10 @@ int ggml_metal_op_mul_mat(ggml_metal_op_t ctx, int idx) { static const int ne11_ext_max = getenv("GGML_METAL_EXT_MAX") ? atoi(getenv("GGML_METAL_EXT_MAX")) : 8; static const bool q1_0_ext_enable = getenv("GGML_METAL_Q1_0_EXT_ENABLE") != NULL; static const int q1_0_mv_max = getenv("GGML_METAL_Q1_0_MV_MAX") ? atoi(getenv("GGML_METAL_Q1_0_MV_MAX")) : 16; + // GGML_METAL_Q2_0_NR1 >= 2 routes Q2_0 ne11 2..8 off the ext path and onto the + // experimental multi-column mul_mv variants (see ggml-metal-device.cpp); the + // default keeps Q2_0 on the ext path + static const int q2_0_nr1 = getenv("GGML_METAL_Q2_0_NR1") ? atoi(getenv("GGML_METAL_Q2_0_NR1")) : 0; // narrow-N tensor-path mul_mm for q1_0 mid-size batches (spec-decode verify): // GGML_METAL_Q1_0_NB_MIN/_NB_MAX - ne11 range routed to the nb kernels (min 0 disables) @@ -2101,7 +2105,7 @@ int ggml_metal_op_mul_mat(ggml_metal_op_t ctx, int idx) { op->src[0]->type == GGML_TYPE_F16 || op->src[0]->type == GGML_TYPE_BF16 || (op->src[0]->type == GGML_TYPE_Q1_0 && q1_0_ext_enable) || - op->src[0]->type == GGML_TYPE_Q2_0 || + (op->src[0]->type == GGML_TYPE_Q2_0 && q2_0_nr1 < 2) || op->src[0]->type == GGML_TYPE_Q4_0 || op->src[0]->type == GGML_TYPE_Q4_1 || op->src[0]->type == GGML_TYPE_Q5_0 || diff --git a/ggml/src/ggml-metal/ggml-metal.cpp b/ggml/src/ggml-metal/ggml-metal.cpp index a1003b3acff8..105714d8b433 100644 --- a/ggml/src/ggml-metal/ggml-metal.cpp +++ b/ggml/src/ggml-metal/ggml-metal.cpp @@ -7,9 +7,16 @@ #include "ggml-metal-context.h" #include "ggml-metal-ops.h" +#include +#include #include #include +#include +#if TARGET_OS_OSX +#include +#endif + #define GGML_METAL_NAME "MTL" #define GGML_METAL_MAX_DEVICES 16 @@ -923,7 +930,29 @@ ggml_backend_reg_t ggml_backend_metal_reg(void) { if (!initialized) { // workaround macOS limitation (kIOGPUCommandBufferCallbackErrorImpactingInteractivity) until proper fix becomes possible // ref: https://github.com/ggml-org/llama.cpp/issues/20141#issuecomment-4272947703 - setenv("AGX_RELAX_CDM_CTXSTORE_TIMEOUT", "1", true); + // + // The override fixes long-context command-buffer timeouts on + // M1/M2, but on M5/current macOS it prevents + // MTLCreateSystemDefaultDevice() from returning a device at all. + // Keep it on by default and disable it only on M5 (sysctl needs + // no Metal device); GGML_METAL_RELAX_CDM_CTXSTORE_TIMEOUT=0/1 + // forces either way. + bool relax_cdm_ctxstore = true; +#if TARGET_OS_OSX + { + char brand[128] = { 0 }; + size_t brand_len = sizeof(brand) - 1; + if (sysctlbyname("machdep.cpu.brand_string", brand, &brand_len, NULL, 0) == 0 && strstr(brand, " M5") != NULL) { + relax_cdm_ctxstore = false; + } + } +#endif + if (const char * env = getenv("GGML_METAL_RELAX_CDM_CTXSTORE_TIMEOUT")) { + relax_cdm_ctxstore = atoi(env) != 0; + } + if (relax_cdm_ctxstore) { + setenv("AGX_RELAX_CDM_CTXSTORE_TIMEOUT", "1", true); + } static ggml_backend_metal_reg_ptr reg_ctx(ggml_backend_metal_reg_init()); diff --git a/ggml/src/ggml-metal/ggml-metal.metal b/ggml/src/ggml-metal/ggml-metal.metal index 1154bdd98c8a..f5836fca94c1 100644 --- a/ggml/src/ggml-metal/ggml-metal.metal +++ b/ggml/src/ggml-metal/ggml-metal.metal @@ -3783,7 +3783,34 @@ kernel void kernel_mul_mv_q1_0_f32_nr1_4( kernel_mul_mv_q1_0_f32_impl(args, src0, src1, dst, nullptr, tgpig, tiisg, sgitg); } -template + +// dot of an SW-element yl slice against the matching SW 2-bit codes of a +// q2_0 block (bytes pre-loaded so they can be reused across src1 columns). +// Same lo/hi bit decomposition as block_q_n_dot_y(block_q2_0) above; each +// accumulator adds in ascending element order, so the SW=16 form is +// bit-identical to that helper. +template +static inline float q2_0_dot_y(thread const uint8_t * b, const float d, const float sumy, thread const float * yl) { + float acc_lo = 0.0f; + float acc_hi = 0.0f; + + FOR_UNROLL (short i = 0; i < SW; i++) { + acc_lo += select(0.0f, yl[i], bool(b[i/4] & (1u << (2*(i%4) + 0)))); + acc_hi += select(0.0f, yl[i], bool(b[i/4] & (1u << (2*(i%4) + 1)))); + } + + return d * (acc_lo + 2.0f*acc_hi - sumy); +} + +// nr0: src0 rows per simdgroup, nr1: src1 columns per threadgroup-y slot, +// tpb: threads cooperating on one q2_0 block (slice width SW = QK2_0/tpb). +// Same structure as kernel_mul_mv_q1_0_f32_impl above: nr1 > 1 reads the +// streamed, bandwidth-dominant q2_0 weights ONCE for nr1 output columns. +// This is the spec-decode verify path: the generic mul_mv_ext route runs at +// roughly half this kernel's effective weight bandwidth per pass (measured +// n=3 at 2.08x the n=1 cost on M5 Pro), and mul_mm only pays off for +// ne11 >~ 32. Register budget mirrors q1_0: keep nr1*SW <= 32. +template void kernel_mul_mv_q2_0_f32_impl( args_t args, device const char * src0, @@ -3794,6 +3821,7 @@ void kernel_mul_mv_q2_0_f32_impl( ushort tiisg, ushort sgitg) { const short NSG = FC_mul_mv_nsg; + const short SW = QK2_0/tpb; // y-slice elements per thread const int nb = args.ne00/QK2_0; @@ -3802,50 +3830,98 @@ void kernel_mul_mv_q2_0_f32_impl( const int im = tgpig.z; const int first_row = (r0 * NSG + sgitg) * nr0; + const int c0 = r1 * nr1; const uint i12 = im%args.ne12; const uint i13 = im/args.ne12; - const uint64_t offset1 = r1*args.nb11 + (i12)*args.nb12 + (i13)*args.nb13; - - device const float * y = (device const float *) (src1 + offset1); - device const block_q2_0 * ax[nr0]; for (int row = 0; row < nr0; ++row) { const uint64_t offset0 = (first_row + row)*args.nb01 + (i12/args.r2)*args.nb02 + (i13/args.r3)*args.nb03; ax[row] = (device const block_q2_0 *) ((device char *) src0 + offset0); } - float yl[16]; - float sumf[nr0] = {0.f}; - - const short ix = (tiisg/8); - const short il = (tiisg%8)*16; + float yl[nr1][SW]; + float sumy[nr1]; + float sumf[nr0][nr1]; + FOR_UNROLL (short row = 0; row < nr0; row++) { + FOR_UNROLL (short c = 0; c < nr1; c++) { + sumf[row][c] = 0.f; + } + } - device const float * yb = y + ix*QK2_0 + il; + const short ix = (tiisg/tpb); // block in flight + const short il = (tiisg%tpb)*SW; // element offset within the block - for (int ib = ix; ib < nb; ib += N_SIMDWIDTH/8) { - float sumy = 0.f; + device const float * yb[nr1]; + FOR_UNROLL (short c = 0; c < nr1; c++) { + // tail columns are clamped (results computed but not stored) + const int ic = MIN(c0 + c, args.ne11 - 1); + const uint64_t offset1 = (uint64_t)ic*args.nb11 + (i12)*args.nb12 + (i13)*args.nb13; + yb[c] = (device const float *) (src1 + offset1) + ix*QK2_0 + il; + } - FOR_UNROLL (short i = 0; i < 16; i++) { - yl[i] = yb[i]; - sumy += yb[i]; + for (int ib = ix; ib < nb; ib += N_SIMDWIDTH/tpb) { + FOR_UNROLL (short c = 0; c < nr1; c++) { + sumy[c] = 0.f; + FOR_UNROLL (short i = 0; i < SW; i++) { + yl[c][i] = yb[c][i]; + sumy[c] += yb[c][i]; + } } FOR_UNROLL (short row = 0; row < nr0; row++) { - sumf[row] += block_q_n_dot_y(ax[row] + ib, sumy, yl, il); + device const block_q2_0 * qb = ax[row] + ib; + device const uint8_t * qs = qb->qs + il/4; + const float d = qb->d; + + uint8_t b[SW/4]; + FOR_UNROLL (short i = 0; i < SW/4; i++) { + b[i] = qs[i]; + } + + if (nr1 > 1) { + // multi-column: the select-form dot is ALU-bound (2 conditional + // adds per element PER COLUMN), which is what makes the ext + // route scale ~linearly in n. Expand the 2-bit codes ONCE into + // float weights {0..3} and leave a single FMA per + // column-element: sum((q-1)*d*y) = d*(sum(q*y) - sumy). + float w[SW]; + FOR_UNROLL (short i = 0; i < SW; i++) { + w[i] = (float) ((b[i/4] >> (2*(i%4))) & 3); + } + + FOR_UNROLL (short c = 0; c < nr1; c++) { + float acc = 0.0f; + FOR_UNROLL (short i = 0; i < SW; i++) { + acc = fma(w[i], yl[c][i], acc); + } + sumf[row][c] += d*(acc - sumy[c]); + } + } else { + // single column: keep the exact select-form accumulation order + // of the original kernel (bit-identical AR/decode path) + sumf[row][0] += q2_0_dot_y(b, d, sumy[0], yl[0]); + } } - yb += QK2_0 * (N_SIMDWIDTH/8); + FOR_UNROLL (short c = 0; c < nr1; c++) { + yb[c] += QK2_0 * (N_SIMDWIDTH/tpb); + } } - device float * dst_f32 = (device float *) dst + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0; + device float * dst_f32 = (device float *) dst + (uint64_t)im*args.ne0*args.ne1; - for (int row = 0; row < nr0; ++row) { - const float tot = simd_sum(sumf[row]); + for (short c = 0; c < nr1; c++) { + if (c0 + c >= args.ne11) { + break; + } + for (int row = 0; row < nr0; ++row) { + const float tot = simd_sum(sumf[row][c]); - if (tiisg == 0 && first_row + row < args.ne01) { - dst_f32[first_row + row] = tot; + if (tiisg == 0 && first_row + row < args.ne01) { + dst_f32[(uint64_t)(c0 + c)*args.ne0 + first_row + row] = tot; + } } } } @@ -3859,7 +3935,43 @@ kernel void kernel_mul_mv_q2_0_f32( uint3 tgpig[[threadgroup_position_in_grid]], ushort tiisg[[thread_index_in_simdgroup]], ushort sgitg[[simdgroup_index_in_threadgroup]]) { - kernel_mul_mv_q2_0_f32_impl(args, src0, src1, dst, nullptr, tgpig, tiisg, sgitg); + kernel_mul_mv_q2_0_f32_impl(args, src0, src1, dst, nullptr, tgpig, tiisg, sgitg); +} + +[[host_name("kernel_mul_mv_q2_0_f32_nr1_2")]] +kernel void kernel_mul_mv_q2_0_f32_nr1_2( + constant ggml_metal_kargs_mul_mv & args, + device const char * src0, + device const char * src1, + device char * dst, + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiisg[[thread_index_in_simdgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]]) { + kernel_mul_mv_q2_0_f32_impl(args, src0, src1, dst, nullptr, tgpig, tiisg, sgitg); +} + +[[host_name("kernel_mul_mv_q2_0_f32_nr1_3")]] +kernel void kernel_mul_mv_q2_0_f32_nr1_3( + constant ggml_metal_kargs_mul_mv & args, + device const char * src0, + device const char * src1, + device char * dst, + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiisg[[thread_index_in_simdgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]]) { + kernel_mul_mv_q2_0_f32_impl(args, src0, src1, dst, nullptr, tgpig, tiisg, sgitg); +} + +[[host_name("kernel_mul_mv_q2_0_f32_nr1_4")]] +kernel void kernel_mul_mv_q2_0_f32_nr1_4( + constant ggml_metal_kargs_mul_mv & args, + device const char * src0, + device const char * src1, + device char * dst, + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiisg[[thread_index_in_simdgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]]) { + kernel_mul_mv_q2_0_f32_impl(args, src0, src1, dst, nullptr, tgpig, tiisg, sgitg); } kernel void kernel_mul_mv_q4_0_f32( @@ -10839,7 +10951,7 @@ template [[host_name("kernel_mul_mv_id_bf16_f32_4")]] kernel kernel_mul_mv_id_4 template [[host_name("kernel_mul_mv_id_q8_0_f32")]] kernel kernel_mul_mv_id_t kernel_mul_mv_id>>; template [[host_name("kernel_mul_mv_id_q1_0_f32")]] kernel kernel_mul_mv_id_t kernel_mul_mv_id>>; -template [[host_name("kernel_mul_mv_id_q2_0_f32")]] kernel kernel_mul_mv_id_t kernel_mul_mv_id>>; +template [[host_name("kernel_mul_mv_id_q2_0_f32")]] kernel kernel_mul_mv_id_t kernel_mul_mv_id>>; template [[host_name("kernel_mul_mv_id_q4_0_f32")]] kernel kernel_mul_mv_id_t kernel_mul_mv_id>>; template [[host_name("kernel_mul_mv_id_q4_1_f32")]] kernel kernel_mul_mv_id_t kernel_mul_mv_id>>; template [[host_name("kernel_mul_mv_id_q5_0_f32")]] kernel kernel_mul_mv_id_t kernel_mul_mv_id>>; From a1648218fbb0fe29e6c7da4b972b70a7aad613c0 Mon Sep 17 00:00:00 2001 From: bri-prism <288398250+bri-prism@users.noreply.github.com> Date: Mon, 13 Jul 2026 19:16:54 +0100 Subject: [PATCH 2/2] speculative: Metal DSpark Markov resample + quantized markov heads Adds a Metal device path for the block Markov resample, alongside the existing CUDA path. It builds one dependency-chain graph for the whole draft block (each step's GPU argmax feeds the next step's get_rows) and submits it once, reading the drafter's still-device-resident logits, so the sequential Markov dependency stays exact with a single sync. Falls back to the host path when the head type or backend is unsupported, or when DSPARK_MARKOV_CPU=1. Also teaches llama_model_dspark_get_markov to dequantize quantized markov head tensors (Q4_0/Q5_0/Q8_0) for the host/CUDA path, instead of rejecting everything but f32/f16/bf16 -- drafters that ship quantized heads previously had the correction silently disabled (has_markov=0). Validated on Metal: accept counts byte-identical to the CPU-forced path; the Metal resample runs ~1.8x the single-thread CPU Markov path. --- common/speculative.cpp | 24 +++++- src/llama-context.cpp | 161 +++++++++++++++++++++++++++++++++++++++++ src/llama-context.h | 7 ++ src/llama-ext.h | 11 +++ src/llama-model.cpp | 20 ++++- 5 files changed, 219 insertions(+), 4 deletions(-) diff --git a/common/speculative.cpp b/common/speculative.cpp index 230c65c542cc..fc4b8f8ded41 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -1297,9 +1297,31 @@ struct common_speculative_impl_draft_dspark : public common_speculative_impl { } #endif + // Metal path, same contract as the CUDA block above: one + // dependency-chain graph for the whole block (each step's GPU + // argmax feeds the next step's get_rows), reading the drafter's + // still-device-resident logits. Returns false when the head or + // backend is unsupported (or DSPARK_MARKOV_CPU=1) -> host path. + bool did_metal = false; + if (!did_cuda && has_markov) { + result.resize((size_t) block_size); + if (llama_dspark_markov_resample(ctx_dft, block_size, dp.id_last, result.data())) { + static bool warned_metal_mask = false; + for (int32_t k = 1; k < block_size && !warned_metal_mask; ++k) { + if (result[(size_t) (k - 1)] == mask_token_id) { + LOG_WRN("%s: metal markov resample sampled mask_token_id at a chained position\n", __func__); + warned_metal_mask = true; + } + } + did_metal = true; + } else { + result.clear(); + } + } + llama_token prev_token = dp.id_last; - if (!did_cuda) + if (!did_cuda && !did_metal) for (int32_t k = 0; k < block_size; ++k) { // prev_token is the token SAMPLED at step k-1 (assigned from best_id // at the end of this loop), never a draft input id -- that is the diff --git a/src/llama-context.cpp b/src/llama-context.cpp index d19888efee72..0c90ad3021fc 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -2555,6 +2555,159 @@ ggml_status llama_context::graph_compute( return status; } +bool llama_context::dspark_markov_resample( + uint32_t n_rows, + llama_token prev_token, + llama_token * result) { + if (n_rows == 0 || result == nullptr || getenv("DSPARK_MARKOV_CPU") != nullptr) { + return false; + } + + const ggml_tensor * head_a = model.dspark_markov_head_a; + const ggml_tensor * head_b = model.dspark_markov_head_b; + if (head_a == nullptr || head_b == nullptr || gf_res_prev == nullptr) { + return false; + } + + const ggml_tensor * t_logits = gf_res_prev->get_logits(); + const int64_t n_vocab = model.vocab.n_tokens(); + if (t_logits == nullptr || t_logits->type != GGML_TYPE_F32 || t_logits->data == nullptr || + t_logits->ne[0] != n_vocab || t_logits->ne[1] < (int64_t) n_rows || + head_a->ne[0] != head_b->ne[0] || head_a->ne[1] != head_b->ne[1] || + head_a->ne[1] != n_vocab) { + return false; + } + + const auto supported_head_type = [](ggml_type type) { + return type == GGML_TYPE_F32 || type == GGML_TYPE_F16 || type == GGML_TYPE_BF16 || + type == GGML_TYPE_Q4_0 || type == GGML_TYPE_Q5_0 || type == GGML_TYPE_Q8_0; + }; + if (!supported_head_type(head_a->type) || !supported_head_type(head_b->type)) { + return false; + } + + const ggml_backend_dev_t dev = model.dev_output(); + if (dev == nullptr || ggml_backend_dev_type(dev) == GGML_BACKEND_DEVICE_TYPE_CPU) { + return false; + } + + if (!dspark_markov_sched) { + if (backend_ptrs.empty() || backend_buft.size() != backend_ptrs.size()) { + return false; + } + + dspark_markov_sched.reset(ggml_backend_sched_new( + backend_ptrs.data(), backend_buft.data(), (int) backend_ptrs.size(), + /* graph_size = */ 64, /* parallel = */ false, cparams.op_offload)); + if (!dspark_markov_sched) { + return false; + } + } + + // The decode graph is asynchronous. Synchronize it once before the + // dedicated scheduler reads its logits output tensor. + synchronize(); + + // Build one graph covering rows [k0, k0 + n_chain). Each step consumes the + // previous step's GPU argmax tensor as the row id for head_a, so the + // sequential Markov dependency remains exact while Metal executes the + // whole chain in one scheduler submission. + const auto resample_chain = [&](uint32_t k0, uint32_t n_chain, llama_token tok0) -> bool { + ggml_init_params params = { + /*.mem_size =*/ 128*ggml_tensor_overhead() + ggml_graph_overhead_custom(64, false), + /*.mem_buffer =*/ nullptr, + /*.no_alloc =*/ true, + }; + ggml_context_ptr ctx { ggml_init(params) }; + if (!ctx) { + return false; + } + + ggml_tensor * ids = ggml_new_tensor_1d(ctx.get(), GGML_TYPE_I32, 1); + ggml_set_input(ids); + + // Do not use ggml_view_1d on t_logits: its parent edge would recursively + // pull the completed decode graph into this tiny graph. Each base tensor + // is a detached, read-only alias of one already-computed logits row. + std::vector sampled_rows; + sampled_rows.reserve(n_chain); + + ggml_tensor * prev_ids = ids; + for (uint32_t k = k0; k < k0 + n_chain; ++k) { + ggml_tensor * base = ggml_new_tensor_1d(ctx.get(), GGML_TYPE_F32, n_vocab); + base->buffer = t_logits->buffer; + base->data = (char *) t_logits->data + (size_t) k * (size_t) n_vocab * sizeof(float); + + ggml_tensor * emb = ggml_get_rows(ctx.get(), const_cast(head_a), prev_ids); + ggml_tensor * bias = ggml_mul_mat(ctx.get(), const_cast(head_b), emb); + ggml_tensor * logits = ggml_add(ctx.get(), base, bias); + ggml_tensor * sampled = ggml_argmax(ctx.get(), logits); + + ggml_set_output(sampled); + sampled_rows.push_back(sampled); + prev_ids = sampled; + } + + ggml_cgraph * gf = ggml_new_graph_custom(ctx.get(), 64, false); + ggml_build_forward_expand(gf, sampled_rows.back()); + + for (int i = 0; i < ggml_graph_n_nodes(gf); ++i) { + if (!ggml_backend_dev_supports_op(dev, ggml_graph_node(gf, i))) { + return false; + } + } + + ggml_backend_sched_reset(dspark_markov_sched.get()); + if (!ggml_backend_sched_alloc_graph(dspark_markov_sched.get(), gf)) { + return false; + } + + ggml_backend_t ids_backend = ggml_backend_sched_get_tensor_backend(dspark_markov_sched.get(), ids); + ggml_backend_t out_backend = ggml_backend_sched_get_tensor_backend(dspark_markov_sched.get(), sampled_rows.back()); + if (ids_backend == nullptr || out_backend == nullptr || + ggml_backend_get_device(out_backend) != dev) { + return false; + } + + const int32_t id = (int32_t) tok0; + ggml_backend_tensor_set(ids, &id, 0, sizeof(id)); + + const ggml_status status = ggml_backend_sched_graph_compute(dspark_markov_sched.get(), gf); + if (status != GGML_STATUS_SUCCESS) { + return false; + } + + // One scheduler synchronization covers the entire sequential chain. + for (uint32_t k = 0; k < n_chain; ++k) { + int32_t sampled_id = -1; + ggml_backend_tensor_get(sampled_rows[k], &sampled_id, 0, sizeof(sampled_id)); + if (sampled_id < 0 || sampled_id >= n_vocab) { + return false; + } + + result[k0 + k] = (llama_token) sampled_id; + } + + return true; + }; + + // A/B toggle: emulate the pre-fusion behavior — one graph build, scheduler + // submission, synchronization, and host readback per draft step, with the + // sampled token fed back through the host between steps. + if (getenv("DSPARK_MARKOV_PER_STEP") != nullptr) { + llama_token tok = prev_token; + for (uint32_t k = 0; k < n_rows; ++k) { + if (!resample_chain(k, 1, tok)) { + return false; + } + tok = result[k]; + } + return true; + } + + return resample_chain(0, n_rows, prev_token); +} + llm_graph_cb llama_context::graph_get_cb() const { return [&](const llama_ubatch & ubatch, ggml_tensor * cur, const char * name, int il) { if (il >= 0) { @@ -4274,3 +4427,11 @@ llama_memory_breakdown llama_get_memory_breakdown(const struct llama_context * c llama_context * llama_get_ctx_other(struct llama_context * ctx) { return ctx->get_cparams().ctx_other; } + +bool llama_dspark_markov_resample( + struct llama_context * ctx, + int32_t n_rows, + llama_token prev_token, + llama_token * result) { + return ctx != nullptr && ctx->dspark_markov_resample((uint32_t) n_rows, prev_token, result); +} diff --git a/src/llama-context.h b/src/llama-context.h index a1dbdbb06b89..656d7dcf652b 100644 --- a/src/llama-context.h +++ b/src/llama-context.h @@ -257,6 +257,10 @@ struct llama_context { // returns the result of ggml_backend_sched_graph_compute_async execution ggml_status graph_compute(ggml_cgraph * gf, bool batched); + // Run the DSpark vanilla Markov resample on a dedicated backend scheduler, + // leaving the main decode scheduler and its graph allocations untouched. + bool dspark_markov_resample(uint32_t n_rows, llama_token prev_token, llama_token * result); + // reserve a graph with a dummy ubatch of the specified size ggml_cgraph * graph_reserve( uint32_t n_tokens, uint32_t n_seqs, uint32_t n_outputs, const llama_memory_context_i * mctx, bool split_only = false, size_t * sizes = nullptr); @@ -377,6 +381,9 @@ struct llama_context { std::vector backend_buf_exp_size; // expected buffer sizes llm_graph_result_ptr gf_res_prev; + + // dedicated scheduler for the DSpark Metal Markov resample (see dspark_markov_resample) + ggml_backend_sched_ptr dspark_markov_sched; llm_graph_result_ptr gf_res_reserve; // host buffer for the model output (logits and embeddings) diff --git a/src/llama-ext.h b/src/llama-ext.h index 06f227099f13..e5de37dbbf60 100644 --- a/src/llama-ext.h +++ b/src/llama-ext.h @@ -186,3 +186,14 @@ LLAMA_API bool llama_model_dspark_get_markov( const struct llama_model * model, std::vector & w1, std::vector & w2); + +// Run the sequential vanilla Markov resample on the drafter backend. The +// latest decode graph's logits remain device-resident; this helper gathers the +// previous-token row from markov_head_a, multiplies by markov_head_b, adds the +// corresponding logits row, and returns one argmax per block position. Returns +// false when the head/backend is unsupported so callers can use the host path. +LLAMA_API bool llama_dspark_markov_resample( + struct llama_context * ctx, + int32_t n_rows, + llama_token prev_token, + llama_token * result); diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 0d18752defe8..d1c3c79fe33f 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -2727,9 +2727,23 @@ bool llama_model_dspark_get_markov( return true; } default: - LLAMA_LOG_ERROR("%s: unsupported markov head tensor type %s (only f32/f16/bf16 supported)\n", - __func__, ggml_type_name(t->type)); - return false; + if (!ggml_is_quantized(t->type)) { + LLAMA_LOG_ERROR("%s: unsupported markov head tensor type %s\n", + __func__, ggml_type_name(t->type)); + return false; + } + + const auto * qtype = ggml_get_type_traits(t->type); + if (qtype == nullptr || qtype->to_float == nullptr) { + LLAMA_LOG_ERROR("%s: quantized markov head tensor type %s has no dequantizer\n", + __func__, ggml_type_name(t->type)); + return false; + } + + std::vector raw(ggml_nbytes(t)); + ggml_backend_tensor_get(t, raw.data(), 0, raw.size()); + qtype->to_float(raw.data(), out.data(), n); + return true; } };