Skip to content

Commit 666a1c9

Browse files
authored
feat(kda): add KDA (Kimi Delta Attention) kernels, backend, and tests (#1051)
- Add KDA Pallas kernel (python/sgl_jax/srt/kernels/kda/) with naive reference - Add KDA attention backend (linear/kda_backend.py) and short convolution layer - Wire KDA into hybrid_linear_attn_backend and radix_linear_attention - Add unit tests for TP and DP paths (test_kda_attention*.py, test_short_conv.py) - Register tests in unit-test-tpu-v6e-4 suite
1 parent e2c5f91 commit 666a1c9

11 files changed

Lines changed: 3917 additions & 4 deletions

File tree

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
from sgl_jax.srt.kernels.kda.kda import chunk_kda_fwd as chunk_kda
2+
from sgl_jax.srt.kernels.kda.naive import naive_recurrent_kda
3+
4+
__all__ = ["chunk_kda", "naive_recurrent_kda"]

0 commit comments

Comments
 (0)