fix(integrations): prevent device mismatch in msa_attention
Author: leaderofARS
Pull request: yes

<!-- ci-dashboard-badge:start -->
[![CI](https://transformers-ci.lor-e.huggingface.cool/badge/pr?pr=46848)](https://transformers-ci.lor-e.huggingface.cool/d/pytest-observability-by-pr/pytest-observability-branch?var-pr=46848)
<!-- ci-dashboard-badge:end -->

## What does this PR do?                                                                                                                                   
                                                                                                                                                             
During single-token decoding (when  batch_size == 1  and  cache_position  is provided), the cumulative sequence boundary  cu_seqlens_k  is constructed by concatenating a GPU-based zero tensor with  valid_k  (representing the number of valid keys in the KV cache).                                              
                                                                                                                                                             
Since  cache_position  is typically initialized on the CPU,  valid_k  is created as a CPU tensor by default when cast to  int32 . Concatenating a CPU tensor with a GPU tensor directly causes a device mismatch  RuntimeError  ( RuntimeError: Expected all tensors to be on the same device... ).              
                                                                                                                                                             
This PR fixes the issue by explicitly mapping  valid_k  to the query tensor's device ( device=q.device ) when casting, ensuring all concatenated tensors reside on the same GPU.                                                                                                                                    
                                                                                                                                                             
## Code Agent Policy                                                                                                                                       
                                                                                                                                                             
[✓] I confirm that this is not a pure code agent PR.                                                                                                       
                                                                                                                                                             
## Before submitting                                                                                                                                       
                                                                                                                                                             
[✓] Did you read the contributor guideline https://huggingface.co/docs/transformers/contributing and the Pull Request https://huggingface.co/docs/transformers/pr_checks checks?                                                                                                                     
[ ] Was this discussed/approved via a Github issue or the forum https://discuss.huggingface.co/? Please add a link to it if that's the case.               
[ ] Did you make sure to update the documentation with your changes according to the guidelines https://github.com/huggingface/transformers/tree/main/docs?
[ ] Did you write any new necessary tests https://huggingface.co/docs/transformers/testing?                                                                
                                                                                                                                                             
## Who can review?                                                                                                                                         
                                                                                                                                                             
• Attention / Kernels: @vasqu @drbh

Comments:
Thanks for the review! That's a good point. I don't currently have access to a CUDA device locally, but I'll look into existing test patterns for GPU-specific attention kernels in Transformers and see if I can add a suitable regression test for this path. Any pointers to a relevant test location or similar example would be greatly appreciated.

The two failing checks both stem from a single infrastructure OOM kill ( exit code 137 ) on  test_run_wav2vec2_pretraining  in  examples_torch [shard 0/1]                                                                                                                                                       
 — unrelated to this PR. The wav2vec2 pretraining test is a known memory-heavy test and its CI runner ran out of memory.All checks directly related to this change pass. Could a maintainer re-trigger  examples_torch  or override the status check to unblock the merge?

Hi @vasqu , just following up on this PR. I believe I've addressed all of the review feedback, and the latest CI run completed successfully. Whenever you have a chance, I'd appreciate another look. Thanks again for all the detailed feedback throughout the review!

**[For maintainers]** Suggested jobs to run (before merge)

run-slow: minimax_m3_vl

<!-- ci-dashboard-recap:start -->

---

### CI recap

**Dashboard:** [View test results in Grafana](https://transformers-ci.lor-e.huggingface.cool/d/pytest-observability-by-pr/pytest-observability-branch?var-pr=46848)
**Latest run:** [29269183410:2](https://github.com/huggingface/transformers/actions/runs/29269183410)
**Result:** `success` | **Jobs:** 12 | **Tests:** 163,928 | **Failures:** 0 | **Duration:** 14h 19m

<!-- ci-dashboard-recap:end -->

Hi @vasqu, just following up on this PR. The latest changes address the SM100 guard feedback, and the updated CI completed successfully. Whenever you have a chance, I'd appreciate another look. Thanks!