splash attention not compatible with shard_map(check_vma=True)

[hb-scaler]

current splash attention does not specify vma for the output shape dtype structure, which causes shard_map(check_vma=True) to fail. Using shard_map(check_vma=False) will cause some compiler optimization passes to be disabled