[Bug]: Eagle input triton kernel has bug in index #15908
Description
🐛 Describe the bug
The triton kernel for preparing eagle input in #15729 seems to have a bug where it writes the same value across blocks.
Added a colab notebook to highlight the bug and the fix: https://colab.research.google.com/drive/19RjTr8K2TTa-Zw7OUW7Bw0UfnMunUEFf?usp=sharing
The example assumes 4 draft tokens, block size 2, 0 rejections. The 1st output shows that values are repeated for every 2 places since block size is 2.
@WoosukKwon lmk if I misunderstood the indexing.