[ttx/npu] Optimize lightning_indexer_kernel by deferring k_scale to post-dot

[lyujheng]

Description

Optimize lightning_indexer_kernel by deferring k_scale application from pre-scaling K tensor
to post-dot product scaling, fully leveraging AIC and AIV core memory bandwidth.

Changes

• Remove per-element K tensor pre-scaling (k = k * k_scale[:, None])
• Apply k_scale after QK dot product as a lightweight broadcast multiply

Performance

Benchmark on Ascend 910C NPU and triton-ascend 3.2.0 (device latency):

Shape (B, M/N, K)	dtype	Before (us)	After (us)	Speedup
(128, 256, 128)	bf16	23041.37	3775.62	6.10x
(128, 256, 128)	fp16	23086.30	3776.70	6.11x
(128, 256, 128)	fp32	12731.75	4275.99	2.98x
(24, 1024, 128)	bf16	81850.27	22634.35	3.62x
(24, 1024, 128)	fp16	81893.52	22645.16	3.62x
(24, 1024, 128)	fp32	65051.82	28828.05	2.26x
(24, 16384, 128)	bf16	1410.02	420.78	3.35x
(24, 16384, 128)	fp16	1402.41	420.66	3.33x
(24, 16384, 128)	fp32	1041.73	502.65	2.07x

Average speedup: ~3.6x (bf16/fp16 gains significantly higher than fp32).

Accuracy

All accuracy tests passed across all test shapes and dtypes (bf16/fp16/fp32).

[gemini-code-assist]

Code Review

This pull request optimizes the lightning_indexer_kernel by moving the k_scale multiplication to occur after the dot product instead of before. This change reduces the computational overhead of element-wise multiplications by performing the scaling on the result of the dot product. I have no further feedback to provide.

[lyujheng]

Hi @wwens7, could you help review this PR when you have a chance?

This PR optimizes lightning_indexer_kernel by fusing k_scale into post-dot scaling, achieving ~3.6x average speedup on Ascend 910C NPU. All accuracy tests passed locally across bf16/fp16/fp32.

The change is straightforward — moving the scaling from pre-dot element-wise multiply to post-dot broadcast, reducing AIC-AIV data transfer overhead by avoiding global memory round-trips.

Thanks!

[Neuromancer42]

LGTM