[Qwen3.8-flash-next] QSAIndexer can not train

System Info

In the official technical report, the QSA indexer can be trained using KL loss, but in the current code, the parameters of the indexer cannot be trained.

Who can help?

No response

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
• My own task or dataset (give details below)

Reproduction

selected_block_indices = scores.topk(min(self.block_topk, num_complete_blocks), dim=0).indices. The gradient is truncated here.

Expected behavior

support training the parameters of the QSA indexer.