From 7b90944b61f81101193c9dc5374192ea83255ceb Mon Sep 17 00:00:00 2001
From: kami <kamixtrash@gmail.com>
Date: Sat, 18 Jul 2026 14:28:26 +0400
Subject: [PATCH] feat(inference): pass [[models]] params through to
 llama-server (enables draft-MTP) (#243)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

ModelConfig.params was loaded from config but never consumed — the spawn
command in DefaultModelManager was hardcoded. Thread it through: params (a
flag->value map) flattens to token order on ModelDescriptor.extraArgs and
appends to the llama-server argv. This engages Multi-Token Prediction
speculative decoding purely from config:

  [[models]]
  params = { "-md" = "/path/draft.gguf", "-ngld" = "99", "--spec-type" = "draft-mtp", "--spec-draft-n-max" = "4" }

Each entry stays a distinct argv token (no shell), so paths with spaces are safe.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
---
 .../inference/commons/ModelDescriptor.kt      |  3 +++
 .../llama/cpp/DefaultModelManager.kt          |  2 +-
 .../infrastructure/InfrastructureModule.kt    |  4 ++++
 .../InfrastructureModuleModelTest.kt          | 21 +++++++++++++++++++
 4 files changed, 29 insertions(+), 1 deletion(-)

diff --git a/infrastructure/inference/commons/src/main/kotlin/com/correx/infrastructure/inference/commons/ModelDescriptor.kt b/infrastructure/inference/commons/src/main/kotlin/com/correx/infrastructure/inference/commons/ModelDescriptor.kt
index f8ddf3b9..2af97a64 100644
--- a/infrastructure/inference/commons/src/main/kotlin/com/correx/infrastructure/inference/commons/ModelDescriptor.kt
+++ b/infrastructure/inference/commons/src/main/kotlin/com/correx/infrastructure/inference/commons/ModelDescriptor.kt
@@ -12,4 +12,7 @@ data class ModelDescriptor(
     val idleTimeoutMs: Long = 60_000L,
     val contextSize: Int = 24_576,
     val capabilities: Set<CapabilityScore> = emptySet(),
+    /** Extra llama-server CLI args (already flattened to token order), e.g. draft-MTP speculative
+     * decoding: `-md <draft.gguf> -ngld 99 --spec-type draft-mtp --spec-draft-n-max 4`. */
+    val extraArgs: List<String> = emptyList(),
 )
diff --git a/infrastructure/inference/llama_cpp/src/main/kotlin/com/correx/infrastructure/inference/llama/cpp/DefaultModelManager.kt b/infrastructure/inference/llama_cpp/src/main/kotlin/com/correx/infrastructure/inference/llama/cpp/DefaultModelManager.kt
index 4a31dbc5..8d75144f 100644
--- a/infrastructure/inference/llama_cpp/src/main/kotlin/com/correx/infrastructure/inference/llama/cpp/DefaultModelManager.kt
+++ b/infrastructure/inference/llama_cpp/src/main/kotlin/com/correx/infrastructure/inference/llama/cpp/DefaultModelManager.kt
@@ -66,7 +66,7 @@ class DefaultModelManager(
                 "--ctx-size", descriptor.contextSize.toString(),
                 "--host", host,
                 "--port", port.toString(),
-            )
+            ) + descriptor.extraArgs
             val lp = LlamaProcess(command, logFile)
             lp.start()
             lp
diff --git a/infrastructure/src/main/kotlin/com/correx/infrastructure/InfrastructureModule.kt b/infrastructure/src/main/kotlin/com/correx/infrastructure/InfrastructureModule.kt
index 2e22bd50..1ee130c0 100644
--- a/infrastructure/src/main/kotlin/com/correx/infrastructure/InfrastructureModule.kt
+++ b/infrastructure/src/main/kotlin/com/correx/infrastructure/InfrastructureModule.kt
@@ -170,6 +170,10 @@ object InfrastructureModule {
         residencyMode = ResidencyMode.PERSISTENT,
         contextSize = config.contextSize,
         capabilities = parseCapabilitiesFromConfig(config.capabilities),
+        // [[models]] `params` (a flag->value map) passes straight through to llama-server, e.g.
+        // draft-MTP speculative decoding. Flatten to token order; each stays a distinct argv entry
+        // (no shell), so paths with spaces are safe.
+        extraArgs = config.params.flatMap { (flag, value) -> listOf(flag, value) },
     )
 
     private fun parseCapabilitiesFromConfig(capsMap: Map<String, Double>): Set<CapabilityScore> {
diff --git a/infrastructure/src/test/kotlin/com/correx/infrastructure/InfrastructureModuleModelTest.kt b/infrastructure/src/test/kotlin/com/correx/infrastructure/InfrastructureModuleModelTest.kt
index 94a7b4ae..bf675889 100644
--- a/infrastructure/src/test/kotlin/com/correx/infrastructure/InfrastructureModuleModelTest.kt
+++ b/infrastructure/src/test/kotlin/com/correx/infrastructure/InfrastructureModuleModelTest.kt
@@ -42,6 +42,27 @@ class InfrastructureModuleModelTest {
         assertEquals(0.8, codingScore!!.score)
     }
 
+    @Test
+    fun `modelConfigToDescriptor flattens params into extraArgs for llama-server`() {
+        val config = ModelConfig(
+            id = "mtp-model",
+            modelPath = "/models/main.gguf",
+            params = linkedMapOf(
+                "-md" to "/models/draft.gguf",
+                "-ngld" to "99",
+                "--spec-type" to "draft-mtp",
+                "--spec-draft-n-max" to "4",
+            ),
+        )
+
+        val descriptor = InfrastructureModule.modelConfigToDescriptor(config)
+
+        assertEquals(
+            listOf("-md", "/models/draft.gguf", "-ngld", "99", "--spec-type", "draft-mtp", "--spec-draft-n-max", "4"),
+            descriptor.extraArgs,
+        )
+    }
+
     @Test
     fun `modelConfigToDescriptor handles empty capabilities`() {
         val config = ModelConfig(