Wiring Android's Vulkan Compute to a Quantized On-Device Diffusion Model for Real-Time Image Segmentation

A developer published a pattern for wiring raw Vulkan compute shaders directly to a quantized INT8 diffusion-based segmentation model on Android, targeting the Pixel 9's Adreno 750 GPU. The approach bypasses NNAPI—which silently falls back to CPU on unsupported ops—in favor of explicit descriptor set rings, push constants for dynamic tensor dimensions, and ping-pong buffer scheduling to overlap GPU work with CPU post-processing while staying within the device's VkMemoryHeap budget.

--- title: "Vulkan Compute Shaders for On-Device AI on Android" published: true description: "Wire Vulkan compute to a quantized diffusion model on Android. Bypass NNAPI, manage descriptor sets, and stay within Pixel 9's memory ceiling on Adreno 750." tags: android, mobile, architecture, performance canonical url: https://mvpfactory.co/blog/vulkan-compute-android-on-device-ai --- What we are building Let me show you a pattern I use when NNAPI stops being your friend. We are going to wire raw Vulkan compute directly to a quantized INT8 segmentation model running on Pixel 9 — covering descriptor set management, push constants for dynamic tensor dimensions, ping-pong buffer scheduling, and the VkMemoryHeap budget that keeps you from hitting OOM on Adreno 750. Prerequisites - NDK r26+ and Vulkan headers in your Android project - A quantized INT8 model with custom attention-like ops where NNAPI silent fallback hurts most - Familiarity with VkCommandBuffer and VkQueue basics - Android GPU Inspector installed for verification --- Why not NNAPI? NNAPI was designed for portability across a heterogeneous device ecosystem. That portability costs you: driver-level graph compilation, opaque buffer staging, and an inference runtime that hides scheduling primitives you need to overlap GPU work with CPU post-processing. Here is the gotcha that will save you hours: NNAPI does not crash on unsupported ops — it falls back silently to CPU execution. Your profiling shows wall-clock latency but not why . Here is how the three paths compare on quantized models and Adreno-class GPUs: | Execution path | Dispatch overhead | Memory visibility | Op coverage | |---|---|---|---| | NNAPI HAL delegated | High — compiled graph, opaque staging | None | Partial — silent CPU fallback | | TFLite GPU Delegate | Medium — OpenCL/Vulkan backend | Limited | Good for standard ops | | Raw Vulkan Compute | Minimal — direct VkQueue submit | Full | Complete — you write every shader | Raw Vulkan wins on control. You pay in engineering complexity. That is the right trade-off when targeting a single known GPU family. --- Step 1 — Descriptor set management Here is the minimal setup to get this working. Allocate a single VkDescriptorPool sized to your worst-case pipeline depth. Bind tensor buffers as VK DESCRIPTOR TYPE STORAGE BUFFER . Do not re-create descriptor sets per frame — maintain a small ring of pre-allocated sets and rotate them in sync with your command buffer submissions. glsl // INT8 depthwise conv compute shader bindings layout set = 0, binding = 0 readonly buffer InputTensor { int8 t input ; }; layout set = 0, binding = 1 readonly buffer Weights { int8 t weights ; }; layout set = 0, binding = 2 writeonly buffer Output { int32 t output ; }; For a multi-stage pipeline — encoder, bottleneck, decoder — each shader stage gets its own VkDescriptorSetLayout . Pooling these correctly is the difference between a stable production build and a driver crash under memory pressure. --- Step 2 — Push constants for dynamic tensor dimensions The docs do not mention this clearly, but hard-coding tensor dimensions in shader specialization constants forces recompilation on resize. Push constants — limited to 128 bytes on most Adreno drivers — are how you pass dynamic width, height, stride, and quantization scale at dispatch time with zero overhead. kotlin // Kotlin-side push constant upload before vkCmdDispatch val pushData = ByteBuffer.allocateDirect 16 .order ByteOrder.nativeOrder .apply { putInt inputWidth ; putInt inputHeight putFloat quantScale ; putFloat quantZeroPoint } vkCmdPushConstants cmdBuffer, pipelineLayout, VK SHADER STAGE COMPUTE BIT, 0, 16, pushData --- Step 3 — Ping-pong buffer scheduling A single dispatch-and-wait loop leaves the CPU idle during GPU execution. The fix: maintain two complete sets of input/output buffers. Submit frame N to the GPU, then let the CPU process the mask output from frame N-1 while the GPU works. This requires explicit VkSemaphore synchronization between submissions — one semaphore signaled when the GPU finishes writing output, consumed by the CPU-side mask compositor. The scheduling complexity is real, but latency-to-display drops because you eliminate the synchronous stall on every frame. --- Step 4 — Respect the VkMemoryHeap ceiling on Adreno 750 The Pixel 9's Adreno 750 surfaces multiple VkMemoryHeap entries. The device-local heap is the fast path — but it is bounded. Allocate your full model weight set, activations, and ping-pong buffers there simultaneously and you will breach the budget and trigger driver-level eviction or OOM. The safe approach: place read-only weights in the device-local heap once. Activations and ping-pong buffers go into a host-visible heap allocation pool. Query vkGetPhysicalDeviceMemoryProperties at init time and gate your allocation strategy on the reported heap sizes — never assume a fixed budget. --- Gotchas Silent NNAPI fallback is invisible in wall-clock profiling. Use Android GPU Inspector to verify which ops actually dispatch to the accelerator before committing to NNAPI for any custom op layer. Descriptor set recreation per frame kills performance. Pre-allocate your ring and rotate. Specialization constants on resize trigger recompilation. Push constants exist for exactly this reason — use them for any per-inference parameter that changes. Guessing heap budget constants is always wrong. Query at runtime, every time, on every device. --- Conclusion NNAPI is a comfortable abstraction until it silently degrades on a custom op. Dropping to raw Vulkan compute with managed descriptor sets, push constants for dynamic tensor dims, and ping-pong buffer scheduling cut end-to-end inference latency on this pipeline and kept memory safely under the Adreno 750's usable VkMemoryHeap budget. The engineering overhead is real — but for a known GPU family at production cadence, it is the right call. Further reading: - Vulkan Compute — Khronos spec https://registry.khronos.org/vulkan/ - Android GPU Inspector https://developer.android.com/agi - VkMemoryHeap docs https://registry.khronos.org/vulkan/specs/latest/man/html/VkMemoryHeap.html