chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

[FLA] Add explicit kernel arguments to kda.py for Kimi Linear support (#14561)

Browse Source
This commit is contained in:
Alison Shao
2025-12-06 22:45:41 -08:00
committed by GitHub
parent 3c7886ec4c
commit 41d61faa99
1 changed files with 3 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
python/sglang/srt/layers/attention/fla/kda.py
View File

@@ -102,8 +102,11 @@ def fused_recurrent_kda_fwd(
# stride_final_state_token=stride_final_state_token,
# stride_indices_seq=stride_indices_seq,
# stride_indices_tok=stride_indices_tok,
USE_INITIAL_STATE=initial_state is not None,
STORE_FINAL_STATE=final_state is not None,
IS_BETA_HEADWISE=beta.ndim == v.ndim,
USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel,
IS_VARLEN=cu_seqlens is not None,
# INPLACE_FINAL_STATE=inplace_final_state,
IS_KDA=True,
num_warps=num_warps,