revert change to cu_seqlen_k and max_k when preparing from position_ids

[winglian]

What does this PR do?

follow up to PR #39622 as these new lines from #39474 causes this during training w multi-gpu

[rank2]:   File "/root/miniconda3/envs/py3.11/lib/python3.11/site-packages/flash_attn/flash_attn_interface.py", line 165, in _flash_attn_varlen_forward
[rank2]:     out, softmax_lse, S_dmask, rng_state = flash_attn_gpu.varlen_fwd(
[rank2]:                                            ^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]: RuntimeError: cu_seqlens_k must have shape (batch_size + 1)

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?

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.

[ArthurZucker]

okay, as we discussed, generate now prepares the inputs for us so does not use this branch