High inference time for Softmax decomposition for Ethos-U backend

[eladc4]

🐛 Describe the bug

During model quantiation for Ethos-U, Softmax ops are decomposed to atomic ops and a maximum substitution is added before the Softmax to improve activation quantization range. The maximum substitution uses aten::amax which is very slow on the NPU.

Quantizing a toy model:
python softmax_script.py

Results show a very large inference time:

Batch Inference time 1057.09 ms

with aten::amax using 85% of the NPU cycles with very low Util% (MAC) value.

Note: the tensor sizes in the script are taken from the whisper-tiny encoder model.

softmax_script.py

Versions

PyTorch version: 2.9.0
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: macOS 26.3 (arm64)
GCC version: Could not collect
Clang version: 17.0.0 (clang-1700.6.4.2)
CMake version: version 4.2.1
Libc version: N/A

Python version: 3.12.10 (v3.12.10:0cc81280367, Apr 8 2025, 08:46:59) [Clang 13.0.0 (clang-1300.0.29.30)] (64-bit runtime)
Python platform: macOS-26.3-arm64-arm-64bit
Is CUDA available: False
CUDA runtime version: No CUDA
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
Is XPU available: False
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True
Caching allocator config: N/A

CPU:
Apple M4 Max

Versions of relevant libraries:
[pip3] executorch==1.0.0
[pip3] numpy==2.3.4
[pip3] onnx==1.20.0
[pip3] onnx-ir==0.1.14
[pip3] onnxruntime==1.23.2
[pip3] onnxscript==0.5.7
[pip3] torch==2.9.0
[pip3] torchao==0.14.0
[pip3] torchaudio==2.9.0
[pip3] torchcodec==0.9.1
[pip3] torchmetrics==1.8.2
[pip3] torchvision==0.24.0
[conda] Could not collect

cc @digantdesai @freddan80 @per @zingo @oscarandersson8218 @mansnils @Sebastian-Larsson @robell

[gggekov]

Hi @eladc4 ,
Thanks for raising that. The amax operator will indeed use the elementwise engine of the NPU rather than the MACs.

That being said, I'd like to understand the impact on performance of softmax relative to a full model. How did you conclude that softmax is slow ? Note that the Vela performance estimator is very rough and cannot be trusted for performance analysis. To evaluate the performance of an operator or model on the Ethos-U, you either need to run the operator/model on the FVP, or on silicon.

How many NPU cycles does the softmax consume and how many cycles does the full model consume ? It would be useful to get that metric to understand the impact of softmax from the full network. I believe the FVP layer-per-layer profiling feature($ ./examples/arm/run.sh --target=ethos-u85-128 --model_name=mv2 --memory_mode=Shared_Sram --model_explorer --perf_overlay --visualize_tosa) can be useful to get such data.

[eladc4]

Hi @gggekov ,

I didn't run the FVP yet, so the analysis on the Vela output. The tables below are the summary of the Vela output per operator on the Whisper-tiny. As you can see, it is dominated by ReducedMax (and AvgPool, which I opened a separate issue for #18316).

Compile Spec

ethos_compile_spec = {
  'target': 'ethos-u85-128',
  'system_config': 'Ethos_U85_SYS_DRAM_Low',
  'memory_mode': 'Shared_Sram',
  'extra_flags': [ "--output-format=raw", "--debug-force-regor", "--verbose-all"]}

encoder

Summary by op type:

op type	count	Cycles	macs per cycle	op time (s)
AvgPool	18	2,613,172,050	0.003968	5.23
ReduceMax	4	2,268,378,360	0.023806	4.54
Transpose	34	406,270,832	0.303285	0.81
ReduceSum	13	238,480,614	0.248171	0.48
Conv2D	27	163,041,796	72.239440	0.33
Sub	22	77,184,000	0.000000	0.15
MatMul	48	54,737,856	126.274584	0.11
Mul	34	37,632,000	0.000000	0.08
MemoryCopy	30	17,295,360	0.999112	0.03
LUT	10	12,322,568	0.891048	0.02
Add	9	2,592,000	0.000000	0.01
Rescale	3	288,576	5.988024	0.00
Gather	1	0	0.000000	0.00

Batch Inference time 13072.52 ms, 0.08 inferences/s (batch size 1)

decoder

Summary by op type:

op type	count	Cycles	macs per cycle	op time (s)
AvgPool	26	566,722,650	0.003964	1.13
ReduceMax	8	392,346,576	0.023742	0.78
Transpose	68	125,411,676	0.230081	0.25
ReduceSum	21	42,616,178	0.244935	0.09
Conv2D	42	29,909,736	123.497339	0.06
Sub	34	13,543,200	0.000000	0.03
Mul	58	10,056,000	0.000000	0.02
MatMul	96	9,614,160	124.017075	0.02
MemoryCopy	39	3,879,168	0.994853	0.01
LUT	12	3,709,052	0.375622	0.01
Add	16	2,138,400	0.000000	0.00
Rescale	8	116,736	5.921053	0.00

Batch Inference time 2523.43 ms, 0.40 inferences/s (batch size 1)

[gggekov]

Could you run the model with the layer-per-layer profiling on the FVP ? The Vela output is not accurate from performance standpoint(you shouldn't trust the cycles column in the table)