perf: prefetch page indices for mla kernel

[yzh119]

Followup of #952

cc @abcdabcd987

Before this PR

Config: batch_size=64, seq_len=1024, num_heads=64
Memory bandwidth: 1509.87 GB/s
FLOPs: 163.25 TFLOPs
Config: batch_size=64, seq_len=1024, num_heads=128
Memory bandwidth: 1766.19 GB/s
FLOPs: 345.46 TFLOPs
Config: batch_size=128, seq_len=1024, num_heads=64
Memory bandwidth: 2307.97 GB/s
FLOPs: 249.55 TFLOPs
Config: batch_size=128, seq_len=1024, num_heads=128
Memory bandwidth: 1975.24 GB/s
FLOPs: 386.35 TFLOPs
Config: batch_size=768, seq_len=1024, num_heads=64
Memory bandwidth: 2871.63 GB/s
FLOPs: 310.49 TFLOPs
Config: batch_size=768, seq_len=1024, num_heads=128
Memory bandwidth: 2225.07 GB/s
FLOPs: 435.21 TFLOPs
Config: batch_size=64, seq_len=2048, num_heads=64
Memory bandwidth: 1948.15 GB/s
FLOPs: 222.38 TFLOPs
Config: batch_size=64, seq_len=2048, num_heads=128
Memory bandwidth: 1973.36 GB/s
FLOPs: 426.74 TFLOPs
Config: batch_size=128, seq_len=2048, num_heads=64
Memory bandwidth: 2625.63 GB/s
FLOPs: 299.72 TFLOPs
Config: batch_size=128, seq_len=2048, num_heads=128
Memory bandwidth: 2121.92 GB/s
FLOPs: 458.86 TFLOPs
Config: batch_size=768, seq_len=2048, num_heads=64
Memory bandwidth: 2996.11 GB/s
FLOPs: 342.01 TFLOPs
Config: batch_size=768, seq_len=2048, num_heads=128
Memory bandwidth: 2146.40 GB/s
FLOPs: 464.16 TFLOPs
Config: batch_size=64, seq_len=8192, num_heads=64
Memory bandwidth: 2717.28 GB/s
FLOPs: 323.71 TFLOPs
Config: batch_size=64, seq_len=8192, num_heads=128
Memory bandwidth: 2129.24 GB/s
FLOPs: 500.04 TFLOPs
Config: batch_size=128, seq_len=8192, num_heads=64
Memory bandwidth: 3002.75 GB/s
FLOPs: 357.72 TFLOPs
Config: batch_size=128, seq_len=8192, num_heads=128
Memory bandwidth: 2101.93 GB/s
FLOPs: 493.63 TFLOPs
Config: batch_size=768, seq_len=8192, num_heads=64
Memory bandwidth: 3083.42 GB/s
FLOPs: 367.33 TFLOPs
Config: batch_size=768, seq_len=8192, num_heads=128
Memory bandwidth: 2064.96 GB/s
FLOPs: 484.95 TFLOPs

After this PR

Config: batch_size=64, seq_len=1024, num_heads=64
Memory bandwidth: 1596.98 GB/s
FLOPs: 172.67 TFLOPs
Config: batch_size=64, seq_len=1024, num_heads=128
Memory bandwidth: 1685.22 GB/s
FLOPs: 329.62 TFLOPs
Config: batch_size=128, seq_len=1024, num_heads=64
Memory bandwidth: 2280.49 GB/s
FLOPs: 246.58 TFLOPs
Config: batch_size=128, seq_len=1024, num_heads=128
Memory bandwidth: 1917.53 GB/s
FLOPs: 375.06 TFLOPs
Config: batch_size=768, seq_len=1024, num_heads=64
Memory bandwidth: 2869.03 GB/s
FLOPs: 310.21 TFLOPs
Config: batch_size=768, seq_len=1024, num_heads=128
Memory bandwidth: 2208.35 GB/s
FLOPs: 431.94 TFLOPs
Config: batch_size=64, seq_len=2048, num_heads=64
Memory bandwidth: 2047.44 GB/s
FLOPs: 233.72 TFLOPs
Config: batch_size=64, seq_len=2048, num_heads=128
Memory bandwidth: 1936.08 GB/s
FLOPs: 418.67 TFLOPs
Config: batch_size=128, seq_len=2048, num_heads=64
Memory bandwidth: 2617.48 GB/s
FLOPs: 298.79 TFLOPs
Config: batch_size=128, seq_len=2048, num_heads=128
Memory bandwidth: 2105.97 GB/s
FLOPs: 455.41 TFLOPs
Config: batch_size=768, seq_len=2048, num_heads=64
Memory bandwidth: 2999.55 GB/s
FLOPs: 342.40 TFLOPs
Config: batch_size=768, seq_len=2048, num_heads=128
Memory bandwidth: 2181.54 GB/s
FLOPs: 471.75 TFLOPs
Config: batch_size=64, seq_len=8192, num_heads=64
Memory bandwidth: 2780.86 GB/s
FLOPs: 331.29 TFLOPs
Config: batch_size=64, seq_len=8192, num_heads=128
Memory bandwidth: 2176.12 GB/s
FLOPs: 511.05 TFLOPs
Config: batch_size=128, seq_len=8192, num_heads=64
Memory bandwidth: 3031.58 GB/s
FLOPs: 361.15 TFLOPs
Config: batch_size=128, seq_len=8192, num_heads=128
Memory bandwidth: 2165.73 GB/s
FLOPs: 508.61 TFLOPs
Config: batch_size=768, seq_len=8192, num_heads=64
Memory bandwidth: 3126.37 GB/s
FLOPs: 372.45 TFLOPs
Config: batch_size=768, seq_len=8192, num_heads=128
Memory bandwidth: 2142.42 GB/s
FLOPs: 503.14 TFLOPs

[abcdabcd987]

Significant boost

[muoshuosha]

Hi,I have a question about this pr.
Why the performance of some cases is drop and some of other cases are improved after this pr?
Why can't "prefetch page indices" improve the performance of all cases?