From 6282fa7ef09f00d2ebc600496508fddd4f3f45c0 Mon Sep 17 00:00:00 2001 From: Mengsheng Wu Date: Thu, 16 Jul 2026 14:35:26 +0800 Subject: [PATCH 1/2] feat(sdk): add MTP speculative decoding to the llama_cpp backend Wire llama.cpp's common_speculative (draft-mtp) into the llama_cpp plugin: load an optional draft/MTP GGUF into an MTP context that shares the target KV cache (pinned to the target's device), and replace the per-token decode loop with a draft-verify loop when a draft model is configured. Draft accept counts are reported in the profile. New public config on geniex_ModelConfig (spec_draft_model, spec_n_draft) and geniex_ProfileData (draft_n_total, draft_n_accepted); mirrored across the Go, Python, and Android FFI surfaces. CLI gains --draft-model / --draft-tokens and prints the acceptance rate. Signed-off-by: Mengsheng Wu --- .../app/src/main/cpp/android_utils.cpp | 39 +-- .../android/app/src/main/cpp/jniutils.cpp | 22 +- .../java/com/geniex/sdk/bean/ModelConfig.kt | 6 + .../java/com/geniex/sdk/bean/ProfilingData.kt | 2 + bindings/go/common.go | 9 + bindings/python/geniex/_ffi/_types.py | 4 + bindings/python/geniex/auto.py | 4 +- bindings/python/geniex/generation/output.py | 7 +- cli/cmd/geniex/common/process.go | 19 +- cli/cmd/geniex/infer.go | 20 +- sdk/include/geniex.h | 9 + sdk/plugins/llama_cpp/include/llm.h | 17 ++ sdk/plugins/llama_cpp/include/profiler.h | 3 + sdk/plugins/llama_cpp/src/llm.cpp | 228 ++++++++++++++++-- sdk/plugins/llama_cpp/src/profiler.cpp | 7 + 15 files changed, 348 insertions(+), 48 deletions(-) diff --git a/bindings/android/app/src/main/cpp/android_utils.cpp b/bindings/android/app/src/main/cpp/android_utils.cpp index 12518aa96..fdb62d930 100644 --- a/bindings/android/app/src/main/cpp/android_utils.cpp +++ b/bindings/android/app/src/main/cpp/android_utils.cpp @@ -16,7 +16,7 @@ #define MAX_PATH_LEN 512 namespace geniex_android_sdk { -void throw_runtime_exception(JNIEnv *env, const char *_Nonnull format, ...) { +void throw_runtime_exception(JNIEnv* env, const char* _Nonnull format, ...) { va_list va; va_start(va, format); char errMsg[128]; @@ -29,12 +29,12 @@ void throw_runtime_exception(JNIEnv *env, const char *_Nonnull format, ...) { } } -jobject create_java_profile_data(JNIEnv *env, geniex_ProfileData data) { +jobject create_java_profile_data(JNIEnv* env, geniex_ProfileData data) { jclass cls = env->FindClass("com/geniex/sdk/bean/ProfilingData"); if (!cls) return nullptr; - // (DDDJJJDDDLjava/lang/String;)V - jmethodID ctor = env->GetMethodID(cls, "", "(DDDJJJDDDLjava/lang/String;)V"); + // (DDDJJJDDDJJLjava/lang/String;)V + jmethodID ctor = env->GetMethodID(cls, "", "(DDDJJJDDDJJLjava/lang/String;)V"); if (!ctor) return nullptr; const auto ttft_ms = static_cast(data.ttft / 1000.0); @@ -60,6 +60,9 @@ jobject create_java_profile_data(JNIEnv *env, geniex_ProfileData data) { LOGd("prefill_speed=%.6f tok/s, decoding_speed=%.6f tok/s, rtf=%.4f", prefill_speed, decoding_speed, rtf); + const auto draft_n_total = static_cast(data.draft_n_total); + const auto draft_n_accepted = static_cast(data.draft_n_accepted); + jstring jStopReason = env->NewStringUTF(data.stop_reason ? data.stop_reason : ""); jobject obj = env->NewObject(cls, @@ -72,15 +75,17 @@ jobject create_java_profile_data(JNIEnv *env, geniex_ProfileData data) { audio_ms, // J J J prefill_speed, decoding_speed, - rtf, // D D D - jStopReason // String + rtf, // D D D + draft_n_total, + draft_n_accepted, // J J + jStopReason // String ); env->DeleteLocalRef(jStopReason); return obj; } -bool check_jni_exception(JNIEnv *env, const char *where) { +bool check_jni_exception(JNIEnv* env, const char* where) { if (env->ExceptionCheck()) { LOGe("Exception at %s", where); env->ExceptionDescribe(); @@ -98,7 +103,7 @@ bool check_jni_exception(JNIEnv *env, const char *where) { * @param fieldName * @return */ -const char *getStringField(JNIEnv *env, jclass cls, jobject inputObj, const char *fieldName) { +const char* getStringField(JNIEnv* env, jclass cls, jobject inputObj, const char* fieldName) { jfieldID fid = env->GetFieldID(cls, fieldName, "Ljava/lang/String;"); if (check_jni_exception(env, "GetFieldID failed") || !fid) { LOGe("field '%s' not found", fieldName); @@ -111,7 +116,7 @@ const char *getStringField(JNIEnv *env, jclass cls, jobject inputObj, const char } std::string s = jniutils::jstring2str(env, jstr); env->DeleteLocalRef(jstr); - const char *c = jniutils::hold_c_str(s); + const char* c = jniutils::hold_c_str(s); LOGd("%s = %s", fieldName, c); return c; } @@ -124,7 +129,7 @@ const char *getStringField(JNIEnv *env, jclass cls, jobject inputObj, const char * @param fieldName * @return */ -jint getIntField(JNIEnv *env, jclass cls, jobject obj, const char *fieldName) { +jint getIntField(JNIEnv* env, jclass cls, jobject obj, const char* fieldName) { jfieldID fieldId = env->GetFieldID(cls, fieldName, "Ljava/lang/Integer;"); if (check_jni_exception(env, "GetFieldID failed") || !fieldId) { LOGe("field '%s' not found", fieldName); @@ -155,7 +160,7 @@ jint getIntField(JNIEnv *env, jclass cls, jobject obj, const char *fieldName) { * @param fieldName * @return */ -jfloat getFloatField(JNIEnv *env, jclass cls, jobject obj, const char *fieldName) { +jfloat getFloatField(JNIEnv* env, jclass cls, jobject obj, const char* fieldName) { jfieldID fieldId = env->GetFieldID(cls, fieldName, "Ljava/lang/Float;"); if (!fieldId) { LOGe("field '%s' not found", fieldName); @@ -178,7 +183,7 @@ jfloat getFloatField(JNIEnv *env, jclass cls, jobject obj, const char *fieldName return result; } -jobject getObjectField(JNIEnv *env, jclass cls, jobject obj, const char *name, const char *sig) { +jobject getObjectField(JNIEnv* env, jclass cls, jobject obj, const char* name, const char* sig) { jfieldID configId = env->GetFieldID(cls, name, sig); if (check_jni_exception(env, "GetFieldID failed") || !configId) { LOGe("field '%s' not found", name); @@ -194,7 +199,7 @@ jobject getObjectField(JNIEnv *env, jclass cls, jobject obj, const char *name, c } } -jboolean getBoolField(JNIEnv *env, jclass cls, jobject obj, const char *fieldName) { +jboolean getBoolField(JNIEnv* env, jclass cls, jobject obj, const char* fieldName) { jfieldID fieldId = env->GetFieldID(cls, fieldName, "Z"); if (check_jni_exception(env, "GetFieldID failed") || !fieldId) { LOGe("field '%s' not found", fieldName); @@ -213,7 +218,7 @@ jboolean getBoolField(JNIEnv *env, jclass cls, jobject obj, const char *fieldNam * Get primitive float field (for non-nullable Kotlin Float) * JNI signature: "F" */ -jfloat getPrimitiveFloatField(JNIEnv *env, jclass cls, jobject obj, const char *fieldName) { +jfloat getPrimitiveFloatField(JNIEnv* env, jclass cls, jobject obj, const char* fieldName) { jfieldID fieldId = env->GetFieldID(cls, fieldName, "F"); if (check_jni_exception(env, "GetFieldID failed") || !fieldId) { LOGe("primitive float field '%s' not found", fieldName); @@ -232,7 +237,7 @@ jfloat getPrimitiveFloatField(JNIEnv *env, jclass cls, jobject obj, const char * * Get primitive int field (for non-nullable Kotlin Int) * JNI signature: "I" */ -jint getPrimitiveIntField(JNIEnv *env, jclass cls, jobject obj, const char *fieldName) { +jint getPrimitiveIntField(JNIEnv* env, jclass cls, jobject obj, const char* fieldName) { jfieldID fieldId = env->GetFieldID(cls, fieldName, "I"); if (check_jni_exception(env, "GetFieldID failed") || !fieldId) { LOGe("primitive int field '%s' not found", fieldName); @@ -247,7 +252,7 @@ jint getPrimitiveIntField(JNIEnv *env, jclass cls, jobject obj, const char *fiel return result; } -jobject create_string_list(JNIEnv *env, const char **strings, int count) { +jobject create_string_list(JNIEnv* env, const char** strings, int count) { jclass array_list_class = env->FindClass("java/util/ArrayList"); jmethodID array_list_constructor = env->GetMethodID(array_list_class, "", "()V"); jmethodID array_list_add = env->GetMethodID(array_list_class, "add", "(Ljava/lang/Object;)Z"); @@ -265,7 +270,7 @@ jobject create_string_list(JNIEnv *env, const char **strings, int count) { return list; } -jobject create_float_list(JNIEnv *env, float *floats, int count) { +jobject create_float_list(JNIEnv* env, float* floats, int count) { jclass array_list_class = env->FindClass("java/util/ArrayList"); jmethodID array_list_constructor = env->GetMethodID(array_list_class, "", "()V"); jmethodID array_list_add = env->GetMethodID(array_list_class, "add", "(Ljava/lang/Object;)Z"); diff --git a/bindings/android/app/src/main/cpp/jniutils.cpp b/bindings/android/app/src/main/cpp/jniutils.cpp index d5a086519..64c4ce071 100644 --- a/bindings/android/app/src/main/cpp/jniutils.cpp +++ b/bindings/android/app/src/main/cpp/jniutils.cpp @@ -258,14 +258,23 @@ geniex_ModelConfig extract_model_config(JNIEnv* env, jobject configObj) { fid = env->GetFieldID(cls, "verbose", "Z"); config.verbose = env->GetBooleanField(configObj, fid); + // spec_draft_model + fid = env->GetFieldID(cls, "spec_draft_model", "Ljava/lang/String;"); + jstr = (jstring)env->GetObjectField(configObj, fid); + config.spec_draft_model = jstr ? hold_c_str(jstring2str(env, jstr)) : nullptr; + + // spec_n_draft + fid = env->GetFieldID(cls, "spec_n_draft", "I"); + config.spec_n_draft = env->GetIntField(configObj, fid); + return config; } jobject extract_profiling_data(JNIEnv* env, const geniex_ProfileData& data) { jclass cls = env->FindClass("com/geniex/sdk/bean/ProfilingData"); if (!cls) return nullptr; - // (DDDJJJDDDLjava/lang/String;)V - jmethodID ctor = env->GetMethodID(cls, "", "(DDDJJJDDDLjava/lang/String;)V"); + // (DDDJJJDDDJJLjava/lang/String;)V + jmethodID ctor = env->GetMethodID(cls, "", "(DDDJJJDDDJJLjava/lang/String;)V"); if (!ctor) return nullptr; const jdouble ttft_ms = static_cast(data.ttft / 1000.0); @@ -291,6 +300,9 @@ jobject extract_profiling_data(JNIEnv* env, const geniex_ProfileData& data) { LOGe("prefill_speed=%.6f tok/s, decoding_speed=%.6f tok/s, rtf=%.4f", prefill_speed, decoding_speed, rtf); + const jlong draft_n_total = static_cast(data.draft_n_total); + const jlong draft_n_accepted = static_cast(data.draft_n_accepted); + jstring jStopReason = env->NewStringUTF(data.stop_reason ? data.stop_reason : ""); jobject obj = env->NewObject(cls, @@ -303,8 +315,10 @@ jobject extract_profiling_data(JNIEnv* env, const geniex_ProfileData& data) { audio_ms, // J J J prefill_speed, decoding_speed, - rtf, // D D D - jStopReason // String + rtf, // D D D + draft_n_total, + draft_n_accepted, // J J + jStopReason // String ); env->DeleteLocalRef(jStopReason); diff --git a/bindings/android/app/src/main/java/com/geniex/sdk/bean/ModelConfig.kt b/bindings/android/app/src/main/java/com/geniex/sdk/bean/ModelConfig.kt index 764f199f4..8c923bff3 100644 --- a/bindings/android/app/src/main/java/com/geniex/sdk/bean/ModelConfig.kt +++ b/bindings/android/app/src/main/java/com/geniex/sdk/bean/ModelConfig.kt @@ -44,4 +44,10 @@ data class ModelConfig( val enable_thinking: Boolean = false, val verbose: Boolean = false, + + /** Path to a draft / MTP model GGUF to enable speculative decoding (llama_cpp only; "" = disabled) */ + val spec_draft_model: String = "", + + /** Draft tokens per step for speculative decoding (0 = plugin default) */ + val spec_n_draft: Int = 0, ) diff --git a/bindings/android/app/src/main/java/com/geniex/sdk/bean/ProfilingData.kt b/bindings/android/app/src/main/java/com/geniex/sdk/bean/ProfilingData.kt index efbf3b5e4..abce5024e 100644 --- a/bindings/android/app/src/main/java/com/geniex/sdk/bean/ProfilingData.kt +++ b/bindings/android/app/src/main/java/com/geniex/sdk/bean/ProfilingData.kt @@ -10,5 +10,7 @@ data class ProfilingData( val prefillSpeed: Double, /* Prefill speed (tokens/sec) */ val decodingSpeed: Double, /* Decoding speed (tokens/sec) */ val realTimeFactor: Double, /* Real-Time Factor(RTF) (1.0 = real-time, >1.0 = faster, <1.0 = slower) */ + val draftNTotal: Long, /* Speculative decoding: draft tokens generated (0 when disabled) */ + val draftNAccepted: Long, /* Speculative decoding: draft tokens accepted by the target model */ val stopReason: String /* Stop reason: "eos", "length", "user", "stop_sequence" */ ) diff --git a/bindings/go/common.go b/bindings/go/common.go index 9cbaf3c04..6691604b2 100644 --- a/bindings/go/common.go +++ b/bindings/go/common.go @@ -23,6 +23,8 @@ type ProfileData struct { PrefillSpeed float64 DecodingSpeed float64 RealTimeFactor float64 + DraftNTotal int64 + DraftNAccepted int64 StopReason string } @@ -46,6 +48,8 @@ func newProfileDataFromCPtr(c C.geniex_ProfileData) ProfileData { PrefillSpeed: float64(c.prefill_speed), DecodingSpeed: float64(c.decoding_speed), RealTimeFactor: float64(c.real_time_factor), + DraftNTotal: int64(c.draft_n_total), + DraftNAccepted: int64(c.draft_n_accepted), StopReason: C.GoString(c.stop_reason), } } @@ -163,6 +167,8 @@ type ModelConfig struct { NGpuLayers int32 ChatTemplatePath string ChatTemplateContent string + SpecDraftModel string + SpecNDraft int32 } // fillC writes mc into an embedded C struct; pair with freeCModelConfig to @@ -179,6 +185,8 @@ func (mc ModelConfig) fillC(out *C.geniex_ModelConfig) { n_gpu_layers: C.int32_t(mc.NGpuLayers), chat_template_path: cStringIfSet(mc.ChatTemplatePath), chat_template_content: cStringIfSet(mc.ChatTemplateContent), + spec_draft_model: cStringIfSet(mc.SpecDraftModel), + spec_n_draft: C.int32_t(mc.SpecNDraft), } } @@ -188,6 +196,7 @@ func freeCModelConfig(c *C.geniex_ModelConfig) { } cFreeIfSet(unsafe.Pointer(c.chat_template_path)) cFreeIfSet(unsafe.Pointer(c.chat_template_content)) + cFreeIfSet(unsafe.Pointer(c.spec_draft_model)) } // LCOV_EXCL_STOP diff --git a/bindings/python/geniex/_ffi/_types.py b/bindings/python/geniex/_ffi/_types.py index d53114ee0..ef105c6de 100644 --- a/bindings/python/geniex/_ffi/_types.py +++ b/bindings/python/geniex/_ffi/_types.py @@ -40,6 +40,8 @@ class geniex_ProfileData(Structure): ('prefill_speed', c_double), ('decoding_speed', c_double), ('real_time_factor', c_double), + ('draft_n_total', c_int64), + ('draft_n_accepted', c_int64), ('stop_reason', c_char_p), ] @@ -109,6 +111,8 @@ class geniex_ModelConfig(Structure): ('max_tokens', c_int32), ('enable_thinking', c_bool), ('verbose', c_bool), + ('spec_draft_model', c_char_p), + ('spec_n_draft', c_int32), ] diff --git a/bindings/python/geniex/auto.py b/bindings/python/geniex/auto.py index b050e96a5..e1d8ae15d 100644 --- a/bindings/python/geniex/auto.py +++ b/bindings/python/geniex/auto.py @@ -221,9 +221,9 @@ def _build_model_config(plugin_id: str | None, n_ctx: int, n_gpu_layers: int, ** _logger.warning('qairt runtime does not consume n_ctx=%d; forcing 0', n_ctx) n_ctx = 0 cfg = geniex_ModelConfig(n_ctx=n_ctx, n_gpu_layers=n_gpu_layers) - _int_fields = {'n_threads', 'n_threads_batch', 'n_batch', 'n_ubatch', 'n_seq_max', 'max_tokens'} + _int_fields = {'n_threads', 'n_threads_batch', 'n_batch', 'n_ubatch', 'n_seq_max', 'max_tokens', 'spec_n_draft'} _bool_fields = {'enable_thinking', 'verbose'} - _str_fields = {'chat_template_path', 'chat_template_content', 'system_prompt'} + _str_fields = {'chat_template_path', 'chat_template_content', 'system_prompt', 'spec_draft_model'} for k, v in kwargs.items(): if k in _int_fields: setattr(cfg, k, int(v)) diff --git a/bindings/python/geniex/generation/output.py b/bindings/python/geniex/generation/output.py index de8cb7a32..bb368734e 100644 --- a/bindings/python/geniex/generation/output.py +++ b/bindings/python/geniex/generation/output.py @@ -30,6 +30,8 @@ class ProfileData: generated_tokens: int = 0 prefill_speed: float = 0.0 decode_speed: float = 0.0 + draft_n_total: int = 0 + draft_n_accepted: int = 0 stop_reason: str | None = None backend: str | None = None device: str | None = None @@ -47,6 +49,8 @@ def from_c(cls, c: geniex_ProfileData) -> 'ProfileData': generated_tokens=c.generated_tokens, prefill_speed=c.prefill_speed, decode_speed=c.decoding_speed, + draft_n_total=c.draft_n_total, + draft_n_accepted=c.draft_n_accepted, stop_reason=stop, ) @@ -60,7 +64,8 @@ def __repr__(self) -> str: f'generated_tokens={self.generated_tokens} tok, ' f'prefill_speed={self.prefill_speed:.1f} tok/s, ' f'decode_speed={self.decode_speed:.1f} tok/s, ' - f'stop_reason={self.stop_reason}, ' + + (f'draft_accept={self.draft_n_accepted}/{self.draft_n_total}, ' if self.draft_n_total > 0 else '') + + f'stop_reason={self.stop_reason}, ' f'backend={self.backend}, ' f'device={self.device}, ' f'quant={self.quant}, ' diff --git a/cli/cmd/geniex/common/process.go b/cli/cmd/geniex/common/process.go index bff6341f6..e2bbfdf8d 100644 --- a/cli/cmd/geniex/common/process.go +++ b/cli/cmd/geniex/common/process.go @@ -242,10 +242,10 @@ func (p *Processor) fsmInit() { {STATE_START, "assistant"}: {STATE_ASSISTANT, nil}, // gemma4 - {STATE_ASSISTANT, "<|channel>"}: {STATE_GEMMA_CHANNEL, nil}, - {STATE_GEMMA_CHANNEL, "thought"}: {STATE_GEMMA_CHANNEL_THOUGHT, nil}, + {STATE_ASSISTANT, "<|channel>"}: {STATE_GEMMA_CHANNEL, nil}, + {STATE_GEMMA_CHANNEL, "thought"}: {STATE_GEMMA_CHANNEL_THOUGHT, nil}, {STATE_GEMMA_CHANNEL_THOUGHT, "\n"}: {STATE_THINK, thinkStart(true)}, - {STATE_THINK, ""}: {STATE_NORMAL, thinkEnd(true)}, + {STATE_THINK, ""}: {STATE_NORMAL, thinkEnd(true)}, } p.fsmState = STATE_ASSISTANT } @@ -292,6 +292,13 @@ stop reason: %s pd.StopReason, ) + if pd.DraftNTotal > 0 { + text += fmt.Sprintf("\ndraft accept: %d/%d (%.1f%%)", + pd.DraftNAccepted, + pd.DraftNTotal, + 100.0*float64(pd.DraftNAccepted)/float64(pd.DraftNTotal)) + } + } else { if pd.AudioDuration > 0 { // ASR TTS text = fmt.Sprintf("processing_time %.2fs | audio_duration %.2fs | RTF %.2f (%.1fx realtime)", @@ -301,10 +308,14 @@ stop reason: %s 1.0/pd.RealTimeFactor) } else if pd.DecodingSpeed != 0 { - text = fmt.Sprintf("— %.1f tok/s • %d tok • %.1f s first token —", + text = fmt.Sprintf("— %.1f tok/s • %d tok • %.1f s first token", pd.DecodingSpeed, pd.GeneratedTokens, float64(pd.TTFT)/1e6) + if pd.DraftNTotal > 0 { + text += fmt.Sprintf(" • %.0f%% accept", 100.0*float64(pd.DraftNAccepted)/float64(pd.DraftNTotal)) + } + text += " —" } else { if pd.TotalTimeUs() != 0 { diff --git a/cli/cmd/geniex/infer.go b/cli/cmd/geniex/infer.go index af2313bfc..9084365b6 100644 --- a/cli/cmd/geniex/infer.go +++ b/cli/cmd/geniex/infer.go @@ -39,6 +39,8 @@ var ( systemPrompt string computeUnit string slidingWindow bool + draftModel string + draftTokens int32 // sampler config temperature float32 @@ -87,6 +89,8 @@ var ( llmFlags.StringArrayVarP(&prompt, "prompt", "p", nil, "pass prompt") llmFlags.StringVarP(&tokenFile, "token-file", "t", "", "path to token file (space-separated token IDs) (llama_cpp only)") llmFlags.BoolVarP(&slidingWindow, "sliding-window", "", false, "evict oldest context on overflow instead of erroring (qairt only)") + llmFlags.StringVarP(&draftModel, "draft-model", "", "", "path to a draft/MTP model GGUF to enable speculative decoding (llama_cpp only)") + llmFlags.Int32VarP(&draftTokens, "draft-tokens", "", 3, "draft tokens per step for speculative decoding (llama_cpp only)") return llmFlags }() vlmFlags = func() *pflag.FlagSet { @@ -264,6 +268,16 @@ func inferLLM(paths *geniex_sdk.ModelPaths) error { return err } + // Speculative decoding is llama_cpp-only; ignore the draft model for other runtimes. + specDraftModel := draftModel + if paths.RuntimeID != geniex_sdk.RuntimeLlamaCpp { + if specDraftModel != "" { + fmt.Println(render.GetTheme().Warning.Sprintf( + "Warning: --draft-model is only supported by llama_cpp; ignoring for runtime %s", paths.RuntimeID)) + } + specDraftModel = "" + } + spin := render.NewSpinner("loading model...") spin.Start() @@ -273,8 +287,10 @@ func inferLLM(paths *geniex_sdk.ModelPaths) error { RuntimeID: paths.RuntimeID, DeviceID: deviceID, Config: geniex_sdk.ModelConfig{ - NCtx: nctxResolved, - NGpuLayers: nglResolved, + NCtx: nctxResolved, + NGpuLayers: nglResolved, + SpecDraftModel: specDraftModel, + SpecNDraft: draftTokens, }, }) spin.Stop() diff --git a/sdk/include/geniex.h b/sdk/include/geniex.h index 28c421279..d66f00ac9 100644 --- a/sdk/include/geniex.h +++ b/sdk/include/geniex.h @@ -363,6 +363,9 @@ typedef struct { double decoding_speed; /* Decoding speed (tokens/sec) */ double real_time_factor; /* Real-Time Factor(RTF) (1.0 = real-time, >1.0 = faster, <1.0 = slower) */ + int64_t draft_n_total; /* Speculative decoding: draft tokens generated (0 when disabled) */ + int64_t draft_n_accepted; /* Speculative decoding: draft tokens accepted by the target model */ + const char* stop_reason; /* Stop reason: "eos", "length", "user", "stop_sequence", "context_length" */ } geniex_ProfileData; @@ -425,6 +428,12 @@ typedef struct { int32_t max_tokens; // max tokens to generate bool enable_thinking; // enable thinking mode for Qwen models bool verbose; // verbose logging + + // Speculative decoding (llama_cpp only; ignored by qairt). Enabled when + // spec_draft_model is a non-empty path to a draft/MTP GGUF (e.g. a Gemma-4 + // `mtp-` assistant). Other backends ignore these fields. + geniex_Path spec_draft_model; // path to the draft / MTP head GGUF ("" = disabled) + int32_t spec_n_draft; // draft tokens per step (0 = plugin default of 3) } geniex_ModelConfig; /* ==================== LLM Handle ======================================== */ diff --git a/sdk/plugins/llama_cpp/include/llm.h b/sdk/plugins/llama_cpp/include/llm.h index 75ae2ce21..4e69eeac9 100644 --- a/sdk/plugins/llama_cpp/include/llm.h +++ b/sdk/plugins/llama_cpp/include/llm.h @@ -3,12 +3,17 @@ #pragma once +#include #include +#include #include "htp_session.h" #include "llama.h" +#include "params.h" #include "plugin/ILlm.h" +#include "profiler.h" #include "sampling.h" +#include "speculative.h" #include "threadpool.h" namespace geniex { @@ -20,6 +25,12 @@ class LlamaLlm : public ILlm { Threadpools pools_; std::optional chat_template_str = std::nullopt; + // Speculative decoding (MTP). All null unless a draft model was loaded. + llama_model* draft_model = nullptr; + llama_context* draft_ctx = nullptr; + common_speculative* spec = nullptr; + int32_t spec_n_max = 0; + int n_past_global = 0; int n_past = 0; // for context shifting std::vector past_prompt_tokens; // for prefix match @@ -48,6 +59,12 @@ class LlamaLlm : public ILlm { private: void set_sampler(const geniex_SamplerConfig* cfg); + + int32_t setup_speculative(const geniex_ModelConfig& config, Device device, const char* device_id); + void teardown_speculative(); + int32_t decode_speculative(const geniex_GenerationConfig& cfg, const std::vector& prompt_ids, + const std::function& emit, const std::function& n_generated, + ::common::Profiler& profiler); }; } // namespace geniex diff --git a/sdk/plugins/llama_cpp/include/profiler.h b/sdk/plugins/llama_cpp/include/profiler.h index 355c49bff..cd663741d 100644 --- a/sdk/plugins/llama_cpp/include/profiler.h +++ b/sdk/plugins/llama_cpp/include/profiler.h @@ -32,6 +32,7 @@ class Profiler { void record_ttft(); void update_prompt_tokens(uint32_t); void update_generated_tokens(uint32_t); + void set_draft_stats(int64_t n_total, int64_t n_accepted); void set_stop_reason(StopReason); StopReason get_stop_reason() const; void end(); @@ -56,6 +57,8 @@ class Profiler { StopReason stop_reason = StopReason::GENIEX_STOP_REASON_UNKNOWN; uint32_t prompt_tokens = 0; uint32_t generated_tokens = 0; + int64_t draft_n_total = 0; + int64_t draft_n_accepted = 0; }; } // namespace common diff --git a/sdk/plugins/llama_cpp/src/llm.cpp b/sdk/plugins/llama_cpp/src/llm.cpp index f26d9ee81..004347cd5 100644 --- a/sdk/plugins/llama_cpp/src/llm.cpp +++ b/sdk/plugins/llama_cpp/src/llm.cpp @@ -20,7 +20,10 @@ namespace geniex { LlamaLlm::~LlamaLlm() { + if (spec) common_speculative_free(spec); if (sampler) common_sampler_free(sampler); + if (draft_ctx) llama_free(draft_ctx); + if (draft_model) llama_model_free(draft_model); if (ctx) llama_free(ctx); if (model) llama_model_free(model); // pools_ frees its threadpools in its own destructor, after ctx is freed. @@ -94,6 +97,16 @@ int32_t LlamaLlm::create_impl(const geniex_LlmCreateInput* input) { return tp_ret; } + // Speculative decoding (MTP): optional, keyed on a non-empty draft model path. + // Failure to set it up is non-fatal — we log and fall back to plain decoding. + if (config.spec_draft_model && config.spec_draft_model[0] != '\0') { + int32_t spec_ret = setup_speculative(config, device, input->device_id); + if (spec_ret != GENIEX_SUCCESS) { + GENIEX_LOG_WARN("speculative decoding setup failed; falling back to plain decoding"); + teardown_speculative(); + } + } + // Load chat template if path is provided if (config.chat_template_content) { try { @@ -345,54 +358,61 @@ int32_t LlamaLlm::generate(const geniex_LlmGenerateInput* input, geniex_LlmGener std::vector generated_tokens; std::stringstream full_text; - while (res == GENIEX_SUCCESS && (int)generated_tokens.size() < cfg.max_tokens) { - llama_token id = common_sampler_sample(this->sampler, this->ctx, -1); - common_sampler_accept(this->sampler, id, /* accept_grammar= */ true); - - // Record TTFT on first token generation + // Emit one sampled token: records TTFT, applies EOS / stop-sequence / user + // callback checks, and appends to the output. Returns false when generation + // should stop (the stop reason is set on the profiler). Shared by the plain + // and speculative decode loops. + auto emit = [&](llama_token id) -> bool { if (!first_token_generated) { profiler.record_ttft(); first_token_generated = true; - GENIEX_LOG_DEBUG("First token generated, TTFT recorded"); } - // Check EOS token if (llama_vocab_is_eog(vocab, id)) { - GENIEX_LOG_DEBUG("EOS token generated, stopping generation"); profiler.set_stop_reason(common::StopReason::GENIEX_STOP_REASON_EOS); - break; + return false; } - // Convert token to string char token_buf[64]; int n = llama_token_to_piece(vocab, id, token_buf, sizeof(token_buf) - 1, 0, this->allow_special_tokens); if (n < 0) { res = GENIEX_ERROR_LLM_GENERATION_FAILED; - break; + return false; } token_buf[n] = '\0'; - // Check stop sequences const bool stop_matched = std::any_of( cfg.stop, cfg.stop + cfg.stop_count, [&](const char* s) { return s && strcmp(token_buf, s) == 0; }); if (stop_matched) { - GENIEX_LOG_DEBUG("Stop sequence matched"); profiler.set_stop_reason(common::StopReason::GENIEX_STOP_REASON_STOP_SEQUENCE); - break; + return false; } generated_tokens.push_back(id); - // Call the callback directly (UTF-8 validation is now handled at bridge - // level) if (input->on_token && !input->on_token(token_buf, input->user_data)) { GENIEX_LOG_WARN("User callback requested stop during token generation"); profiler.set_stop_reason(common::StopReason::GENIEX_STOP_REASON_USER); - break; + return false; } full_text << token_buf; + return true; + }; + + if (this->spec) { + auto n_generated = [&]() { return (int)generated_tokens.size(); }; + res = decode_speculative(cfg, prompt_ids, emit, n_generated, profiler); + } else { + while (res == GENIEX_SUCCESS && (int)generated_tokens.size() < cfg.max_tokens) { + llama_token id = common_sampler_sample(this->sampler, this->ctx, -1); + common_sampler_accept(this->sampler, id, /* accept_grammar= */ true); - res = process(&id, 1); + if (!emit(id)) { + break; + } + + res = process(&id, 1); + } } // update output and profiler data @@ -429,6 +449,111 @@ int32_t LlamaLlm::get_model_info(geniex_LlmModelInfo* output) { // Private namespace geniex { +// Speculative (MTP) decode loop. Each step the MTP head drafts up to spec_n_max +// tokens, the target verifies them in one batch, and the accepted prefix is +// committed at once. Mirrors llama.cpp's speculative-simple accounting for a +// single sequence. Assumes the whole prompt was already prefilled on this->ctx +// and this->n_past is the number of prefilled tokens. +// +// id_last is the running committed token that is not yet in the KV cache. It is +// re-decoded (with the drafts) each step and emitted at the top of the next +// step, so every produced token is emitted exactly once. +// +// Partial acceptance relies on plain KV-tail removal (llama_memory_seq_rm), +// which the CPU/GPU/HTP memory backends we target support; the checkpoint dance +// the server uses for recurrent contexts is intentionally omitted. +int32_t LlamaLlm::decode_speculative(const geniex_GenerationConfig& cfg, const std::vector& prompt_ids, + const std::function& emit, const std::function& n_generated, common::Profiler& profiler) { + const llama_seq_id seq_id = 0; + auto* mem_tgt = llama_get_memory(this->ctx); + auto* mem_dft = llama_get_memory(this->draft_ctx); + + // The MTP head reads the target's nextn embeddings; enable them for every + // decode. This is a static property of the speculator, so set it once. + llama_set_embeddings(this->ctx, common_speculative_need_embd_nextn(this->spec)); + + // Local view of the committed tokens the drafter reads; grows as we accept. + std::vector prompt = prompt_ids; + + common_speculative_begin(this->spec, seq_id, prompt); + + // Sample the first token from the prefill; keep it separate as id_last. + llama_token id_last = common_sampler_sample(this->sampler, this->ctx, -1); + common_sampler_accept(this->sampler, id_last, /* accept_grammar= */ true); + + llama_batch batch = llama_batch_init(this->spec_n_max + 1, /*embd=*/0, /*n_seq_max=*/1); + + int64_t draft_n_total = 0; + int64_t draft_n_accepted = 0; + int32_t res = GENIEX_SUCCESS; + bool stop = false; + std::vector draft; + + while (!stop && res == GENIEX_SUCCESS && n_generated() < cfg.max_tokens) { + // Emit the running committed token (always valid: sampled by the target). + if (!emit(id_last)) { + break; + } + + // Draft the tokens that (probably) follow id_last. + draft.clear(); + common_speculative_get_draft_params(this->spec, seq_id) = { + /* .drafting = */ true, + /* .n_max = */ this->spec_n_max, + /* .n_past = */ this->n_past, + /* .id_last = */ id_last, + /* .prompt = */ &prompt, + /* .result = */ &draft, + }; + common_speculative_draft(this->spec); + draft_n_total += (int64_t)draft.size(); + + // Verification batch: [id_last, draft0, draft1, ...], all needing logits. + common_batch_clear(batch); + llama_pos pos = this->n_past; + common_batch_add(batch, id_last, pos++, {seq_id}, /*logits=*/true); + for (llama_token t : draft) { + common_batch_add(batch, t, pos++, {seq_id}, /*logits=*/true); + } + + if (llama_decode(this->ctx, batch) != 0 || !common_speculative_process(this->spec, batch)) { + res = GENIEX_ERROR_LLM_GENERATION_FAILED; + break; + } + + // Accept the longest draft prefix the target agrees with. ids always has + // at least one entry (the target's own next token); ids.size()-1 drafts + // were accepted. + std::vector ids = common_sampler_sample_and_accept_n(this->sampler, this->ctx, draft); + const size_t n_accept = ids.size() - 1; + draft_n_accepted += (int64_t)n_accept; + common_speculative_accept(this->spec, seq_id, (uint16_t)n_accept); + + // Commit id_last + accepted drafts. Emit the accepted drafts now; the + // last id becomes the next id_last, emitted at the top of the next step. + prompt.push_back(id_last); + for (size_t i = 0; i < n_accept; ++i) { + if (!emit(ids[i])) { + stop = true; + break; + } + prompt.push_back(ids[i]); + } + this->n_past += (int)n_accept + 1; + + // Drop any rejected draft tail from both KV caches. + llama_memory_seq_rm(mem_tgt, seq_id, this->n_past, -1); + llama_memory_seq_rm(mem_dft, seq_id, this->n_past, -1); + + id_last = ids.back(); + } + + llama_batch_free(batch); + + profiler.set_draft_stats(draft_n_total, draft_n_accepted); + return res; +} + void LlamaLlm::set_sampler(const geniex_SamplerConfig* cfg) { if (this->sampler) { common_sampler_free(this->sampler); @@ -438,4 +563,71 @@ void LlamaLlm::set_sampler(const geniex_SamplerConfig* cfg) { this->sampler = common_sampler_init(this->model, s); } +void LlamaLlm::teardown_speculative() { + if (this->spec) { + common_speculative_free(this->spec); + this->spec = nullptr; + } + if (this->draft_ctx) { + llama_free(this->draft_ctx); + this->draft_ctx = nullptr; + } + if (this->draft_model) { + llama_model_free(this->draft_model); + this->draft_model = nullptr; + } + this->spec_n_max = 0; +} + +// Load the draft (MTP head) model and wire an MTP draft context that shares the +// target KV cache. Mirrors the manual setup in llama.cpp's server-context.cpp +// for the `--spec-type draft-mtp` + separate-drafter-GGUF path. +int32_t LlamaLlm::setup_speculative(const geniex_ModelConfig& config, Device device, const char* device_id) { + llama_model_params dmpar = build_model_params(config, device); + + // Pin the drafter to the same device(s) as the target; without this it falls + // back to the default placement (which can grab HTP and break the MTP graph). + auto selection = resolve_devices(device_id); + if (selection && !selection->empty()) { + dmpar.devices = selection->data(); + } + + this->draft_model = llama_model_load_from_file(config.spec_draft_model, dmpar); + if (!this->draft_model) { + GENIEX_LOG_ERROR("failed to load draft model: {}", config.spec_draft_model); + return GENIEX_ERROR_COMMON_MODEL_LOAD; + } + + // MTP draft context: same context tuning as the target, plus the MTP wiring + // (shared KV via ctx_other, no rollback snapshots). + llama_context_params dcpar = build_context_params(config, /*n_ctx_default=*/4096, device); + dcpar.ctx_type = LLAMA_CONTEXT_TYPE_MTP; + dcpar.ctx_other = this->ctx; + dcpar.n_rs_seq = 0; + + this->draft_ctx = llama_init_from_model(this->draft_model, dcpar); + if (!this->draft_ctx) { + GENIEX_LOG_ERROR("failed to create MTP draft context"); + return GENIEX_ERROR_COMMON_MODEL_LOAD; + } + + this->spec_n_max = config.spec_n_draft > 0 ? config.spec_n_draft : 3; + + common_params_speculative spar; + spar.types = {COMMON_SPECULATIVE_TYPE_DRAFT_MTP}; + spar.draft.n_max = this->spec_n_max; + spar.draft.ctx_tgt = this->ctx; + spar.draft.ctx_dft = this->draft_ctx; + + this->spec = common_speculative_init(spar, /*n_seq=*/1); + if (!this->spec) { + GENIEX_LOG_ERROR("failed to initialize speculative context"); + return GENIEX_ERROR_COMMON_MODEL_LOAD; + } + + GENIEX_LOG_INFO( + "speculative decoding (MTP) enabled: draft_model={}, n_draft={}", config.spec_draft_model, this->spec_n_max); + return GENIEX_SUCCESS; +} + } // namespace geniex diff --git a/sdk/plugins/llama_cpp/src/profiler.cpp b/sdk/plugins/llama_cpp/src/profiler.cpp index 053bd2a50..e2d807c94 100644 --- a/sdk/plugins/llama_cpp/src/profiler.cpp +++ b/sdk/plugins/llama_cpp/src/profiler.cpp @@ -43,6 +43,11 @@ void Profiler::update_prompt_tokens(uint32_t count) { prompt_tokens = count; } void Profiler::update_generated_tokens(uint32_t count) { generated_tokens = count; } +void Profiler::set_draft_stats(int64_t n_total, int64_t n_accepted) { + draft_n_total = n_total; + draft_n_accepted = n_accepted; +} + void Profiler::set_stop_reason(StopReason reason) { stop_reason = reason; } StopReason Profiler::get_stop_reason() const { return stop_reason; } @@ -66,6 +71,8 @@ void Profiler::to_profile_data(ProfileData& pd) { pd.prompt_tokens = prompt_tokens; pd.generated_tokens = generated_tokens; + pd.draft_n_total = draft_n_total; + pd.draft_n_accepted = draft_n_accepted; pd.stop_reason = stop_reason_to_string(stop_reason); } From 155f0ebde1f0fdb9e0054bc92e8fe6b737e7e2e7 Mon Sep 17 00:00:00 2001 From: Mengsheng Wu Date: Thu, 16 Jul 2026 14:35:26 +0800 Subject: [PATCH 2/2] test(sdk): add --draft-model/--draft-tokens to geniex-bench for MTP Signed-off-by: Mengsheng Wu --- sdk/benchmark/benchmark.c | 26 ++++++++++++++++++++++---- 1 file changed, 22 insertions(+), 4 deletions(-) diff --git a/sdk/benchmark/benchmark.c b/sdk/benchmark/benchmark.c index f66d398d5..cf273e921 100644 --- a/sdk/benchmark/benchmark.c +++ b/sdk/benchmark/benchmark.c @@ -108,6 +108,9 @@ typedef struct { int32_t n_threads; int32_t ngl_override; /* -1 = use resolved alias default; >=0 overrides */ + const char* draft_model; /* draft/MTP GGUF for speculative decoding (llama_cpp); NULL = disabled */ + int32_t draft_tokens; /* draft tokens per step (0 = plugin default) */ + const char* output_json; const char* output_md; const char* cell_id; @@ -217,6 +220,8 @@ static void usage(const char* argv0) { " -t, --threads N generation threads (0 = SDK default)\n" " -ngl, --n-gpu-layers N llama_cpp layers to offload; overrides the\n" " device alias default (-1 = all layers)\n" + " --draft-model PATH draft/MTP GGUF enabling speculative decoding (llama_cpp)\n" + " --draft-tokens N draft tokens per step (0 = plugin default)\n" " --warmup N default 1\n" " --no-warmup equivalent to --warmup 0\n" " --temperature F default 0.0\n" @@ -704,6 +709,10 @@ static void parse_args(int argc, char** argv, options_t* o) { o->n_threads = atoi(arg_value(argc, argv, &i, a)); } else if (strcmp(a, "-ngl") == 0 || strcmp(a, "--n-gpu-layers") == 0) { o->ngl_override = atoi(arg_value(argc, argv, &i, a)); + } else if (strcmp(a, "--draft-model") == 0) { + o->draft_model = arg_value(argc, argv, &i, a); + } else if (strcmp(a, "--draft-tokens") == 0) { + o->draft_tokens = atoi(arg_value(argc, argv, &i, a)); } else if (strcmp(a, "--output-json") == 0) { o->output_json = arg_value(argc, argv, &i, a); } else if (strcmp(a, "--output-md") == 0) { @@ -986,10 +995,12 @@ static void fill_gen_config(geniex_GenerationConfig* g, geniex_SamplerConfig* s, static void fill_model_config(geniex_ModelConfig* c, const options_t* o, int32_t ngl) { memset(c, 0, sizeof(*c)); - c->n_ctx = o->n_ctx; - c->n_threads = o->n_threads; - c->n_gpu_layers = ngl; - c->max_tokens = o->max_new_tokens; + c->n_ctx = o->n_ctx; + c->n_threads = o->n_threads; + c->n_gpu_layers = ngl; + c->max_tokens = o->max_new_tokens; + c->spec_draft_model = o->draft_model; /* may be NULL */ + c->spec_n_draft = o->draft_tokens; } static void run_llm(const options_t* o, const char* device_id, int32_t ngl, run_result_t* out) { @@ -1171,6 +1182,13 @@ static void run_llm(const options_t* o, const char* device_id, int32_t ngl, run_ if (!is_warmup && o->accuracy && gout.full_text) { print_gen_text(gout.full_text); } + if (!is_warmup && gout.profile_data.draft_n_total > 0) { + fprintf(stderr, + "[spec] draft acceptance = %.5f (%lld accepted / %lld generated)\n", + (double)gout.profile_data.draft_n_accepted / (double)gout.profile_data.draft_n_total, + (long long)gout.profile_data.draft_n_accepted, + (long long)gout.profile_data.draft_n_total); + } if (gout.full_text) { geniex_free(gout.full_text); }