`generate_batch` failing
Author: qgallouedec
Pull request: no
Labels: name: bug

### System Info

- `transformers` version: 4.57.0.dev0 (cf084f5b40e19b5a5f946cee75bead6d4247b071)
- Platform: Linux-5.15.0-1048-aws-x86_64-with-glibc2.31
- Python version: 3.12.11
- Huggingface_hub version: 0.34.4
- Safetensors version: 0.5.3
- Accelerate version: 1.10.1
- Accelerate config:    not found
- DeepSpeed version: 0.17.4
- PyTorch version (accelerator?): 2.7.1+cu128 (CUDA)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using distributed or parallel set-up in script?: <fill in>
- Using GPU in script?: <fill in>
- GPU type: NVIDIA H100 80GB HBM3

### Who can help?

@remi-or @McPatate 

### Information

- [ ] The official example scripts
- [ ] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

```python
from transformers import AutoModelForCausalLM, GenerationConfig

model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2-0.5B-Instruct", device_map="cuda")

input_ids = [
    [29038, 787, 4103],
    [29038, 787, 4103],
]

model.generate_batch(input_ids, generation_config=GenerationConfig())
```

```
`eos_token_id` not set in GenerationConfig. Setting to -1 (disabled).
num_blocks = 42646 is too large, setting to self._upper_bound_num_blocks = 4096
max_batch_tokens = 13646 is too large, setting to self._upper_bound_max_batch_tokens = 256
PagedAttentionCache initialized with self.num_blocks = 4096, self.block_size = 32, page_size = 128, self.max_batch_tokens = 256 num_attention_masks = 1
self.cache_shape = (131073, 2, 64) self.key_cache[0].shape = torch.Size([131073, 2, 64]) self.key_cache[0].numel() = 16777344                                  
Error in generation loop: CUDA error: operation failed due to a previous error during capture                                                                  
CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.
Traceback (most recent call last):
  File "/fsx/qgallouedec/transformers/src/transformers/generation/continuous_batching/continuous_api.py", line 686, in warmup
    self._generation_step(batch_processor)
  File "/fsx/qgallouedec/transformers/src/transformers/generation/continuous_batching/continuous_api.py", line 692, in _generation_step
    batch_data = batch_processor.get_model_kwargs()
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/fsx/qgallouedec/transformers/src/transformers/generation/continuous_batching/continuous_api.py", line 239, in get_model_kwargs
    kwargs["attention_mask"][layer_type] = self.attention_mask[:1, :, :t, : seqlens_k[-1]]
                                           ~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: CUDA error: operation not permitted when stream is capturing
CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.


During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/fsx/qgallouedec/transformers/src/transformers/generation/continuous_batching/continuous_api.py", line 781, in _run_generation_loop
    self._inner_generation_loop(batch_processor)
  File "/fsx/qgallouedec/transformers/src/transformers/generation/continuous_batching/continuous_api.py", line 801, in _inner_generation_loop
    self.warmup(batch_processor)
  File "/fsx/qgallouedec/transformers/src/transformers/generation/continuous_batching/continuous_api.py", line 685, in warmup
    with torch.cuda.graph(self.graph, stream=stream):
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/fsx/qgallouedec/miniconda3/envs/trl/lib/python3.12/site-packages/torch/cuda/graphs.py", line 186, in __exit__
    self.cuda_graph.capture_end()
  File "/fsx/qgallouedec/miniconda3/envs/trl/lib/python3.12/site-packages/torch/cuda/graphs.py", line 84, in capture_end
    super().capture_end()
RuntimeError: CUDA error: operation failed due to a previous error during capture
CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.

Generation loop finished.                                                                                                                                      
Generation thread terminated unexpectedly.                                                                                                                     
Solving 2 requests:   0%|                                                                                                           | 0/2 [00:01<?, ?request/s]
```
The bug was introduced in #40688

I was able to track back the issue to this specific commit https://github.com/remi-or/transformers/commit/0fe77e8617f99d5c2d653ae8180223a36d99aead

### Expected behavior

To work

Comments:
Indeed, CB is no longer compatible with cuda graphs since the last changes, which is an issue. I am adding an actual exception to #40828 and restoring cuda graphs are the next PR on my continuous batching roadmap. Will link the PR once it is up.

Hi @remi-or, thanks for clarifying in this thread about CUDA Graphs support. You mentioned that bringing it back is on your roadmap — is there anywhere we could follow along with that roadmap or the expected timeline? Totally understand priorities may shift, but it would be helpful to know if/when this might land. Thanks for all the great work you’re doing!

Hey @mxchampagne , sorry for the late answer but I actually opened a PR about this yesterday. It might be merged next week !

Hey @remi-or, thanks for the follow up!

PR is merged. Let me know if I can close this @qgallouedec @mxchampagne  :)