fix flex attn when optional args aren't passed

[winglian]

What does this PR do?

#37307 (llama4) introduced a regression in make_flex_block_causal_mask with new optional args that llama4 uses. It also added padding, but padding only works if the key_length isn't None. For all models except llama4, flex attention is now broken since none of them pass key_length, and currently they error with.

[rank3]:   File "/root/miniconda3/envs/py3.11/lib/python3.11/site-packages/torch/nn/functional.py", line 5209, in pad                                                                                                                                        
[rank3]:     return torch._C._nn.pad(input, pad, mode, value)                                                                                                                                                                                                
[rank3]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^                                                                                                                                                                                                
[rank3]: TypeError: pad(): argument 'pad' failed to unpack the object at pos 2 with error "type must be tuple of ints,but got NoneType"                                                                                                                      

Also, the batch size was likely mistakenly hard-coded to 1, which probably shouldn't be the case.

Fixes # (issue)

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline,
  Pull Request section?
• Was this discussed/approved via a Github issue or the forum? Please add a link
  to it if that's the case.
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?

Who can review?

@ArthurZucker

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

[github-actions]

Hi 👋, thank you for opening this pull request! The pull request is converted to draft by default. The CI will be paused while the PR is in draft mode. When it is ready for review, please click the Ready for review button (at the bottom of the PR page). This will assign reviewers and trigger CI.

[ArthurZucker]

LGTM

[ArthurZucker]

Thanks for the prompt fix!