Validate the independent FLA 0.5 compatibility path

[morluto]

TL;DR

Test FlashKDA directly against FLA's independent Triton chunk_kda implementation using the matching raw-gate, raw-beta, and state-layout contract.

Motivation

The existing FLA comparison does not directly assert FlashKDA against FLA's chunked implementation:

• It uses pre-activated beta because FLA 0.5.0 lacked in-kernel beta activation.
• It omits the matching safe-gate option.
• It collects chunk_kda output without asserting it.
• Its chunk final-state comparison is warning-only.
• Once the correct dispatch flags are supplied, FLA may dispatch chunk_kda back into FlashKDA, producing a self-comparison.

FLA 0.5.1 provides the interface required to exercise the intended compatibility path.

Comparison path

Before

chunk_kda -- eligible dispatch --> FlashKDA
    |
    +-- output collected but not asserted
    +-- final state warning-only

After

FlashKDA output/state ---------+
                               +--> direct comparison
FLA Triton chunk_kda ----------+
  FLA_FLASH_KDA=0, raw g/beta, matching gate flags

Changes

• Declare compatibility with flash-linear-attention>=0.5.1,<0.6.
• Pass raw gate and beta inputs with:
  • use_gate_in_kernel=True
  • use_beta_sigmoid_in_kernel=True
  • safe_gate=True
  • state_v_first=True
• Disable FLA_FLASH_KDA for the independent reference and benchmark.
• Directly assert output and final state for fixed-length and variable-length cases.
• Update the README, test installation command, benchmark invocation, and generated benchmark metadata.

Recurrence scope

This PR does not change the FlashKDA recurrence.

FLA's KDA path applies beta while forming its triangularly transformed u and w terms. The resulting update is algebraically equivalent to FlashKDA's beta-scaled residual followed by the triangular inverse.

Suggested review order

1. tests/test_fwd.py — independent reference selection and direct assertions.
2. README.md and tests/test.sh — declared FLA compatibility contract.
3. benchmarks/bench_fwd.py — independent baseline configuration.
4. benchmarks/generate_benchmark_md.py — synchronized benchmark metadata.

The main risk to inspect is accidental redispatch from FLA back into FlashKDA.

Validation

• python -m py_compile benchmarks/bench_fwd.py benchmarks/generate_benchmark_md.py tests/test_fwd.py
• bash -n tests/test.sh
• git diff --check origin/master...HEAD