Question about Flexprefill

[xiehuanyi]

Hi, I'm trying with Flexprefill algorithm. According to the paper https://arxiv.org/pdf/2502.20766, the flexprefill should include two parts, Vertical Slash Index Search and Query Aware Index Search. And the authors use Jensen-Shannon divergence (JSD) to determine which mask to use. However, I didn't find the implement of the JSD. I'm wondering if I missed something important or if I misunderstood the paper.

Thanks a lot! And Happy Labor Day!

[PiotrNawrot]

Please check out A.1.3 of our paper - we deliberately skipped the second part of their algorithm to be able to measure the effect of adaptive budgeting. You're correct about the original algorithm of FlexPrefill - it's more sophisticated than our implementation. Our goal was to survey and simplify existing algorithms to critical differences and then to ablate them individually, that's why we skipped the Query Aware part - so that's the only difference between FlexPrefill and Vertical-Slash is adaptive budgeting. That being said, I tested the FlexPrefill with Block-Sparse part and it performed on par or worse on our benchmark suite - this is consistent with the the original paper where the inclusion of Block-Sparse part resulted in minor improvements, according to Figure 5 from the FlexPrefill paper.