Vulkan compute shaders for on-device AI on Android
Meta 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 productengineering
TL;DR
NNAPI is a comfortable abstraction — until it isn’t. For a quantized diffusion-based segmentation model running at real-time cadence on Pixel 9, the dispatch overhead and opaque memory management of NNAPI became the bottleneck. Dropping to raw Vulkan compute with carefully managed descriptor sets, push constants for dynamic tensor dims, and a ping-pong buffer schedule cut end-to-end inference latency and kept us safely under the Adreno 750’s usable VkMemoryHeap budget.
Why bypass NNAPI?
NNAPI was designed for portability across a heterogeneous device ecosystem. That portability has a cost: driver-level graph compilation, opaque buffer staging, and an inference runtime that doesn’t expose the scheduling primitives you need to overlap GPU work with CPU post-processing.
The failure mode is subtle. NNAPI doesn’t crash — it falls back silently to CPU execution when an op isn’t accelerated, and your profiling shows wall-clock latency but not why. On a quantized INT8 segmentation model with custom attention-like ops, you will hit that silent fallback more than you expect.
Here’s how the three paths compare on quantized models and Adreno-class GPUs:
| Execution path | Approximate 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 for that control in engineering complexity, which is the right trade-off when you’re targeting a single known GPU family.
Descriptor set management
Vulkan’s descriptor model is explicit by design. For a multi-stage segmentation pipeline — encoder, bottleneck, decoder — each shader stage needs its own VkDescriptorSetLayout. Pooling these correctly is the difference between a stable production build and a driver crash under memory pressure.
In practice, allocate a single VkDescriptorPool sized to your worst-case pipeline depth, and bind tensor buffers as VK_DESCRIPTOR_TYPE_STORAGE_BUFFER. Don’t re-create descriptor sets per frame. Maintain a small ring of pre-allocated sets and rotate them in sync with your command buffer submissions.
// Example: simple INT8 depthwise conv compute shader header
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[]; };Push constants for dynamic tensor dimensions
Quantized models often run at multiple resolutions depending on the input frame crop. Hard-coding tensor dimensions in shader specialization constants forces recompilation on resize. Push constants — limited to 128 bytes on most Adreno drivers — are the right way to pass dynamic width, height, stride, and quantization scale at dispatch time with zero overhead.
// 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)Ping-pong buffer scheduling to hide dispatch latency
A single dispatch-and-wait loop leaves the CPU idle during GPU execution. For a segmentation pipeline that also needs to composite the mask back into the camera preview, that idle time is wasted. The fix is ping-pong scheduling: 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.
The VkMemoryHeap ceiling on Adreno 750
The Pixel 9’s Adreno 750 surfaces multiple VkMemoryHeap entries. The device-local heap that maps to on-chip GPU memory is the fast path — but it’s bounded. Allocate your full model weight set, activations, and ping-pong buffers there simultaneously and you’ll breach the budget and trigger driver-level eviction or OOM.
The safe approach: weights are read-only after load, so place them in the device-local heap once. Activations and ping-pong buffers share a secondary allocation pool from a host-visible heap and rely on Adreno’s efficient cache hierarchy to minimize the penalty. Query vkGetPhysicalDeviceMemoryProperties at init time and gate your allocation strategy on the reported heap sizes — never assume a fixed budget.
Takeaways
Before committing to NNAPI, use the Android GPU Inspector to verify which ops actually dispatch to the accelerator. Silent CPU fallback on even one layer of a segmentation model can tank your frame budget.
Use push constants for any per-inference parameter that changes — quantization scales, input dimensions, crop offsets. It costs nothing at dispatch time and eliminates pipeline recompilation on resize.
Size your VkMemoryHeap allocation strategy to the device’s reported heap budget, not a guessed constant. Query at runtime, separate weights from activations, and use ping-pong buffers to keep the GPU pipeline full without doubling your peak memory footprint.