ROCm: Flash/mem-efficient SDPA hangs during batched autoregressive generation with left-padded sequences
- Dominant language
- Python
- Stars
- 103k
- Forks
- 29.5k
- PR merge metrics
- PR metrics pending
Description
## Summary
Flash and memory-efficient SDPA backends hang indefinitely during batched autoregressive generation (`model.generate()`) when input sequences are left-padded to equal length. Disabling both backends and falling back to the math SDPA kernel resolves the hang. Single-sequence generation (no padding) works correctly with all backends.
## Environment
- **PyTorch:** 2.12.0+rocm7.2
- **ROCm:** 7.2.0 (HIP 7.2.53211)
- **GPU:** AMD Radeon RX 7900 XTX (gfx1100, 24 GB VRAM)
- **OS:** Ubuntu 24.04, kernel 6.8.0-117-generic
- **Python:** 3.10.18
- **transformers:** 4.57.3
## Reproduction
Reproducible with the [Qwen3-TTS](https://huggingface.co/Qwen/Qwen3-TTS-12Hz-1.7B-Base) model using the `qwen_tts` library (v0.1.1). The model uses SDPA attention (via HuggingFace transformers) and supports batched inference with left-padding.
```python
from qwen_tts import Qwen3TTSModel
import torchaudio
model = Qwen3TTSModel.from_pretrained(
"Qwen/Qwen3-TTS-12Hz-1.7B-Base",
device_map="cuda:0",
)
ref_audio, sr = torchaudio.load("reference.wav") # any 16kHz mono WAV
ref_text = "Transcription of the reference audio."
# Single generation — works fine with all SDPA backends
audio_single = model.generate_voice_clone(
text="This is a test sentence.",
ref_audio=ref_audio,
ref_text=ref_text,
)
# Batch generation — hangs indefinitely (left-pads shorter sequences)
audio_batch = model.generate_voice_clone(
text=["Short text.", "This is a significantly longer sentence that will cause left-padding of the shorter text."],
ref_audio=ref_audio,
ref_text=ref_text,
)
```
### What the model does internally
When `text` is a list, [`Qwen3TTSModel.generate()`](https://github.com/QwenLM/Qwen3-TTS/blob/main/qwen_tts/core/models/modeling_qwen3_tts.py) builds per-sequence `talker_input_embeds`, then left-pads them to equal length with `torch.nn.utils.rnn.pad_sequence(..., padding_value=0.0)`. It creates a proper 2D attention mask (0 for padded positions, 1 for real) and passes both to `self.talker.generate()` — the standard HuggingFace `GenerationMixin.generate()` autoregressive loop.
The hang occurs inside this autoregressive loop. The model never emits EOS and generation runs indefinitely at high GPU utilization. There is no crash, no error, no timeout — it simply never returns.
## Workaround
Disabling flash and memory-efficient SDPA backends before batch generation forces the math backend, which handles the left-padded batched generation correctly:
```python
import torch
# Save original state
flash_was = torch.backends.cuda.flash_sdp_enabled()
mem_was = torch.backends.cuda.mem_efficient_sdp_enabled()
# Disable for batch generation
torch.backends.cuda.enable_flash_sdp(False)
torch.backends.cuda.enable_mem_efficient_sdp(False)
audio_batch = model.generate_voice_clone(
text=["Short text.", "A much longer sentence requiring left-padding."],
ref_audio=ref_audio,
ref_text=ref_text,
)
# Restore
torch.backends.cuda.enable_flash_sdp(flash_was)
torch.backends.cuda.enable_mem_efficient_sdp(mem_was)
```
## Observations
- **Single-sequence generation** (no left-padding): Works with all SDPA backends.
- **Batched generation** (left-padded): Hangs with flash and/or mem-efficient SDPA enabled. Works with math-only SDPA.
- The hang is **deterministic** — same inputs always hang.
- The hang occurs regardless of batch size (tested with batch sizes 2, 5, 10, 50).
- The hang occurs regardless of text length — both short and long texts hang in batch mode.
- **NVIDIA users are not affected** — this appears specific to ROCm's SDPA kernel implementations.
- Previously worked with PyTorch + ROCm 6.3. The regression appeared after upgrading to ROCm 7.2 / PyTorch 2.12.0+rocm7.2.
## Notes
- A standalone minimal reproducer (without the model) using raw `F.scaled_dot_product_attention` with left-padded tensors and explicit masks does *not* reproduce the hang in isolation. The issue manifests specifically during the iterative KV-cache-based autoregressive generation loop, suggesting it may involve the interaction between the SDPA kernels and growing KV cache with padded positions over many steps.
- The model architecture uses GQA (16 Q-heads, 2 KV-heads, head_dim=64) which may be relevant to the kernel path selection.
cc @jeffdaily @sunway513 @jithunnair-amd @pruthvistony @ROCmSupport @jataylo @hongxiayang @naromero77amd @pragupta @jerrymannil @xinyazhang @drisspg @liangel-02 @howardzhang-cv @include-yy
Contributor guide
Assessment
This issue has not been assessed yet.