Use Vulkan

Dockerfile

@@ -1,4 +1,4 @@
-FROM nvidia/cuda:12.2.0-runtime-ubuntu22.04 AS base
+FROM nvidia/cuda:12.6.0-runtime-ubuntu24.04 AS base
 
 # Install base packages
 RUN apt-get update && \
@@ -11,7 +11,7 @@ RUN apt-get update && \
     git \
     gnupg \
     jq \
-    libcurl4-nss-dev \
+    libcurl4-gnutls-dev \
     libvulkan1 \
     libvulkan-dev \
     mesa-vulkan-drivers \
@@ -46,8 +46,6 @@ RUN mkdir -p /usr/local/lib/llama-cpp-python
 RUN tar -xf /tmp/vulkan-libs.tar -C /usr/local/lib/llama-cpp-python
 RUN cp -f /usr/local/lib/llama-cpp-python/*.so* /usr/local/lib/
 RUN rm /tmp/vulkan-libs.tar
-# Set Vulkan environment to use CPU fallback (LavaPipe)
-ENV VK_ICD_FILENAMES=/usr/share/vulkan/icd.d/lvp_icd.x86_64.json
 
 # Create directory structure
 RUN mkdir -p \

scripts/local.sh

@@ -21,6 +21,13 @@ docker build \
     --tag sia \
     .
 
+GPU_FLAGS=""
+# Check if AMD GPU devices are available
+if [ -e /dev/dri ] && [ -e /dev/kfd ]; then
+    echo "AMD GPU detected, enabling GPU acceleration"
+    GPU_FLAGS="--device=/dev/dri --device=/dev/kfd --group-add video"
+fi
+
 docker run \
     --init \
     --rm \
@@ -28,6 +35,7 @@ docker run \
     -p 8080:8080 \
     -p 8000:8000 \
     --env-file .env \
+    $GPU_FLAGS \
     -v /$(pwd)/models/:/root/models/ \
     -v /$(pwd)/iterations/:/root/data/iterations/ \
     -v /$(pwd)/tasks/:/root/data/tasks/ \

sia/config.py

@@ -4,7 +4,7 @@ from pathlib import Path
 from typing import Dict
 import argparse
 import os
-import tomli as tomllib
+import tomllib
 
 @dataclass
 class Config:

sia/web_agent.py

@@ -49,6 +49,7 @@ class WebAgent(BaseAgent):
 
         self._update_compiled_context()
         self._working_memory.add_change_handler(self._update_compiled_context)
+        self._parser._io_buffer.add_change_handler(lambda *_: self._update_compiled_context())
     
     @property
     def response_buffer(self) -> ResponseBuffer:

tools/gemma_infer/pyproject.toml

@@ -16,6 +16,7 @@ dependencies = [
     "sentencepiece>=0.2.0",
     "tiktoken>=0.9.0",
     "transformers>=4.0.0",
+    "vulkan",
     "xml_schema_validator @ file:///root/sia/lib/xml_schema_validator",
 ]
 

tools/gemma_infer/src/gemma_infer/gemma_llm_engine.py

@@ -23,11 +23,12 @@ class GemmaLlmEngine(LlmEngine):
         self._model = model
         self._tokenizer = AutoTokenizer.from_pretrained(tokenizer)
         self._token_limit = token_limit
+
         self._llm = Llama(
             model_path=model,
-            n_gpu_layers=100,
+            n_gpu_layers=0,
             n_ctx=token_limit,
-            #verbose=False,  # Disable most logging
+            flash_attn=True,
         )
 
     def infer_xml(self, schema: Path, system: str, context: str, prefix: str) -> Iterator[str]:
@@ -56,7 +57,9 @@ class GemmaLlmEngine(LlmEngine):
         yield from skip_prefix(content_generator(), prefix)
 
     def token_count(self, system: str, context: str) -> int:
-        return len(self._format_messages(system, context, None))
+        prompt = self._format_messages(system, context, None)
+        tokens = self._tokenizer.encode(prompt)
+        return len(tokens)
 
     def token_limit(self) -> int:
         return self._token_limit