QVAC-20987 feat[api]: add llm reasoning budget caps

.github/workflows/cpp-tests-llm.yml

@@ -37,7 +37,7 @@ jobs:
             platform: linux
             runner: qvac-ubuntu2404-x64
             arch: x64
-          - os: macos-14
+          - os: macos-15-xlarge
             platform: darwin
             arch: arm64
           - os: windows-2025

packages/llm-llamacpp/CMakeLists.txt

@@ -172,7 +172,7 @@ if(BUILD_CLI)
   target_include_directories(cli_tool PRIVATE ${QVAC_LIB_INFERENCE_ADDON_CPP_INCLUDE_DIRS})
   # Set C++20 standard for std::views support
   target_compile_features(cli_tool PRIVATE cxx_std_20)
-  
+
   # Install ggml backends next to the CLI tool binary in build folder
   if((ANDROID OR UNIX) AND NOT APPLE)
     foreach(_backend ${GGML_AVAILABLE_BACKENDS})
@@ -190,7 +190,7 @@ if(BUILD_TESTING)
   find_package(GTest CONFIG REQUIRED)
   include(GoogleTest)
   enable_testing()
-  
+
   # Integration tests for model classes (includes backend selection tests)
   # Pass ENABLE_COVERAGE option to test subdirectory
   add_subdirectory(test/unit)

packages/llm-llamacpp/addon/src/addon/AddonJs.hpp

@@ -1,6 +1,8 @@
 #pragma once
 #include <any>
+#include <cmath>
 #include <functional>
+#include <limits>
 #include <memory>
 #include <string>
 #include <type_traits>
@@ -371,15 +373,18 @@ inline void parseGenerationParams(
   auto reasoningBudget = configObj->getOptionalPropertyAs<js::Number, double>(
       env, "reasoning_budget");
   if (reasoningBudget.has_value()) {
-    // Exact compare (0 and -1 are representable doubles) so fractional
-    // inputs like 0.5 are rejected, not silently truncated.
-    if (*reasoningBudget != 0 && *reasoningBudget != -1) {
+    // Reject fractional inputs (0.5, -1.1, 32.7) by requiring the value to
+    // round-trip through int. -1 = unrestricted, 0 = disabled, N>0 caps the
+    // reasoning channel at N tokens via the budget sampler.
+    const double value = *reasoningBudget;
+    if (value < -1 || value != std::floor(value) ||
+        value > static_cast<double>(std::numeric_limits<int>::max())) {
       throw StatusError(
           general_error::InvalidArgument,
-          "generationParams.reasoning_budget must be -1 (unrestricted) "
-          "or 0 (disabled)");
+          "generationParams.reasoning_budget must be -1 (unrestricted), "
+          "0 (disabled), or a positive integer (token cap)");
     }
-    overrides.reasoning_budget = static_cast<int>(*reasoningBudget);
+    overrides.reasoning_budget = static_cast<int>(value);
   }
 }
 

packages/llm-llamacpp/addon/src/model-interface/GenerationParamsApply.cpp

@@ -29,6 +29,21 @@ void applyGenerationOverridesToSampling(
   setIf(overrides.presence_penalty, sampling.penalty_present);
   setIf(overrides.repeat_penalty, sampling.penalty_repeat);
 
+  // Forward reasoning_budget into the budget-sampler's token cap. The
+  // template-specific start/end/forced vectors are refreshed after prompt
+  // formatting, where common_chat_params exposes the thinking tags.
+  if (overrides.reasoning_budget) {
+    const int budget = *overrides.reasoning_budget;
+    if (budget > 0) {
+      sampling.reasoning_budget_tokens = budget;
+    } else {
+      sampling.reasoning_budget_tokens = -1;
+      sampling.reasoning_budget_start.clear();
+      sampling.reasoning_budget_end.clear();
+      sampling.reasoning_budget_forced.clear();
+    }
+  }
+
   // `json_schema` and `grammar` are mutually exclusive at the JS boundary
   // and in `AddonJs::runJob::parseText`, so reaching this branch with both
   // set means a caller bypassed those checks (most likely the C++ unit

packages/llm-llamacpp/addon/src/model-interface/LlamaModel.cpp

@@ -98,7 +98,8 @@ void LlamaModel::resolveShardPaths(
 void LlamaModel::tuneConfigMap(
     std::unordered_map<std::string, std::string>& configFilemap,
     const ModelMetaData& metadata, const std::optional<int>& adrenoVersion,
-    const FinetuneConfigOverrides& finetuneOverrides, bool isOpenCl) {
+    const FinetuneConfigOverrides& finetuneOverrides, bool isOpenCl,
+    bool isMetal) {
 
   const bool isFinetuning = finetuneOverrides.active;
 
@@ -143,13 +144,11 @@ void LlamaModel::tuneConfigMap(
     QLOG_IF(
         Priority::INFO,
         "[LlamaModel] BitNet model detected: disabling flash attention\n");
-  } else if (isOpenCl && notUserSet("flash-attn", "flash_attn")) {
+  } else if (notUserSet("flash-attn", "flash_attn")) {
     configFilemap.erase("flash_attn");
-    configFilemap["flash-attn"] = "off";
+    configFilemap["flash-attn"] = "on";
     QLOG_IF(
-        Priority::INFO,
-        "[LlamaModel] OpenCL backend selected: disabling flash attention by "
-        "default (not reliably supported on OpenCL)\n");
+        Priority::INFO, "[LlamaModel] Enabling flash attention by default\n");
   }
 
   constexpr int kAdrenoUbatchThreshold = 800;
@@ -213,22 +212,18 @@ void LlamaModel::tuneConfigMap(
     }
   }
 
-  // TurboQuant / PolarQuant KV-cache types (tbq3_0 / tbq4_0 / pq3_0 / pq4_0,
-  // ship Vulkan + CPU implementations only. On the OpenCL backend
-  // the kernels don't exist; on Metal the standalone MUL_MAT pipelines are
-  // explicitly disabled. Silently letting the user proceed lets llama.cpp
-  // commit KV-cache tensors to a backend that can't run the ops on them,
-  // which then aborts in ggml_backend_sched_split_graph at model load.
-  // Surface a clean error here instead.
-#if defined(__APPLE__)
-  constexpr bool kIsMetal = true;
-#else
-  constexpr bool kIsMetal = false;
-#endif
-  if (isOpenCl || kIsMetal) {
+  // Quantized KV-cache types are fragile on OpenCL: standard q-cache types can
+  // fail later during cache shifts, while TBQ/PQ kernels are not implemented.
+  // Surface a clean error here instead of letting llama.cpp commit KV-cache
+  // tensors to a backend that can't run the required ops.
+  if (isOpenCl || isMetal) {
     auto isTurboQuantKvType = [](const std::string& v) {
       return v == "tbq3_0" || v == "tbq4_0" || v == "pq3_0" || v == "pq4_0";
     };
+    auto isQuantizedKvType = [&](const std::string& v) {
+      return isTurboQuantKvType(v) || v == "q4_0" || v == "q4_1" ||
+             v == "q5_0" || v == "q5_1" || v == "q8_0" || v == "iq4_nl";
+    };
     auto checkCacheType = [&](const char* hyphenKey,
                               const char* underscoreKey,
                               const char* side) {
@@ -237,20 +232,30 @@ void LlamaModel::tuneConfigMap(
         it = configFilemap.find(underscoreKey);
       if (it == configFilemap.end())
         return;
-      if (!isTurboQuantKvType(it->second))
+      if (isOpenCl) {
+        if (!isQuantizedKvType(it->second))
+          return;
+      } else if (!isTurboQuantKvType(it->second)) {
         return;
+      }
       const char* backendName = isOpenCl ? "OpenCL" : "Metal";
+      const char* typeName = isTurboQuantKvType(it->second)
+                                 ? "TurboQuant/PolarQuant"
+                                 : "quantized";
+      const char* alternatives =
+          isOpenCl ? "f32/f16/bf16"
+                   : "f32/f16/bf16/q4_0/q4_1/q5_0/q5_1/q8_0/iq4_nl";
       throw qvac_errors::StatusError(
           qvac_errors::general_error::InvalidArgument,
           string_format(
-              "[LlamaModel] cache-type-%s=%s is a TurboQuant/PolarQuant type "
-              "and is not supported on the %s backend (Vulkan and CPU only). "
-              "Either pick a different cache type (f32/f16/bf16/q4_0/q4_1/"
-              "q5_0/q5_1/q8_0/iq4_nl) or switch device to a Vulkan GPU or "
-              "CPU.\n",
+              "[LlamaModel] cache-type-%s=%s is a %s KV-cache type and is not "
+              "supported on the %s backend. Either pick a different cache "
+              "type (%s) or switch device to a Vulkan GPU or CPU.\n",
               side,
               it->second.c_str(),
-              backendName));
+              typeName,
+              backendName,
+              alternatives));
     };
     checkCacheType("cache-type-k", "cache_type_k", "k");
     checkCacheType("cache-type-v", "cache_type_v", "v");
@@ -994,21 +999,23 @@ void LlamaModel::commonParamsParse(
         "embedded chat template is applied\n");
   }
 
-  // reasoning-budget controls whether the model emits a <think> reasoning
-  // channel. -1 (default) leaves it on; 0 disables. `std::from_chars` is used
-  // instead of `std::stoi` because the latter accepts trailing garbage ("0abc"
-  // → 0) and throws an uncaught `std::out_of_range` on overflow.
+  // reasoning-budget controls the size of the model's <think> reasoning
+  // channel: -1 (default) leaves it unrestricted, 0 disables thinking
+  // entirely, any positive N caps the reasoning channel at N tokens (the
+  // budget sampler forces </think> once N reasoning tokens have been
+  // emitted).
   auto parseReasoningBudget = [](const std::string& raw) {
     int value = 0;
     const char* begin = raw.data();
     const char* end = begin + raw.size();
     const auto [ptr, ec] = std::from_chars(begin, end, value);
-    if (ec != std::errc{} || ptr != end || (value != 0 && value != -1)) {
+    if (ec != std::errc{} || ptr != end || value < -1) {
       throw qvac_errors::StatusError(
           ADDON_ID,
           qvac_errors::general_error::toString(
               qvac_errors::general_error::InvalidArgument),
-          "reasoning-budget must be -1 (unrestricted) or 0 (disabled)");
+          "reasoning-budget must be -1 (unrestricted), 0 (disabled), or a "
+          "positive integer (token cap)");
     }
     return value;
   };
@@ -1114,6 +1121,7 @@ void LlamaModel::commonParamsParse(
   }
 
   bool isOpenCl = false;
+  bool isMetal = false;
   {
     using namespace backend_selection;
     const BackendType preferredBackend =
@@ -1180,14 +1188,18 @@ void LlamaModel::commonParamsParse(
 
     isOpenCl = chosenBackend.first == BackendType::GPU &&
                chosenBackend.second.find("opencl") != std::string::npos;
+    isMetal = chosenBackend.first == BackendType::GPU &&
+              (chosenBackend.second.find("metal") != std::string::npos ||
+               chosenBackend.second.rfind("mtl", 0) == 0);
   }
 
   tuneConfigMap(
       configFilemap,
       metadata_,
       outAdrenoVersion,
       pendingFinetuneOverrides_,
-      isOpenCl);
+      isOpenCl,
+      isMetal);
 
   // Handle both reverse-prompt variants
   for (const std::string& key : {"reverse-prompt", "reverse_prompt"}) {
@@ -1310,10 +1322,6 @@ void LlamaModel::commonParamsParse(
   postprocess_cpu_params(params.cpuparams, nullptr);
   postprocess_cpu_params(params.cpuparams_batch, &params.cpuparams);
 
-  postprocess_cpu_params(params.speculative.draft.cpuparams, &params.cpuparams);
-  postprocess_cpu_params(
-      params.speculative.draft.cpuparams_batch, &params.cpuparams_batch);
-
   if (!params.kv_overrides.empty()) {
     params.kv_overrides.emplace_back();
     params.kv_overrides.back().key[0] = 0;

packages/llm-llamacpp/addon/src/model-interface/LlamaModel.hpp

@@ -67,14 +67,15 @@ class LlamaModel : public IModel, public IModelAsyncLoad, public IModelCancel {
   /// @param adrenoVersion Detected Adreno GPU version, if any.
   /// @param finetuneOverrides If set, finetuning mode is active with these
   /// context/batch params and GPU caps.
-  /// @param isOpenCl True when the chosen GPU backend is OpenCL; used to
-  /// disable flash-attn by default since it is not reliably supported on
-  /// the OpenCL backend.
+  /// @param isOpenCl True when the chosen GPU backend is OpenCL; used to reject
+  /// unsupported quantized KV-cache types.
+  /// @param isMetal True when the chosen GPU backend is Metal; used to reject
+  /// unsupported TurboQuant/PolarQuant KV-cache types.
   static void tuneConfigMap(
       std::unordered_map<std::string, std::string>& configFilemap,
       const ModelMetaData& metadata, const std::optional<int>& adrenoVersion,
       const FinetuneConfigOverrides& finetuneOverrides = {},
-      bool isOpenCl = false);
+      bool isOpenCl = false, bool isMetal = false);
 
   /**
    * The Constructor for llama model.

packages/llm-llamacpp/addon/src/model-interface/LlmContext.hpp

@@ -34,10 +34,11 @@ struct GenerationParams {
   // Mutually exclusive with `grammar` — the JS wrapper rejects requests
   // that set both. Mirrors the load-time `--json-schema` flag.
   std::optional<std::string> json_schema;
-  // Reasoning channel budget override. `-1` keeps reasoning on, `0` disables
-  // it for this request. Mirrors the load-time `reasoning-budget` config; the
-  // override is applied to `params_.reasoning_budget` for the duration of the
-  // request and restored on completion.
+  // Reasoning channel budget override. `-1` keeps reasoning unrestricted, `0`
+  // disables it, and positive values cap the reasoning channel at that many
+  // tokens. Mirrors the load-time `reasoning-budget` config; the override is
+  // applied to `params_.reasoning_budget` for the duration of the request and
+  // restored on completion.
   std::optional<int> reasoning_budget;
 
   [[nodiscard]] bool hasOverrides() const {

packages/llm-llamacpp/addon/src/model-interface/MtmdLlmContext.cpp

@@ -216,14 +216,42 @@ void MtmdLlmContext::tokenizeChat(
   if (!tools.empty()) {
     inputs.tools = tools;
   }
-  formattedChat = getPrompt(tmpls_.get(), inputs, &thinkingForcedOpen_);
+  std::string thinkingStartTag;
+  std::string thinkingEndTag;
+  std::string generationPrompt;
+  formattedChat = getPrompt(
+      tmpls_.get(),
+      inputs,
+      &thinkingForcedOpen_,
+      &thinkingStartTag,
+      &thinkingEndTag,
+      &generationPrompt);
+  thinkingForcedOpenText_ =
+      thinkingForcedOpen_
+          ? getThinkingForcedOpenText(generationPrompt, thinkingStartTag)
+          : std::string{};
 
   if (formattedChat.empty()) {
     std::string errorMsg = string_format(
         "[MtmdLlm] %s: formatted chat prompt is empty\n", __func__);
     throw qvac_errors::StatusError(ADDON_ID, toString(EmptyPrompt), errorMsg);
   }
 
+  if (configureReasoningBudgetSampling(
+          params_,
+          modelCtx_.lctx,
+          thinkingStartTag,
+          thinkingEndTag,
+          generationPrompt)) {
+    smpl_.reset(common_sampler_init(modelCtx_.model, params_.sampling));
+    if (!smpl_) {
+      std::string errorMsg = string_format(
+          "[MtmdLlm] %s: failed to initialize sampling subsystem\n", __func__);
+      throw qvac_errors::StatusError(
+          ADDON_ID, toString(UnableToCreateSamplingSystem), errorMsg);
+    }
+  }
+
   QLOG_IF(
       Priority::DEBUG,
       string_format("[MtmdLlm] formatted prompt: %s\n", formattedChat.c_str()));
@@ -492,10 +520,10 @@ bool MtmdLlmContext::generateResponse(
   if (thinkingForcedOpen_ && outputCallback) {
     // MtmdLlmContext doesn't carry a reasoningState_ (no reasoning-aware EOS
     // replacement on the multimodal path today), so unlike TextLlmContext we
-    // only prepend the visible "<think>\n" opener and don't flip an
+    // only prepend the visible reasoning opener and don't flip an
     // inside_reasoning flag. If reasoning state is added here later, mirror
     // TextLlmContext::generateResponse and set it true alongside this emit.
-    outputCallback("<think>\n");
+    outputCallback(thinkingForcedOpenText_);
   }
 
   if (stopGeneration_.load()) {
@@ -717,6 +745,8 @@ void MtmdLlmContext::resetState(bool resetStats) {
 
   // Clear UTF-8 buffer when resetting state
   utf8Buffer_.clear();
+  thinkingForcedOpen_ = false;
+  thinkingForcedOpenText_.clear();
 
   clearSequenceMemory(modelCtx_.lctx);
 

packages/llm-llamacpp/addon/src/model-interface/MtmdLlmContext.hpp

@@ -230,10 +230,11 @@ class MtmdLlmContext : public LlmContext {
   bool isHarmonyModel_ = false;
   llama_token harmonyCallToken_ = LLAMA_TOKEN_NULL;
 
-  // Force-opens the reasoning channel in the prompt suffix to prepend the
-  // matching "<think>\n" opener to the visible stream so consumers see balanced
+  // Force-opens the reasoning channel in the prompt suffix. The text mirrors
+  // the template-specific visible reasoning opener so consumers see balanced
   // tags.
   bool thinkingForcedOpen_ = false;
+  std::string thinkingForcedOpenText_;
 
   std::atomic<bool> stopGeneration_ = false;
 };