Skip to content

[Bug]: Gemma 4: Engine hang during large prefill caused by Interleaved Attention and p-RoPE implementation #39914

Description

@alexcardo

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:

  1. 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.
  2. 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.
  3. 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:

  1. Load Gemma 4 31B with a large max-model-len .
  2. Send a single request with a prompt exceeding ~ 4096 tokens.
  3. 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...

  • Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.

Metadata

Metadata

Assignees

No one assigned

    Labels

    bugSomething isn't working

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions