[V1][Spec Decode][Feature] support tree attention in flash-linear-attenton

[menggeliu1205]

Purpose

Tree attention is not supported in flash-linear-attention, which can verify a tree structure of draft tokens in Qwen3-Next model. To support that, added a new parameter to gdn triton kernel called retrieve_parent_token. This parameter recodes the parent node of each draft token using a 2D (batch_size, max_spec_len+1) tensor. This feature is only enabled if "retrieve_parent_token" is provided. Otherwise, it changes nothing.

And I add two test:

1. tests/v1/spec_decode/test_causal_conv1d_with_eagle_tree.py: test function causal_conv1d_update in speculative decoding with "retrieve_parent_token" or not.
2. tests/v1/spec_decode/test_recurrent_gated_delta_with_eagle_tree.py: test function fused_recurrent_gated_delta_rule in speculative decoding with "retrieve_parent_token" or not.

Test Plan

Test Result

Added test case pass.

TODO

• based on 22752, support speculative decoding with eagle draft tree on Qwen3-Next model.

[gemini-code-assist]

Code Review

This pull request adds support for tree-structured speculative decoding in flash-linear-attention by introducing a retrieve_parent_token parameter to the Triton kernels. The changes primarily affect fused_recurrent.py and causal_conv1d.py, and new tests are included to validate this functionality. My review identifies critical performance and correctness issues in the new Triton kernel implementations. Specifically, both kernels use an inefficient pattern for indexing parent tokens, and one of them contains a potential bug due to using an incorrect variable for indexing. I've provided suggestions to fix these issues, which should improve performance and ensure correctness.