[examples] Read only the causal columns in chunk_fwd_A_diag_anchored

[tarinduj]

Stacked PRs:

• [benchmark] Pretune the linear-attention kernels for CI #3298
• [examples] Add FlashKDA as a second baseline for kda_varlen #3291
• ->[examples] Read only the causal columns in chunk_fwd_A_diag_anchored #3290

---

[examples] Read only the causal columns in chunk_fwd_A_diag_anchored

The intra-chunk score matrix is causal, so everything above its diagonal is structurally zero, and chunk_fwd_A_diag_anchored_helion was computing it anyway: each BC_DIAG-row sub-block read all C columns and then masked most of the result away. It now reads only the columns that can be nonzero, 0 : (i + 1) * BC_DIAG for sub-block i, so the two matmuls span BC_DIAG * (1 + 2 + ... + NC) = C * (NC + 1) / 2 columns instead of NC * C. A and Akk arrive zeroed, so the skipped columns already hold zeros.

The diff reads as a manual unroll because a narrowing width has to be a literal: hl.arange rejects arithmetic on a specialized size, so each tier is written out and guarded on C. FLA's KDA does the same in chunk_intra.py. The tiers cover C of 32 and 64, the two chunk sizes FLA's KDA accepts, and the kernel asserts that.

kda and full_gla share the kernel across four call sites, two of them backward.

Benchmarks

%FLA = 100 * fla_ms / helion_ms on H100, autotuned with HELION_AUTOTUNE_EFFORT=full. Higher is faster than FLA. The baseline row is the same variant's table in #3200.

Forward:

variant	B1_T8192_H96_D128	B2_T16384_H16_D128	B4_T2048_H16_D128	B4_T4096_H64_D128	B8_T2048_H32_D256	B8_T1024_H8_D64
kda (#3200)	118%	117%	109%	116%	109%	104%
kda	129%	126%	98%	128%	111%	95%
full_gla (#3200)	143%	161%	128%	145%	134%	115%
full_gla	154%	169%	130%	155%	137%	108%

Forward + backward:

variant	B1_T8192_H96_D128	B2_T16384_H16_D128	B4_T2048_H16_D128	B4_T4096_H64_D128	B8_T2048_H32_D256	B8_T1024_H8_D64
kda (#3200)	117%	113%	114%	116%	109%	141%
kda	119%	115%	115%	119%	110%	130%
full_gla (#3200)	235%	248%	211%	243%	215%	139%
full_gla	240%	252%	200%	247%	216%	128%