Fix CUDA index out of bounds for q_idx in VLM token type masking for Gemma3, PaliGemma, and example modular
Author: albertvillanova
Pull request: yes

Fix CUDA index out of bounds error that occurs during generation with static caches when using token type IDs for bidirectional image attention masking.

### Background

After PR
- #41505 

changed cache initialization behavior in `generate()`, a **latent bug in the VLM masking code was exposed**. The error manifests as:

```python
CUDA error: device-side assert triggered
/pytorch/aten/src/ATen/native/cuda/IndexKernel.cu:113: operator(): block: [0,0,0], thread: [0,0,0]
Assertion -sizes[i] <= index && index < sizes[i] && "index out of bounds" failed.
```

### Bug

In the `token_type_ids_mask_function` inner mask, the code correctly handles out-of-bounds `kv_idx` values but **fails to handle out-of-bounds `q_idx` values**.

The PR
- #39396 

originally fixed the bidirectional image masking by adding bounds checking for `kv_idx`, but overlooked that `q_idx` needed the same protection.

During generation with static caches:
- Cache shapes can exceed the actual input sequence length (e.g., static cache of 2048 positions with 512 token input)
- The masking function receives both `q_idx` and `kv_idx` that can exceed `token_type_ids.shape[1]`
- Direct indexing like `token_type_ids[batch_idx, q_idx]` causes CUDA index out of bounds errors when `q_idx >= token_type_ids.shape[1]`

The code comment on line 740 already acknowledged this issue:
> "NOTE: static cache shape goes beyond input seq length, while token_type_ids.shape[1] == input seq length"

Bounds checking was implemented for `kv_idx`, but `q_idx` was overlooked.

### Fix

This PR adds the same `torch.where` bounds-checking pattern for `q_idx` that already existed for `kv_idx`:
1. Create `safe_q_idx` to clamp indices within valid range
2. Use safe indices for tensor access
3. Apply `torch.where` to mask out-of-bounds values with appropriate defaults (0 for `token_type_ids`, -1 for `image_group_ids`)

### Affected Models
- `Gemma3ForConditionalGeneration`
- `PaliGemmaForConditionalGeneration`
- Example modular transformer template (`modeling_new_task_model.py`)

### Testing

This PR fixes the downstream failing test in TRL:
> tests/test_grpo_trainer.py::TestGRPOTrainer::test_training_vlm_0_trl_internal_testing_tiny_Gemma3ForConditionalGeneration

See associated issue:
- https://github.com/huggingface/trl/issues/4281

### Related Issues

- Regression exposed by: cache initialization refactor
  - #41505

- This PR completes the fix started in (which fixed `kv_idx` but missed `q_idx`):
  - #39396

- This PR will fix https://github.com/huggingface/trl/issues/4281

CC:
- @gante, who made the PR #41505
- @zucchini-nlp , who made the PR #39396

Comments:
unsubscribe

On Tue, Oct 21, 2025 at 12:32 PM github-actions[bot] <
***@***.***> wrote:

> *github-actions[bot]* left a comment (huggingface/transformers#41757)
> <https://github.com/huggingface/transformers/pull/41757#issuecomment-3425036912>
>
> *[For maintainers]* Suggested jobs to run (before merge)
>
> run-slow: gemma3
>
> —
> Reply to this email directly, view it on GitHub
> <https://github.com/huggingface/transformers/pull/41757#issuecomment-3425036912>,
> or unsubscribe
> <https://github.com/notifications/unsubscribe-auth/ATN5J2YTBP5G47TWRYZQEK33YXLA5AVCNFSM6AAAAACJYBPO2KVHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZTIMRVGAZTMOJRGI>
> .
> You are receiving this because you are subscribed to this thread.Message
> ID: ***@***.***>
>


The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_41757). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

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

run-slow: gemma3, paligemma

Thanks for your review, @zucchini-nlp.

I have run the tests as requested and everything is OK:
```python
pytest -k test_generate_with_static_cache tests/models/gemma3/test_modeling_gemma3.py

tests/models/gemma3/test_modeling_gemma3.py::Gemma3TextModelTest::test_generate_with_static_cache PASSED                                                                                            [ 50%]
tests/models/gemma3/test_modeling_gemma3.py::Gemma3Vision2TextModelTest::test_generate_with_static_cache PASSED                                                                                     [100%]

============================================================================== 2 passed, 373 deselected, 2 warnings in 9.62s ==============================================================================
```

Thanks

Thank YOU, @zucchini-nlp for your fast and efficient review!