hexagon: optimize FA DMA mask cache

[jhen0409]

Overview

This PR is an improvement for flash_attn_ext_f16_thread which replaces the original cache search on every push.

The FA mask blocks are pushed in the same order for every head: (token, block) within 1 mask head. So sizing the cache to that cycle (n_blocks * neq1 * mask->ne[2]) (<= 128) is enough. That dma_cache_push_seq (new function) can find every repeat without a search.

Additional information

Setup: IQ-9075 (Hexagon v73), single NPU HTP0:0, llama-bench -p 512 -n 64 -r 2. master is 260eebc.

model	test	master (t/s)	patch (t/s)	speedup
Llama-3.2-1B-Instruct Q4_0	tg64 @ d0	49.14	49.39	+0.5%
Llama-3.2-1B-Instruct Q4_0	tg64 @ d2048	30.05	31.52	+4.9%
Llama-3.2-1B-Instruct Q4_0	tg64 @ d4096	22.37	24.03	+7.4%
Llama-3.2-1B-Instruct Q4_0	tg64 @ d8192	14.86	16.30	+9.7%
granite-4.2-3b Q4_0	tg64 @ d0	17.92	18.04	+0.7%
granite-4.2-3b Q4_0	tg64 @ d2048	10.60	11.09	+4.7%
granite-4.2-3b Q4_0	tg64 @ d4096	7.79	8.28	+6.2%
granite-4.2-3b Q4_0	tg64 @ d8192	5.09	5.50	+8.2%

Old (before #29199)

model	test	master (t/s)	patch (t/s)	speedup
Llama-3.2-1B-Instruct Q4_0	tg64 @ d0	40.10	41.27	+2.9%
Llama-3.2-1B-Instruct Q4_0	tg64 @ d2048	22.34	26.17	+17.1%
Llama-3.2-1B-Instruct Q4_0	tg64 @ d4096	15.95	19.79	+24.1%
Llama-3.2-1B-Instruct Q4_0	tg64 @ d8192	10.15	13.38	+31.8%
granite-4.2-3b Q4_0	tg64 @ d0	14.49	14.88	+2.7%
granite-4.2-3b Q4_0	tg64 @ d2048	7.72	9.12	+18.2%
granite-4.2-3b Q4_0	tg64 @ d4096	5.40	6.77	+25.4%
granite-4.2-3b Q4_0	tg64 @ d8192	3.37	4.48	+32.9%

PPL not changed and pp512 t/s keep original.

Requirements

• I have read and agree with the contributing guidelines
• AI usage disclosure: YES, I discovered this while learning how this backend works, helped by Claude Code, and did validation and review.

[max-krasnyansky]

@jhen0409 are the perf improvements measurable?
I played with the cache quite a bit in the past and it's small enough that the "search" is really just a few iterations.
Not sure it's worth adding more logic if the improvements are just a few cycles. I got more FA updates coming along with #29199 let's remeasure after that and see. You can use --profile 3 (aka op-trace) to get cycle counts / stats for all FA stages.

[jhen0409]

@jhen0409 are the perf improvements measurable? I played with the cache quite a bit in the past and it's small enough that the "search" is really just a few iterations. Not sure it's worth adding more logic if the improvements are just a few cycles. I got more FA updates coming along with #29199 let's remeasure after that and see. You can use --profile 3 (aka op-trace) to get cycle counts / stats for all FA stages.

Yeah it's measurable. Rebased. TG has only slight gains of 0-10% after #29199. (updated PR description)

GGML_HEXAGON_PROFILE=3 GGML_HEXAGON_OPTRACE=8192 GGML_HEXAGON_OPBATCH=1 \
  test-backend-ops test -b HTP0:0 -o FLASH_ATTN_EXT --test-file cases.txt > trace.log 2>&1
python3 scripts/snapdragon/ggml-hexagon-profile.py trace.log --timeline summary --filter FLASH_ATTN

cases.txt (p-d64-h32-kh8-k4096-nomask-s1024 is reference):

74 0 64 32 1 1 3 1040187392 0 0 4 0 64 1 32 1 4 8192 256 8192 1 64 4096 8 1 2 1024 128 4194304 1 64 4096 8 1 2 1024 128 4194304 1 4096 1 1 1 2 8192 8192 8192 c-llama-1b-d64-h32-kh8-k4096
74 0 64 32 1 1 3 1040187392 0 0 3 0 64 1 32 1 4 8192 256 8192 1 64 4096 8 1 2 1024 128 4194304 1 64 4096 8 1 2 1024 128 4194304 p-d64-h32-kh8-k4096-nomask-s1024
74 0 64 40 1 1 3 1040187392 0 0 4 0 64 1 40 1 4 10240 256 10240 1 64 4096 8 1 2 1024 128 4194304 1 64 4096 8 1 2 1024 128 4194304 1 4096 1 1 1 2 8192 8192 8192 granite3b-d64-h40-kh8-k4096

master (after #29199)

====================================================================================================
FLASH_ATTN_EXT (64:1:32:1 x 64:4096:8:1 x 64:4096:8:1 x 4096:1 -> 64:32 : f32 x f16 x f16 x f16 -> f32) - 1514 usec 2211016 cycles
====================================================================================================
  Thread 0 (HVX)  : DMA 7251185 (328.0%) | O-PROC 7178 (0.3%) | QK-FA 1131059 (51.2%) | SFM-FA 615042 (27.8%)
  Thread 1 (HVX)  : DMA 7255103 (328.1%) | O-PROC 6993 (0.3%) | QK-FA 1131351 (51.2%) | SFM-FA 613936 (27.8%)
  Thread 2 (HVX)  : DMA 7251237 (328.0%) | O-PROC 8957 (0.4%) | QK-FA 1128314 (51.0%) | SFM-FA 613745 (27.8%)
  Thread 3 (HVX)  : DMA 7253115 (328.0%) | O-PROC 7306 (0.3%) | QK-FA 1130485 (51.1%) | SFM-FA 613767 (27.8%)
====================================================================================================
FLASH_ATTN_EXT (64:1:32:1 x 64:4096:8:1 x 64:4096:8:1 -> 64:32 : f32 x f16 x f16 -> f32) - 1312 usec 1915974 cycles
====================================================================================================
  Thread 0 (HVX)  : DMA 3940058 (205.6%) | O-PROC 3379 (0.2%) | QK-FA 1119129 (58.4%) | SFM-FA 616769 (32.2%)
  Thread 1 (HVX)  : DMA 3943469 (205.8%) | O-PROC 4979 (0.3%) | QK-FA 1119807 (58.4%) | SFM-FA 614290 (32.1%)
  Thread 2 (HVX)  : DMA 3937660 (205.5%) | O-PROC 4744 (0.2%) | QK-FA 1118162 (58.4%) | SFM-FA 614352 (32.1%)
  Thread 3 (HVX)  : DMA 3938039 (205.5%) | O-PROC 4209 (0.2%) | QK-FA 1119280 (58.4%) | SFM-FA 614156 (32.1%)
====================================================================================================
FLASH_ATTN_EXT (64:1:40:1 x 64:4096:8:1 x 64:4096:8:1 x 4096:1 -> 64:40 : f32 x f16 x f16 x f16 -> f32) - 1798 usec 2624359 cycles
====================================================================================================
  Thread 0 (HVX)  : DMA 8483523 (323.3%) | O-PROC 2842 (0.1%) | QK-FA 1410761 (53.8%) | SFM-FA 768066 (29.3%)
  Thread 1 (HVX)  : DMA 8499473 (323.9%) | O-PROC 2794 (0.1%) | QK-FA 1411241 (53.8%) | SFM-FA 765946 (29.2%)
  Thread 2 (HVX)  : DMA 8386368 (319.6%) | O-PROC 5493 (0.2%) | QK-FA 1409367 (53.7%) | SFM-FA 767267 (29.2%)
  Thread 3 (HVX)  : DMA 8486016 (323.4%) | O-PROC 2819 (0.1%) | QK-FA 1410937 (53.8%) | SFM-FA 766664 (29.2%)

PR (rebased)

====================================================================================================
FLASH_ATTN_EXT (64:1:32:1 x 64:4096:8:1 x 64:4096:8:1 x 4096:1 -> 64:32 : f32 x f16 x f16 x f16 -> f32) - 1319 usec 1926612 cycles
====================================================================================================
  Thread 0 (HVX)  : DMA 5852762 (303.8%) | O-PROC 7419 (0.4%) | QK-FA 1127447 (58.5%) | SFM-FA 615259 (31.9%)
  Thread 1 (HVX)  : DMA 5862991 (304.3%) | O-PROC 8224 (0.4%) | QK-FA 1127058 (58.5%) | SFM-FA 614687 (31.9%)
  Thread 2 (HVX)  : DMA 5865536 (304.4%) | O-PROC 6409 (0.3%) | QK-FA 1126871 (58.5%) | SFM-FA 614544 (31.9%)
  Thread 3 (HVX)  : DMA 5856998 (304.0%) | O-PROC 6498 (0.3%) | QK-FA 1127033 (58.5%) | SFM-FA 614458 (31.9%)
====================================================================================================
FLASH_ATTN_EXT (64:1:32:1 x 64:4096:8:1 x 64:4096:8:1 -> 64:32 : f32 x f16 x f16 -> f32) - 1305 usec 1905073 cycles
====================================================================================================
  Thread 0 (HVX)  : DMA 3912946 (205.4%) | O-PROC 5806 (0.3%) | QK-FA 1114809 (58.5%) | SFM-FA 615783 (32.3%)
  Thread 1 (HVX)  : DMA 3923764 (206.0%) | O-PROC 5655 (0.3%) | QK-FA 1114679 (58.5%) | SFM-FA 614903 (32.3%)
  Thread 2 (HVX)  : DMA 3916246 (205.6%) | O-PROC 5825 (0.3%) | QK-FA 1114110 (58.5%) | SFM-FA 615540 (32.3%)
  Thread 3 (HVX)  : DMA 3931724 (206.4%) | O-PROC 4865 (0.3%) | QK-FA 1114092 (58.5%) | SFM-FA 615717 (32.3%)
====================================================================================================
FLASH_ATTN_EXT (64:1:40:1 x 64:4096:8:1 x 64:4096:8:1 x 4096:1 -> 64:40 : f32 x f16 x f16 x f16 -> f32) - 1626 usec 2374451 cycles
====================================================================================================
  Thread 0 (HVX)  : DMA 7200304 (303.2%) | O-PROC 6065 (0.3%) | QK-FA 1407369 (59.3%) | SFM-FA 769335 (32.4%)
  Thread 1 (HVX)  : DMA 7208450 (303.6%) | O-PROC 6038 (0.3%) | QK-FA 1407609 (59.3%) | SFM-FA 767589 (32.3%)
  Thread 2 (HVX)  : DMA 7213150 (303.8%) | O-PROC 5983 (0.3%) | QK-FA 1406960 (59.3%) | SFM-FA 766787 (32.3%)
  Thread 3 (HVX)  : DMA 7206171 (303.5%) | O-PROC 6100 (0.3%) | QK-FA 1405005 (59.2%) | SFM-FA 767683 (32.3%)

[max-krasnyansky]

Here is a better version #29282
This one add scalar division into the hot path and introduces more register spills (also regresses for large context where we'll have zero hit rate).
What we really need is a simple direct-mapped cache for the HVX path. HMX path should still use LRU.

[jhen0409]

Here is a better version #29282

This one add scalar division into the hot path and introduces more register spills (also regresses for large context where we'll have zero hit rate).

What we really need is a simple direct-mapped cache for the HVX path. HMX path should still use LRU.

Yeah that's better. Thanks!