Add Vulkan patches and update Docker build process

Dockerfile

@@ -50,10 +50,10 @@ ARG LLAMA_SERVER_VARIANT
 RUN groupadd --system modelrunner && useradd --system --gid modelrunner -G video --create-home --home-dir /home/modelrunner modelrunner
 # TODO: if the render group ever gets a fixed GID add modelrunner to it
 
-COPY scripts/apt-install.sh apt-install.sh
+COPY scripts/ /scripts/
 
 # Install ca-certificates for HTTPS and vulkan
-RUN ./apt-install.sh
+RUN /scripts/apt-install.sh && rm -rf /scripts
 
 WORKDIR /app
 

scripts/0001-Revert-venus-filter-out-venus-incapable-physical-dev.patch

@@ -0,0 +1,49 @@
+diff --git a/src/virtio/vulkan/vn_physical_device.c b/src/virtio/vulkan/vn_physical_device.c
+index 83477d2..00f5615 100644
+--- a/src/virtio/vulkan/vn_physical_device.c
++++ b/src/virtio/vulkan/vn_physical_device.c
+@@ -1500,12 +1500,6 @@ vn_physical_device_init_renderer_extensions(
+
+    vk_free(alloc, exts);
+
+-   /* VK_KHR_external_memory_fd is required for venus memory mapping */
+-   if (!physical_dev->renderer_extensions.KHR_external_memory_fd) {
+-      vk_free(alloc, physical_dev->extension_spec_versions);
+-      return VK_ERROR_INCOMPATIBLE_DRIVER;
+-   }
+-
+    return VK_SUCCESS;
+ }
+
+@@ -1631,6 +1625,10 @@ vn_physical_device_init(struct vn_physical_device *physical_dev)
+    const VkAllocationCallbacks *alloc = &instance->base.vk.alloc;
+    VkResult result;
+
++   result = vn_physical_device_init_renderer_extensions(physical_dev);
++   if (result != VK_SUCCESS)
++      return result;
++
+    vn_physical_device_init_external_memory(physical_dev);
+    vn_physical_device_init_external_fence_handles(physical_dev);
+    vn_physical_device_init_external_semaphore_handles(physical_dev);
+@@ -1662,6 +1660,7 @@ vn_physical_device_init(struct vn_physical_device *physical_dev)
+    return VK_SUCCESS;
+
+ fail:
++   vk_free(alloc, physical_dev->extension_spec_versions);
+    vk_free(alloc, physical_dev->queue_family_properties);
+    return result;
+ }
+@@ -1875,12 +1874,6 @@ filter_physical_devices(struct vn_physical_device *physical_devs,
+          continue;
+       }
+
+-      result = vn_physical_device_init_renderer_extensions(physical_dev);
+-      if (result != VK_SUCCESS) {
+-         vn_physical_device_base_fini(&physical_dev->base);
+-         continue;
+-      }
+-
+       if (supported_count < i)
+          physical_devs[supported_count] = *physical_dev;
+       supported_count++;

scripts/0001-virtio-vulkan-force-16k-alignment-for-allocations-HA.patch

@@ -0,0 +1,50 @@
+diff --git a/src/virtio/vulkan/vn_device_memory.c b/src/virtio/vulkan/vn_device_memory.c
+index 19b4a91..cb0e72c 100644
+--- a/src/virtio/vulkan/vn_device_memory.c
++++ b/src/virtio/vulkan/vn_device_memory.c
+@@ -31,6 +31,10 @@ vn_device_memory_alloc_simple(struct vn_device *dev,
+ {
+    VkDevice dev_handle = vn_device_to_handle(dev);
+    VkDeviceMemory mem_handle = vn_device_memory_to_handle(mem);
++
++   ((VkMemoryAllocateInfo *)alloc_info)->allocationSize =
++      align64(alloc_info->allocationSize, 16384);
++
+    if (VN_PERF(NO_ASYNC_MEM_ALLOC)) {
+       return vn_call_vkAllocateMemory(dev->primary_ring, dev_handle,
+                                       alloc_info, NULL, &mem_handle);
+diff --git a/src/virtio/vulkan/vn_renderer_virtgpu.c b/src/virtio/vulkan/vn_renderer_virtgpu.c
+index 5b199aa..4fe6ed8 100644
+--- a/src/virtio/vulkan/vn_renderer_virtgpu.c
++++ b/src/virtio/vulkan/vn_renderer_virtgpu.c
+@@ -1245,6 +1245,8 @@ virtgpu_bo_create_from_device_memory(
+    const uint32_t blob_flags =
+       virtgpu_bo_blob_flags(gpu, flags, external_handles);
+
++   size = align64(size, 16384);
++
+    uint32_t res_id;
+    uint32_t gem_handle = virtgpu_ioctl_resource_create_blob(
+       gpu, gpu->bo_blob_mem, blob_flags, size, mem_id, &res_id);
+@@ -1295,6 +1297,8 @@ virtgpu_shmem_create(struct vn_renderer *renderer, size_t size)
+ {
+    struct virtgpu *gpu = (struct virtgpu *)renderer;
+
++   size = align64(size, 16384);
++
+    struct vn_renderer_shmem *cached_shmem =
+       vn_renderer_shmem_cache_get(&gpu->shmem_cache, size);
+    if (cached_shmem) {
+diff --git a/src/virtio/vulkan/vn_ring.c b/src/virtio/vulkan/vn_ring.c
+index 2de4978..119a8f9 100644
+--- a/src/virtio/vulkan/vn_ring.c
++++ b/src/virtio/vulkan/vn_ring.c
+@@ -264,7 +264,7 @@ vn_ring_get_layout(size_t buf_size,
+    layout->extra_offset = layout->buffer_offset + layout->buffer_size;
+    layout->extra_size = extra_size;
+
+-   layout->shmem_size = layout->extra_offset + layout->extra_size;
++   layout->shmem_size = align64(layout->extra_offset + layout->extra_size, 16384);
+ }
+
+ struct vn_ring *

scripts/0002-virtio-vulkan-silence-stuck-in-wait-message-HACK.patch

@@ -0,0 +1,15 @@
+diff --git a/src/virtio/vulkan/vn_common.c b/src/virtio/vulkan/vn_common.c
+index 2d389a8..3020ad6 100644
+--- a/src/virtio/vulkan/vn_common.c
++++ b/src/virtio/vulkan/vn_common.c
+@@ -269,8 +269,8 @@ vn_relax(struct vn_relax_state *state)
+
+    if (unlikely(*iter % (1 << warn_order) == 0)) {
+       struct vn_instance *instance = state->instance;
+-      vn_log(instance, "stuck in %s wait with iter at %d", state->reason_str,
+-             *iter);
++      //vn_log(instance, "stuck in %s wait with iter at %d", state->reason_str,
++      //       *iter);
+
+       struct vn_ring *ring = instance->ring.ring;
+       const uint32_t status = vn_ring_load_status(ring);

scripts/apt-install.sh

@@ -1,12 +1,54 @@
 #!/bin/bash
 
+enable_source_repos() {
+  # DEB822 format (Ubuntu 24.04+)
+  for f in /etc/apt/sources.list.d/*.sources; do
+    [ -f "$f" ] && sed -i 's/^Types: deb$/Types: deb deb-src/' "$f"
+  done
+  # Traditional format: uncomment existing deb-src lines
+  [ -f /etc/apt/sources.list ] && sed -i '/^#\s*deb-src/s/^#\s*//' /etc/apt/sources.list
+}
+
+rebuild_and_install_mesa() {
+  enable_source_repos
+  apt-get update
+  apt-get install -y dpkg-dev
+  apt-get build-dep -y mesa
+
+  local build_dir
+  build_dir=$(mktemp -d)
+  pushd "$build_dir"
+
+  apt-get source mesa
+  cd mesa-*/
+
+  local script_dir
+  script_dir="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+  patch -p1 < "$script_dir/0001-Revert-venus-filter-out-venus-incapable-physical-dev.patch"
+  patch -p1 < "$script_dir/0001-virtio-vulkan-force-16k-alignment-for-allocations-HA.patch"
+  patch -p1 < "$script_dir/0002-virtio-vulkan-silence-stuck-in-wait-message-HACK.patch"
+
+  dpkg-buildpackage -us -uc -b
+
+  cd ..
+  dpkg -i mesa-vulkan-drivers_*.deb
+
+  popd
+  rm -rf "$build_dir"
+}
+
 main() {
   set -eux -o pipefail
 
   apt-get update
   local packages=("ca-certificates")
   if [ "$LLAMA_SERVER_VARIANT" = "generic" ] || [ "$LLAMA_SERVER_VARIANT" = "cpu" ]; then
-    packages+=("libvulkan1" "mesa-vulkan-drivers")
+    packages+=("libvulkan1")
+    if [ "$(uname -m)" = "aarch64" ]; then
+      rebuild_and_install_mesa
+    else
+      packages+=("mesa-vulkan-drivers")
+    fi
   fi
 
   apt-get install -y "${packages[@]}"