[examples] Add a variable-length KDA forward

[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] Add a variable-length KDA forward

FLA, vLLM and FlashKDA all take a variable-length batch the same way: one tensor of all the tokens plus cu_seqlens marking the boundaries. The engine had no way to express it, so a ragged batch had to be padded to its longest sequence. helion_chunk_kda now takes cu_seqlens and runs the batch as it is. Forward only: FlashKDA implements no backward, and vLLM is an inference engine.

• cu_seqlens switches helion_chunk_kda to a varlen batch: inputs token-major [1, T_total, H, D], state [N, H, D, DV] one per sequence, and state_v_first for the [N, H, DV, D] vLLM and FlashKDA hold.
• Five varlen kernels, one per stage of the dense diag_anchored pipeline.
• chunk_fwd_h_delta_varlen_helion runs the serial state pass as one program per (head, sequence), looping over that sequence's chunks chunk_offsets[n] : chunk_offsets[n + 1]. The state is initialized inside that loop, so each sequence starts from its own h0 and writes its own final state.
• naive_recurrent_reference takes cu_seqlens and resets its state at each boundary.
• l2norm_fwd_helion tiles one axis over a [N, D] view instead of three over [B, T, H, D].

test_linear_kda_varlen runs the varlen path through the same harness as test_linear_kda. A kda_varlen variant in benchmarks/run_linattn.py benchmarks it on FlashKDA's sequence partitions rather than the dense grid, and is forward-only, so it emits no -bwd dashboard row.

%FLA = 100 * fla_ms / helion_ms on H100, autotuned with HELION_AUTOTUNE_EFFORT=full. Higher is faster than FLA.

Shapes (<partition>_T*_H*_D* encodes the sequence lengths, total tokens, heads, head dim), from FlashKDA's benchmarks/bench_fwd.py:

• fixed_T8192_H96_D128, one sequence: [8192]
• fixed_T8192_H64_D128, one sequence: [8192]
• ragged_T8192_H96_D128, six sequences: [1300, 547, 2048, 963, 271, 3063]
• ragged_T8192_H64_D128, six sequences: [1300, 547, 2048, 963, 271, 3063]
• uniform_T8192_H96_D128, eight sequences: [1024] * 8
• uniform_T8192_H64_D128, eight sequences: [1024] * 8

Forward:

variant	fixed_T8192_H96_D128	fixed_T8192_H64_D128	ragged_T8192_H96_D128	ragged_T8192_H64_D128	uniform_T8192_H96_D128	uniform_T8192_H64_D128
kda_varlen	101%	90%	95%	89%	99%	90%