[CPU] add indices in chunk_gated_delta_rule (#29267)

2026-06-26 07:51:07 +08:00
parent c8c6757bed
commit 1ba7c79761
8 changed files with 49 additions and 17 deletions
@@ -101,6 +101,7 @@ def chunk_gated_delta_rule_cpu(
    cu_seqlens,
    head_first,
    use_qk_l2norm_in_kernel,
+    initial_state_indices,
 ):
    core_attn_out, last_recurrent_state = (
        torch.ops.sgl_kernel.chunk_gated_delta_rule_cpu(
@@ -114,6 +115,7 @@ def chunk_gated_delta_rule_cpu(
            cu_seqlens,
            head_first,
            use_qk_l2norm_in_kernel,
+            initial_state_indices,
        )
    )
    h = None  # Todo: add return h support