[yeahdongcn]

Overview

Follow-up to #29901 and #29990. Only MUSA arch 21 caps static shared memory at 28 KB. With mcc 5.2.0, mp_22 allows 72 KB and mp_31 allows 192 KB, so the 33 KB tile kernel from #29901 only fails to build on mp_21.

This brings the tile kernel back on MUSA. It takes the head pass staging from 8b84aec (thanks @ServeurpersoCom) and limits the two heads per pass to __MUSA_ARCH__ < 220:

• mp_21: tile kernel, two heads per pass, 24984 B static shared memory
• mp_22: tile kernel, single pass, 33176 B
• mp_31: tile kernel, single pass, 33152 B
• CUDA and HIP: unchanged

Testing Done

Built with the default MUSA_ARCHITECTURES (21;22;31) in musa_sdk:5.2.0-devel-ubuntu22.04-s5000. test-backend-ops -o LIGHTNING_INDEXER passes 192/192 on MTT S5000 and MTT S80.

test-backend-ops perf -o LIGHTNING_INDEXER -p 'nh=4,kv=65536,nb=2048,ns=1,nm=1,type_K=f16', time in ms (S5000 is the median of 3 runs):

GPU	master	PR
MTT S5000	85146 us/run	20253 us/run
MTT S80	1564954 us/run	227168 us/run

nb=1 is unchanged, it still uses the vector kernel.

test-backend-ops perf -o LIGHTNING_INDEXER -p 'nh=4,kv=65536,nb=(512|2048),ns=1,', time in ms (S5000 is the median of 3 runs):

nb	type_K	S5000 master	S5000 PR	speedup	S80 master	S80 PR	speedup
512	f16	21.32	5.11	4.17x	391.3	56.8	6.89x
512	bf16	21.34	5.11	4.18x	391.3	58.1	6.73x
512	f32	19.03	5.35	3.56x	391.1	56.6	6.91x
512	q8_0	19.37	5.17	3.74x	393.3	60.2	6.53x
512	q4_0	19.17	4.90	3.91x	392.7	62.2	6.31x
2048	f16	85.15	20.25	4.20x	1565.0	227.2	6.89x
2048	bf16	85.26	20.25	4.21x	1565.0	232.4	6.73x
2048	f32	75.86	21.52	3.53x	1564.1	226.3	6.91x
2048	q8_0	77.35	20.41	3.79x	1572.9	240.6	6.54x
2048	q4_0	76.58	19.26	3.98x	1570.4	248.5	6.32x

Additional information

Requirements

• I have read and agree with the contributing guidelines
• AI usage disclosure: YES (Claude Code)

[ServeurpersoCom]

LGTM, thanks for testing on real hardware! Limiting the head passes to arch 21 keeps CUDA and HIP on the merged kernel, and the red jobs are unrelated

[am17an]

@ServeurpersoCom let a CUDA maintainer merge CUDA PRs. Thanks

[ServeurpersoCom]

Got it, sorry about that, I'll leave CUDA merges to you from now on.