Hello, I'm Seongjae Park from OptAI.
In the previous article, we implemented asynchronous softmax as the custom operator OptSoftmax and recovered accuracy after extensive debugging. Once accuracy was established, however, profiling revealed a more fundamental limitation in the custom-operator approach itself. This article explains that limitation and the decision to change where the optimization was applied.
1. The Structural Limitation: Custom Operators Disrupt HMX Allocation
The QNN compiler is designed to assign MatMul to HMX, the matrix unit. Creating a custom operator reduced HMX utilization because the operation was no longer on the compiler optimization path.
We tested two fusion boundaries. Initially, we wrapped the entire attention computation in a custom operator. Its MatMuls were absorbed into the custom region, preventing the compiler from assigning them to HMX and sending them to HVX instead. The conclusion was to leave MatMul untouched. Keeping MatMul in the graph and customizing only Softmax and mask addition recovered utilization, but it still fell short of the original model.
Both experiments pointed to the same conclusion: as long as we registered a custom operator, we could not match the original model HMX utilization. Leaving the standard-operator path recognized by the compiler reduced hardware acceleration efficiency and introduced unnecessary memory copies.
The fusion boundary for a custom operator is determined by execution-unit allocation, HMX versus HVX, rather than by the visual elegance of the graph. Even a carefully chosen boundary leaves a blind spot where compiler Graph Matching and Fusion cannot reach the custom region.
2. Changing Direction: Applying Async Softmax at Export Time
We changed the approach. Instead of registering a custom operator at runtime, we inserted asynchronous softmax directly into the attention implementation in PyTorch modeling.py, before ONNX export, and exported it as-is. The resulting graph still consists of standard operators, preserving compiler fusion and HMX allocation.
One point is easy to misunderstand. KV cache splitting here does not mean cutting up an ONNX graph after export. Async softmax must split chunks at the point where attention reads the KV cache. The cache must therefore be split inside attention in PyTorch modeling.py, with the tensors needed for chunk-wise partial sums and normalization added before export. Only then does the exported graph fully express this structure as standard operators, allowing the compiler to apply fusion and HMX allocation directly.
| Approach 1 · Custom Op registration | Approach 2 · Export-stage application (selected) | |
|---|---|---|
| Application point | Register OptSoftmax in the QNN graph | Insert async softmax into PyTorch modeling.py attention, then export to ONNX |
| HMX utilization | The compiler cannot recognize the custom region; MatMul leaves HMX, falling below the original | Keep standard operators → preserve fusion · 91 → 94 |
| Cost / remaining work | Hidden costs of kernel implementation, validation, and Genie registration requiring SDK internals | Remove custom registration and validation overhead · TPS still drops by about 10 |
3. Results: HMX Utilization and TPS Diverge
The results matched the design intent. Staying on the standard path raised HMX utilization from the original 91 to 94. Changing only the placement of the optimization exceeded a metric the custom operator had been unable to match. Higher utilization did not, however, translate directly into higher throughput. In the same configuration, TPS fell by about 10.
Measured profiles support this observation.
| Metric | Original model | Export-stage application |
|---|---|---|
| HMX utilization | 91.2 % | 94.9 % |
| HTP node count | 8,606 | 16,066 |
| Decode TPS | Baseline | Approximately −10 |
Notably, actual decode TPS fell even though single-graph profiler metrics improved: graph execute went from 5,652 → 4,836 µs, with 193.6 inf/s. Expressing KV cache splitting as standard operators nearly doubled the HTP node count from 8,606 → 16,066. We suspect the additional split, concat, and reduction operations introduced overhead in the decode loop repeated for every token.
This divergence is itself an important insight from the work. HMX utilization describes how busy the unit is, not how quickly tokens emerge. More important than recovering utilization is the remaining task of finding a point where both utilization and TPS improve. We are currently breaking down the causes of this TPS gap, with narrowing it as the next step.
4. What Genie Integration Actually Involves
Integration with the inference engine was another reason the custom-operator approach was costly. Genie does not officially support adding custom operators, so the op package must be registered directly in backend initialization code rather than in the graph definition. This involved modifying QnnApi.cpp to call backendRegisterOpPackage when the backend handle became available and inject the op package, .so. Doing so required substantial knowledge of Genie SDK internals.
This integration cost was another reason the export-stage approach was attractive. Expressing the optimization through standard operators allows the Genie decode loop to be reused unchanged, removing the maintenance burden of custom registration.
5. Conclusion: Where Matters as Much as What
The main lesson of this port was not the algorithm itself, but where it was inserted. The same asynchronous softmax disrupted compiler HMX allocation when registered as a custom operator, yet increased utilization when expressed through standard operators at export time.
Three points summarize the result.
- To preserve compiler-supported Graph Matching and Fusion, keep standard operators wherever possible. Before modifying the graph, first check whether SDK-supported operation combinations can express the change.
- Custom operators impose unexpected costs in both kernel implementation and validation, including quantization, HVX ordering, and KV cache accumulation, as well as Genie registration.
- HMX utilization and TPS do not necessarily move in the same direction. Recovering utilization should be the starting point for throughput improvement rather than the final measure of success. We are continuing to narrow the TPS gap while retaining the HMX utilization gains.
Thank you for reading!
The Hexagon execution unit specialized for matrix multiplication. It delivers high throughput on large GEMM operations. The graph must keep custom operators from absorbing MatMul so that this unit can be used effectively.
The unit responsible for vector operations, including element-wise softmax operations such as exp and sum. Data ordering depends on hardware behavior.
An optimization in which the compiler recognizes specific patterns of standard operators and fuses them into a single accelerated kernel. Moving off this path through a custom operator prevents that optimization.
References
- Hong, K., et al. (2023). FlashDecoding++: Faster Large Language Model Inference on GPUs. arXiv:2311.01282.
- Qualcomm. (2026). Qualcomm AI Engine Direct (QNN) SDK Documentation.