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

---
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)