Fix CUDA index out of bounds for q_idx in VLM token type masking for Gemma3, PaliGemma, and example modular - #41757

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

• 🚨 [v5] generate delegates default cache initialization to the model #41505

changed cache initialization behavior in generate(), a latent bug in the VLM masking code was exposed. The error manifests as:

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

• [gemma3] fix bidirectional image mask #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:

• CI fails with dev dependencies: torch.AcceleratorError: CUDA error: device-side assert triggered trl#4281

Related Issues

• Regression exposed by: cache initialization refactor

  • 🚨 [v5] generate delegates default cache initialization to the model #41505

• This PR completes the fix started in (which fixed kv_idx but missed q_idx):

  • [gemma3] fix bidirectional image mask #39396

• This PR will fix CI fails with dev dependencies: torch.AcceleratorError: CUDA error: device-side assert triggered trl#4281

CC:

• @gante, who made the PR 🚨 [v5] generate delegates default cache initialization to the model #41505
• @zucchini-nlp , who made the PR [gemma3] fix bidirectional image mask #39396