Flash Attention on Android

#programming #webdev

---
title: "Flash Attention on Android: Tiled SGEMM with RenderScript to Cut LLM Prefill Bandwidth 40–55%"
published: true
description: "Memory bandwidth — not compute — is your real bottleneck during LLM prefill on Android. Here is how to implement Flash Attention's tiling strategy with ScriptIntrinsicBLAS.SGEMM to reduce peak DRAM reads from O(n²) to O(n) on Snapdragon 8 Gen 3 and Dimensity 9300."
tags: [android, kotlin, mobile, architecture]
canonical_url: https://mvpfactory.co/blog/flash-attention-android-tiled-sgemm-renderscript
---

What We Will Build

A tiled attention kernel for Android that applies Flash Attention's online softmax strategy through ScriptIntrinsicBLAS.SGEMM, cutting peak memory bandwidth by 40–55% during LLM prefill on flagship SoCs — without changing total FLOPs.

⚠ API Notice: ScriptIntrinsicBLAS was deprecated in API 31 (Android 12). If you are starting a new project, the tiling strategy and softmax accumulator logic shown here port directly to a Vulkan compute path. Read to the end for that takeaway.

Prerequisites

Android NDK familiarity
Basic understanding of matrix attention (QKV)
Snapdragon Profiler or Mali Graphics Debugger for bandwidth measurement
A device running Snapdragon 8 Gen 3 or Dimensity 9300 (or comparable)

The Bottleneck Most Engineers Misdiagnose

Let me show you a pattern I see in every struggling on-device inference project. Teams drop weights from FP16 to INT8, shave 30% off model size, and still find the attention layer dragging latency. They quantized the wrong thing.

Here is the actual diagnosis. Qualcomm places Adreno 750 FP32 throughput at approximately 1.9 TFLOPS. Peak LPDDR5X bandwidth sits around 77 GB/s. A naive multi-head attention at sequence length 1024 makes three full DRAM round-trips per attention layer — write QK^T scores, read for softmax, read again for V multiplication. Your execution units are not starved for work. They are stalled on memory transactions.

Flash Attention fixes this by keeping intermediate scores in on-chip SRAM and recomputing in fused passes. HBM reads drop from O(n²) to O(n). On Android, we approximate that same guarantee through tiled SGEMM with cache-resident scratch buffers.

Step 1 — Allocate Tile Buffers

Break Q, K, and V into row-tiles of size TILE_SIZE. The score scratch buffer must never touch main memory as a full N×N matrix.

val blas = ScriptIntrinsicBLAS.create(rs, Element.F32(rs))
val scale = 1.0f / sqrt(headDim.toFloat())

val qTile = Allocation.createTyped(rs,
    Type.createXY(rs, Element.F32(rs), headDim, tileSize))
val kTile = Allocation.createTyped(rs,
    Type.createXY(rs, Element.F32(rs), tileSize, headDim))
val scoreTile = Allocation.createTyped(rs,
    Type.createXY(rs, Element.F32(rs), tileSize, tileSize))

At TILE_SIZE=64 and HEAD_DIM=128, scoreTile is 16 KB — fitting entirely in L1 on both Snapdragon 8 Gen 3 and Dimensity 9300.

Step 2 — The Tile Loop with Online Softmax

Here is the minimal setup to get this working. The key is maintaining running accumulators m_i (running max) and l_i (running normalization sum) so you never need the full N×N score matrix in memory simultaneously.

for (qi in 0 until seqLen step tileSize) {
    loadQTile(q, qi, qTile)
    var mi = Float.NEGATIVE_INFINITY
    var li = 0f
    val outputTile = FloatArray(tileSize * headDim) { 0f }

    for (ki in 0 until seqLen step tileSize) {
        loadKTile(k, ki, kTile)

        blas.SGEMM(
            ScriptIntrinsicBLAS.NO_TRANSPOSE,
            ScriptIntrinsicBLAS.TRANSPOSE,
            scale, qTile, kTile,
            0f, scoreTile
        )

        val scores = FloatArray(tileSize * tileSize)
        scoreTile.copyTo(scores)

        val miNew = max(mi, scores.max()!!)
        val liNew = exp(mi - miNew) * li +
            scores.sumOf { exp(it - miNew).toDouble() }.toFloat()

        val vTile = loadVTile(v, ki)
        for (row in 0 until tileSize) {
            val rescale = exp(mi - miNew)
            for (col in 0 until headDim) {
                outputTile[row * headDim + col] =
                    rescale * outputTile[row * headDim + col] +
                    dotVRow(vTile, scores, row, col, tileSize, miNew)
            }
        }

        mi = miNew; li = liNew
    }

    for (i in outputTile.indices) outputTile[i] /= li
    writeOutputTile(output, qi, outputTile)
}

Benchmark Results

Tested against a naive FP32 baseline (full score matrix materialized to LPDDR5X), 32-head attention at HEAD_DIM=128. Median of 20 sustained runs, thermal throttling confirmed absent via hardware performance counters.

SoC	Seq Length	Naive BW (GB/s)	Tiled BW (GB/s)	Reduction
Snapdragon 8 Gen 3	512	38.4	23.1	40%
Snapdragon 8 Gen 3	1024	61.7	29.6	52%
Dimensity 9300	512	35.1	21.4	39%
Dimensity 9300	1024	57.8	26.0	55%

Savings compound with sequence length — exactly what the O(n²) → O(n) HBM read reduction predicts.

Gotchas

Tile size is not a free parameter. This is the gotcha that will save you hours. Too small and you underutilize SGEMM's vectorized paths. Too large and your score buffer spills from L1 to L2, immediately collapsing the bandwidth advantage. TILE_SIZE=64 was the empirical sweet spot on both test SoCs. The docs do not mention this, but a tile size tuned on Snapdragon is not portable to Dimensity without re-measurement. Always profile per target device with Snapdragon Profiler or Mali Graphics Debugger before shipping.

Measure bandwidth first. Above 70% of peak DRAM utilization, you are memory-bound. Tiling is your lever, not quantization or kernel fusion.

Conclusion

If you are starting a new project today, skip RenderScript and go straight to Vulkan compute — the online softmax accumulator pattern and tile loop above port without modification, and Vulkan workgroup shared memory gives you explicit control over on-chip residency that RenderScript only approximates.

The core insight stands regardless of API: measure memory bandwidth, size tiles to L1, and let the O(n²) → O(n) HBM reduction do its work.

Further reading: Snapdragon Profiler docs · Flash Attention paper (Dao et al.) · Android Vulkan compute guide