Commit 666a1c9
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 suite1 parent e2c5f91 commit 666a1c9
11 files changed
Lines changed: 3917 additions & 4 deletions
File tree
- python/sgl_jax
- srt
- kernels/kda
- layers/attention
- linear
- test
- test/srt
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
| 1 | + | |
| 2 | + | |
| 3 | + | |
| 4 | + | |
0 commit comments