feat: Add support for overriding tensor buffer type

extensions/llamacpp-extension/src/index.ts

@@ -39,6 +39,7 @@ type LlamacppConfig = {
   auto_unload: boolean
   chat_template: string
   n_gpu_layers: number
+  override_tensor_buffer_t: string
   ctx_size: number
   threads: number
   threads_batch: number
@@ -1262,6 +1263,14 @@ export default class llamacpp_extension extends AIEngine {
     args.push('--jinja')
     args.push('--reasoning-format', 'none')
     args.push('-m', modelPath)
+    // For overriding tensor buffer type, useful where
+    // massive MOE models can be made faster by keeping attention on the GPU
+    // and offloading the expert FFNs to the CPU.
+    // This is an expert level settings and should only be used by people
+    // who knows what they are doing.
+    // Takes a regex with matching tensor name as input
+    if (cfg.override_tensor_buffer_t)
+      args.push('--override-tensor', cfg.override_tensor_buffer_t)
     args.push('-a', modelId)
     args.push('--port', String(port))
     if (modelConfig.mmproj_path) {
@@ -1340,8 +1349,8 @@ export default class llamacpp_extension extends AIEngine {
 
       return sInfo
     } catch (error) {
-      logger.error('Error loading llama-server:\n', error)
-      throw new Error(`Failed to load llama-server: ${error}`)
+      logger.error('Error in load command:\n', error)
+      throw new Error(`Failed to load model:\n${error}`)
     }
   }
 

web-app/src/containers/ModelSetting.tsx

@@ -22,60 +22,60 @@
 }
 
 export function ModelSetting({
   model,
   provider,
   smallIcon,
 }: ModelSettingProps) {
   const { updateProvider } = useModelProvider()
   const { t } = useTranslation()
 
   // Create a debounced version of stopModel that waits 500ms after the last call
   const debouncedStopModel = debounce((modelId: string) => {
     stopModel(modelId)
   }, 500)
 
   const handleSettingChange = (
     key: string,
     value: string | boolean | number
   ) => {
     if (!provider) return
 
     // Create a copy of the model with updated settings
     const updatedModel = {
       ...model,
       settings: {
         ...model.settings,
         [key]: {
           ...(model.settings?.[key] != null ? model.settings?.[key] : {}),
           controller_props: {
             ...(model.settings?.[key]?.controller_props ?? {}),
             value: value,
           },
         },
       },
     }
 
     // Find the model index in the provider's models array
     const modelIndex = provider.models.findIndex((m) => m.id === model.id)
 
     if (modelIndex !== -1) {
       // Create a copy of the provider's models array
       const updatedModels = [...provider.models]
 
       // Update the specific model in the array
       updatedModels[modelIndex] = updatedModel as Model
 
       // Update the provider with the new models array
       updateProvider(provider.provider, {
         models: updatedModels,
       })
 
       // Call debounced stopModel only when updating ctx_len or ngl
       if (key === 'ctx_len' || key === 'ngl' || key === 'chat_template') {
         debouncedStopModel(model.id)
       }
     }
   }
 
   return (
     <Sheet>
@@ -106,8 +106,10 @@
               <div key={key} className="space-y-2">
                 <div
                   className={cn(
-                    'flex items-start justify-between gap-8',
-                    key === 'chat_template' && 'flex-col gap-1'
+                    'flex items-start justify-between gap-8 last:mb-2',
+                    (key === 'chat_template' ||
+                      key === 'override_tensor_buffer_t') &&
+                      'flex-col gap-1 w-full'
                   )}
                 >
                   <div className="space-y-1 mb-2">

web-app/src/hooks/useModelProvider.ts

@@ -276,9 +276,34 @@ export const useModelProvider = create<ModelProviderState>()(
           })
         }
 
+        // Migration for override_tensor_buffer_type key (version 2 -> 3)
+        if (version === 2 && state?.providers) {
+          state.providers.forEach((provider) => {
+            if (provider.models) {
+              provider.models.forEach((model) => {
+                // Initialize settings if it doesn't exist
+                if (!model.settings) {
+                  model.settings = {}
+                }
+
+                // Add missing override_tensor_buffer_type setting if it doesn't exist
+                if (!model.settings.override_tensor_buffer_t) {
+                  model.settings.override_tensor_buffer_t = {
+                    ...modelSettings.override_tensor_buffer_t,
+                    controller_props: {
+                      ...modelSettings.override_tensor_buffer_t
+                        .controller_props,
+                    },
+                  }
+                }
+              })
+            }
+          })
+        }
+
         return state
       },
-      version: 2,
+      version: 3,
     }
   )
 )

web-app/src/lib/predefined.ts

@@ -133,4 +133,15 @@ export const modelSettings = {
       textAlign: 'right',
     },
   },
+  override_tensor_buffer_t: {
+    key: 'override_tensor_buffer_t',
+    title: 'Override Tensor Buffer Type',
+    description: 'Override the tensor buffer type for the model',
+    controller_type: 'input',
+    controller_props: {
+      value: '',
+      placeholder: 'e.g., layers\\.\\d+\\.ffn_.*=CPU',
+      type: 'text',
+    },
+  },
 }