[Engineering] Porting FlashDecoding++ (2): Custom Operators and Device Debugging

Hello, I'm Seongjae Park from OptAI.

To deploy the asynchronous softmax from the previous article on the Qualcomm Hexagon NPU, we started with the most straightforward approach: implementing attention directly as a Custom Op. This article documents that implementation and, more importantly, the device debugging that followed. The main lesson was that validating the custom kernel proved much more demanding than implementing it.

1. First Attempt: All of Attention in One Custom Operator

Our first step was to combine the entire attention sequence — MatMul → Add(Mask) → Softmax → MatMul — into a single custom operator. The expectation was simple and natural: keeping intermediate tensors inside the operator should reduce memory round trips.

DRAM reads fell substantially, as expected. Latency, however, became worse.

The cause was disrupted HMX allocation. To the QNN compiler, a custom operator is an opaque black box. Once the two attention MatMuls were placed inside it, the compiler could no longer map them to the HMX matrix accelerator and sent them both to HVX, the vector unit. HMX utilization plunged from 91% → 31%, while the work shifted to HVX drove the HTP node count from 8,606 → 30,961. Using fewer intermediate tensors reduced DRAM reads from 295M → 164M, but losing HMX execution for MatMul outweighed that benefit: throughput fell from 159.9 → 113.8 inf/s.

The lesson was clear: MatMul must stay on HMX. We therefore narrowed the fusion boundary. Both MatMuls would remain native QNN operations for HMX to execute, while only the memory-bound mask addition and softmax between them would be fused. The result was OptSoftmax, discussed next.

2. The OptSoftmax Custom Operator

The conventional attention structure is as follows.

scores = Q × Kᵀ
scores += mask
prob    = softmax(scores)
output  = prob × V

This structure keeps the entire [B, H, Q, K] tensor in memory for softmax. As KV length grows, memory bandwidth and VTCM usage increase, and the global dependency prevents streaming. We split softmax into chunks and restructured it around global normalization, translating the asynchronous idea from part 1 into a 2-pass approach. Mask addition and softmax were then fused into the custom operator OptSoftmax.

Before:  MatMul → Add(Mask) → Softmax → MatMul
After:  MatMul → OptSoftmax(Custom Op) → MatMul

*Before: the standard graph with separate MatMul, Add(Mask), Softmax, and MatMul nodes.*

*After: Add and Softmax are fused into a single OptSoftmax node.*

Inside the kernel, each block uses the precomputed unified maximum fdpp_phi instead of repeatedly searching for a maximum as in online softmax. Two VTCM buffers, row_buf[2], overlap computation in the current loop with loading for the next.

*Alternating between two buffers keeps loading and computation overlapped.*

3. Repeated Tokens: Where Debugging Began

The most persistent symptom during the initial port was output that repeated the same word indefinitely. To narrow down the cause, we established a diagnostic rule: if the first token is wrong, investigate encodings and the tokenizer; if the output breaks from the second token onward, investigate the KV cache, mask, RoPE, and attention. Our output was already wrong at the first token, so we investigated encodings and quantization first.

4. Validating the Logic on the CPU Backend First

We first validated the FP32 logic in the host environment using the CPU backend, libQnnCpu.so. The missing --input_layout option in qairt-converter had forced the LLM 4D tensors into a vision-oriented image format, shuffling the token order. Specifying NONTRIVIAL corrected this, after which the data count, mask padding, and argmax were all correct. The CPU code itself was sound; the problem was the HTP quantization path.

5. Quantization Was Changing the Meaning of the Mask

Log analysis showed that mask values were collapsing near -3000 rather than their intended values, preventing masking from working correctly. QNN HTP stores tensors as uint16 integers, but softmax exp, sum, and divide must be computed in floating point for accuracy. Correcting the dequant/quant path to account properly for the score and mask scale and offset restored masking.

uint16 input → dequantize → float → softmax → quantize → uint16 output
real = (quant - zero_point) × scale

6. Physical HVX Behavior Was Shuffling the Indices

After fixing the mask, probability values were nearly identical, yet the indices of the highest probabilities were still shifted by a few positions. The cause was the even/odd separation performed by Q6_Wuw_vzxt_Vuh, the intrinsic used to zero-extend uint16 to uint32. Expanding the 64 uint16 values in a 128-byte vector to uint32 produces 256 bytes, split across two vectors. Because of its shuffle-network structure, HVX places even and odd indices in separate vectors. This is the simplest operation for the hardware wiring, but it directly disrupted the ordering of softmax output.

*Zero-extension separates uint16 values into even and odd lanes, changing output order.*

7. Small Biases Accumulating Through the KV Cache

Even after restoring the order and matching argmax, inference was not fully stable. We compared the original and custom models prefill KV caches layer by layer. Layer 0 was correct, but values diverged in the middle layers and worsened across iterations. The cosine similarity of the early hidden state fell from 0.990 → 0.965 → 0.561 over repeated iterations. A small OptSoftmax bias was fed back through the contaminated KV cache and grew progressively.

*The early-layer hidden state drops to 0.561 at iter_2, measured layer by layer during prefill.*

When a custom operator produces correct-looking values but incorrect results, the cause is usually ordering, quantization, or execution-unit allocation. A scalar fallback path can also verify whether removing HVX-specific rearrangement is actually reflected in the deployed .so. The time spent on all of these checks, however, was itself the hidden cost of the custom-kernel approach.

References

1. Hong, K., et al. (2023). FlashDecoding++: Faster Large Language Model Inference on GPUs. arXiv:2311.01282.
2. Qualcomm. (2026). Qualcomm AI Engine Direct (QNN) SDK Documentation.