FlexAttention ModIndex misses cache hit for autograd func

[drisspg]

Summary

vllm-project/vllm#16078, while working on this Richard and I noticed that we are missing cache on repeated runs to "compile_block_mask" because of mod_index autograd func

pytorch/torch/_dynamo/_trace_wrapped_higher_order_op.py

Line 141 in 21c2565

return mod_index(args[0], index_args)

Fix is to check if grad_mod is enabled / x requries grad. If so run func else: call contents of foward

cc @chauhang @penguinwu @zou3519 @ydwu4 @bdhirsh @Chillee @yanboliang @BoyuanFeng