EngineX/xc-llm-kunlun
3
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

[Kernel] Optimize the recurrent op

Browse Source
This commit is contained in:
ldh2020
2025-12-21 11:18:00 +08:00
committed by GitHub
parent 58c1db5073
commit 004e164bdb
1 changed files with 2 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
vllm_kunlun/ops/fla/fused_recurrent.py
View File

@@ -44,6 +44,7 @@ class FusedRecurrentFunction(torch.autograd.Function):
h0_indices=ssm_state_indices, h0_indices=ssm_state_indices,
num_accepted_tokens=num_accepted_tokens, num_accepted_tokens=num_accepted_tokens,
use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
is_h0_transposed=True
) )
return o, final_state return o, final_state