QUESTION: How to implement a tree attention with flashinfer about flashinfer

[zhyncs]

In order to support the feature of token position disruption brought by speculative decoding, two adjustments need to be made: one is the cos/sin matrix of RoPE, and the other is replacing casual mask with tree mask. With this, it will be very convenient to implement algorithms such as Medusa, EAGLE. From the document at https://docs.flashinfer.ai/index.html, it is currently not supported yet.

[jpf888]

We also need to support MEDUSA when we use MLC-LLM again, and we have seen that Tensorrtllm supports MEDUSA

[zhyncs]

We also need to support MEDUSA when we use MLC-LLM again, and we have seen that Tensorrtllm supports MEDUSA

The current implementation of Medusa in TensorRT-LLM is not fully functional, nor is it a SOTA implementation. By the way, if Medusa is not implemented based on tree mask, you can directly add a verification module at the location of model output without modifying the kernel code in the project. However, performance will be slightly worse and there will be redundant validation.

[UranusSeven]

In order to support the feature of token position disruption brought by speculative decoding, two adjustments need to be made: one is the cos/sin matrix of RoPE, and the other is replacing casual mask with tree mask. With this, it will be very convenient to implement algorithms such as Medusa, EAGLE. From the document at https://docs.flashinfer.ai/index.html, it is currently not supported yet.

Agree. But by using BatchPrefillWithPagedKVCacheWrapper, we can kind of sidestep the whole attention mask thing by just turning one draft sequence into a batch.