Your current environment
Gemma 4 engine completely freezes during the prefill phase when the prompt length exceeds ~4,000 tokens. This occurs due to a suspected failure in the Interleaved Attention/p-RoPE implementation, while incremental token generation remains stable.
🐛 Describe the bug
Environment:
- Model:
google/gemma-4-31b-it (and quantized versions)
- Hardware: NVIDIA RTX 5090 / RTX Pro 6000 Blackwell
- CUDA: 13.0
- Observed behavior: Total engine freeze during prefill of sequences exceeding ~4,000 tokens.
Technical Architecture of Gemma 4:
The model utilizes a complex hybrid attention mechanism that is not fully supported by the current vLLM scheduler/kernels:
- Interleaved Attention: The architecture consists of a mixture of layers: 5 local attention layers (with a sliding window of 1024 tokens for the 31B version) for every 1 global attention layer.
- Dual RoPE: The model employs two different rotary positional embeddings:
- Standard RoPE for local sliding-window layers.
- Proportional RoPE (p-RoPE) for global attention layers to enable context extension up to 256K.
- Shared KV Cache: The model implements a shared KV cache where subsequent layers reuse key and value projections from previous ones.
Observed Issues:
1. Prefill Hang vs. Incremental Growth:
- Incremental Scenario: When the context is built gradually (single stream, token-by-token), the engine remains stable and can handle 200K+ tokens without issues.
- Prefill Scenario: When a large prompt (>~4096 tokens) is sent in a single request, the engine hangs during the prefill phase. The scheduler/router stops processing, and the system enters a deadlock state.
2. Scheduler/Kernel Conflict:
The hang occurs specifically when the prefill process hits the global attention layers. It appears that the transition from local sliding-window attention to global attention (utilizing p-RoPE) fails during large-batch prefill. The engine is unable to handle the position offsets and the shared KV cache projections in a single prefill pass, leading to a complete freeze.
Steps to Reproduce:
- Load Gemma 4 31B with a large
max-model-len .
- Send a single request with a prompt exceeding ~ 4096 tokens.
- Observe the engine freeze during the prefill stage.
Expected Behavior:
The prefill process should correctly handle the interleaved structure (Local $\rightarrow$ Global) and the p-RoPE calculations without deadlocking the scheduler, regardless of whether the context is built incrementally or via a large initial prompt.
Before submitting a new issue...
Your current environment
Gemma 4 engine completely freezes during the prefill phase when the prompt length exceeds ~4,000 tokens. This occurs due to a suspected failure in the Interleaved Attention/p-RoPE implementation, while incremental token generation remains stable.
🐛 Describe the bug
Environment:
google/gemma-4-31b-it(and quantized versions)Technical Architecture of Gemma 4:
The model utilizes a complex hybrid attention mechanism that is not fully supported by the current vLLM scheduler/kernels:
Observed Issues:
1. Prefill Hang vs. Incremental Growth:
2. Scheduler/Kernel Conflict:
The hang occurs specifically when the prefill process hits the global attention layers. It appears that the transition from local sliding-window attention to global attention (utilizing p-RoPE) fails during large-batch prefill. The engine is unable to handle the position offsets and the shared KV cache projections in a single prefill pass, leading to a complete freeze.
Steps to Reproduce:
max-model-len.Expected Behavior:$\rightarrow$ Global) and the p-RoPE calculations without deadlocking the scheduler, regardless of whether the context is built incrementally or via a large initial prompt.
The prefill process should correctly handle the interleaved structure (Local
Before submitting a new issue...