[examples] Apply the KDA input preamble in the kernels

[tarinduj]

Stacked PRs:

• [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] Add a variable-length KDA forward #3289
• ->[examples] Apply the KDA input preamble in the kernels #3288

---

[examples] Apply the KDA input preamble in the kernels

FlashKDA and vLLM both fold KDA's input preamble into the kernel: the q/k L2 norm, the gate activated from a raw projection, and the beta sigmoid. The engine required all three already applied. helion_chunk_kda now folds them in too, behind the same three flags FLA uses. The flags are forward only: FlashKDA implements no backward, and vLLM is an inference engine.

What's here

• use_qk_l2norm_in_kernel normalizes q and k in a new kernel, l2norm_fwd_helion.
• use_gate_in_kernel applies the gate inside chunk_cumsum_gc_helion, in fp32 before the sum.
• use_beta_sigmoid_in_kernel applies torch.sigmoid on the host.
• naive_recurrent_reference applies all three transforms itself, so it checks the kernels on the same pre-activation inputs.

Tests

test_linear_kda_fused_preamble runs the flags through the same harness as test_linear_kda. A kda_fused variant in benchmarks/run_linattn.py benchmarks it on the same shapes as kda; being forward-only it emits no -bwd dashboard row, and the linattn benchmark matrix goes from 7 shards to 8.

Benchmarks

%FLA = 100 * fla_ms / helion_ms on H100, autotuned with HELION_AUTOTUNE_EFFORT=full. Higher is faster than FLA. kda takes the transformed inputs, kda_fused the raw ones with the flags on.

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	119%	116%	104%	117%	109%	98%
kda_fused	116%	114%	123%	113%	107%	127%

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	118%	113%	113%	117%	109%	120%