vllm - 💡(How to fix) Fix [RFC]: Prefill Context Parallel for Qwen3.5 Hybrid Attention [3 comments, 3 participants]

Summary

This RFC describes the implementation of Prefill Context Parallel (PCP) for Qwen3.5's hybrid attention architecture, which combines full attention and linear attention (GatedDeltaNet) layers. PCP reduces long-context prefill latency by:

• Full attention layers: Distributing sequence tokens across ranks using zigzag ring attention (cp_size speedup per layer)
• Linear attention: Splitting batch dimension instead of sequence, eliminating replicated computation and all_gather overhead
• MoE layers: Spliting tokens dimension

Motivation

Long-context prefill is a major bottleneck for LLM inference:

1. Quadratic attention complexity: Full attention is O(N^2) in sequence length, making 256K+ contexts extremely slow
2. Memory bandwidth : Large KV caches stress memory subsystem

Tensor Parallelism (TP) helps by distributing model weights, but each GPU still processes the full sequence. Context Parallel (CP) addresses this by distributing the sequence itself across GPUs.

Challenge: Hybrid Attention

Qwen3.5 uses a hybrid architecture:

• 15 full attention layers: Standard transformer attention
• 45 linear attention layers: GatedDeltaNet (GDN) with recurrent state

Linear attention has a key constraint: the recurrent state at position is depends on ALL previous tokens. This means naive token splitting breaks correctness for linear attention layers.

Design

Architecture Overview

The complete picture showing data distribution and communication across all layer types:

Token Distribution

PCP uses zigzag distribution to balance tokens across ranks:

Input sequence: [0, 1, 2, 3, 4, 5, 6, 7, 8, ...]

With cp_size = C:

Rank 0: [0, C, C+1, 2C, 2C+1, ...]  (positions where pos % C in {0, C-1, ...})
Rank 1: [1, C-1, C+2, 2C-1, ...]  (positions where pos % C in {1, C-2, ...})
...
Rank r: [r, ...]  (positions where pos % C == r or similar pattern)

This interleaving ensures each rank has tokens from all parts of the sequence, improving load balance for causal attention.

For cp_size = 2 specifically:

Rank 0: [0, 3, 4, 7, 8, ...]  (positions 0, 3, 4, 7, ...)
Rank 1: [1, 2, 5, 6, 9, ...]  (positions 1, 2, 5, 6, ...)

Layer-Type Handling

Full Attention Layers (Ring Attention)

Full attention layers use zigzag ring attention:

┌──────────────────────────────────────────────────────────────────┐
│                     Ring Attention (Full Attn)                     │
│                                                                    │
│  Rank r                                        Rank (r+1) mod C   │
│  ┌────────────┐                               ┌────────────┐      │
│  │ Q (local)  │                               │ Q (local)  │      │
│  │ K (local)  │ ─────────── P2P send ───────► │ K (recv)   │      │
│  │ V (local)  │ ─────────── P2P send ────────► │ V (recv)   │      │
│  └────────────┘                               └────────────┘      │
│       │                                              │              │
│       ▼                                              ▼              │
│  Attention(Q_local, K_all, V_all)        Attention(Q_local, ...) │
│       │                                              │              │
│       └──────────────── LSE merge ──────────────────┘              │
└──────────────────────────────────────────────────────────────────┘
where C = cp_size (number of PCP ranks)

Each rank:

1. Computes local Q @ local K/V
2. Sends K/V to next rank via P2P
3. Receives K/V from previous rank
4. Accumulates attention outputs with LSE merge
5. Repeats for cp_size rounds

Benefit: Each rank has N/cp_size query tokens, but still attends to all N keys via ring exchange. Computation per rank: (N/cp_size) × N = N²/cp_size → cp_size speedup (not cp_size², because each Q still needs all K/V for causal).

Linear Attention Layers (Batch Split)

Linear attention requires the full sequence per item due to recurrent state dependency. We split by batch dimension instead of sequence

• Linear attention has no cross-item dependency: each item's recurrent state only depends on that item's history
• Each rank has B/C items, each with FULL sequence S

Key insight: Split by batch, not sequence. This gives each rank the full sequence it needs without any all_gather communication.

Work Plan

Each item below is scoped for a separate pull request.