A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
[Qwen3] Enable dual-chunk-attention support for Qwen3 models. (#21924)
Signed-off-by: Tao He <linzhu.ht@alibaba-inc.com>
Tao He committed
Aug 7, 2025 at 10:58 UTC
7377131a2ccb49cae71aa503ee5be520aa080904
2 files changed
+60
-31
vllm/model_executor/models/qwen3.py
+40
-24
@@ -23,7 +23,7 @@
23
# limitations under the License.
24
"""Inference-only Qwen3 model compatible with HuggingFace weights."""
25
from collections.abc import Iterable
26
-from typing import Optional, Union
26
+from typing import Any, Optional, Union
27
28
import torch
29
from torch import nn
@@ -47,27 +47,31 @@ from vllm.sequence import IntermediateTensors
47
from .interfaces import SupportsLoRA, SupportsPP
48
from .qwen2 import Qwen2MLP as Qwen3MLP
49
from .qwen2 import Qwen2Model
50
-from .utils import AutoWeightsLoader, PPMissingLayer, maybe_prefix
50
+from .utils import (AutoWeightsLoader, PPMissingLayer, extract_layer_index,
51
+ maybe_prefix)
52
53
logger = init_logger(__name__)
54
55
56
class Qwen3Attention(nn.Module):
57
57
- def __init__(self,
58
- hidden_size: int,
59
- num_heads: int,
60
- num_kv_heads: int,
61
- max_position: int = 4096 * 32,
62
- head_dim: Optional[int] = None,
63
- rms_norm_eps: float = 1e-06,
64
- qkv_bias: bool = False,
65
- rope_theta: float = 10000,
66
- cache_config: Optional[CacheConfig] = None,
67
- quant_config: Optional[QuantizationConfig] = None,
68
- rope_scaling: Optional[tuple] = None,
69
- prefix: str = "",
70
- attn_type: str = AttentionType.DECODER) -> None:
58
+ def __init__(
59
+ self,
60
+ hidden_size: int,
61
+ num_heads: int,
62
+ num_kv_heads: int,
63
+ max_position: int = 4096 * 32,
64
+ head_dim: Optional[int] = None,
65
+ rms_norm_eps: float = 1e-06,
66
+ qkv_bias: bool = False,
67
+ rope_theta: float = 10000,
68
+ cache_config: Optional[CacheConfig] = None,
69
+ quant_config: Optional[QuantizationConfig] = None,
70
+ rope_scaling: Optional[tuple] = None,
71
+ prefix: str = "",
72
+ attn_type: str = AttentionType.DECODER,
73
+ dual_chunk_attention_config: Optional[dict[str, Any]] = None,
74
+ ) -> None:
75
super().__init__()
76
self.hidden_size = hidden_size
77
tp_size = get_tensor_model_parallel_world_size()
@@ -89,6 +93,7 @@ class Qwen3Attention(nn.Module):
93
self.kv_size = self.num_kv_heads * self.head_dim
94
self.scaling = self.head_dim**-0.5
95
self.rope_theta = rope_theta
96
+ self.dual_chunk_attention_config = dual_chunk_attention_config
97
98
self.qkv_proj = QKVParallelLinear(
99
hidden_size,
@@ -113,15 +118,22 @@ class Qwen3Attention(nn.Module):
118
max_position=max_position,
119
base=self.rope_theta,
120
rope_scaling=rope_scaling,
121
+ dual_chunk_attention_config=dual_chunk_attention_config,
122
+ )
123
+ self.attn = Attention(
124
+ self.num_heads,
125
+ self.head_dim,
126
+ self.scaling,
127
+ num_kv_heads=self.num_kv_heads,
128
+ cache_config=cache_config,
129
+ quant_config=quant_config,
130
+ prefix=f"{prefix}.attn",
131
+ attn_type=attn_type,
132
+ **{
133
+ "layer_idx": extract_layer_index(prefix),
134
+ "dual_chunk_attention_config": dual_chunk_attention_config,
135
+ } if dual_chunk_attention_config else {},
136
)
117
- self.attn = Attention(self.num_heads,
118
- self.head_dim,
119
- self.scaling,
120
- num_kv_heads=self.num_kv_heads,
121
- cache_config=cache_config,
122
- quant_config=quant_config,
123
- prefix=f"{prefix}.attn",
124
- attn_type=attn_type)
137
self.q_norm = RMSNorm(self.head_dim, eps=rms_norm_eps)
138
self.k_norm = RMSNorm(self.head_dim, eps=rms_norm_eps)
139
@@ -161,6 +173,9 @@ class Qwen3DecoderLayer(nn.Module):
173
# Requires transformers > 4.32.0
174
rope_theta = getattr(config, "rope_theta", 1000000)
175
rope_scaling = getattr(config, "rope_scaling", None)
176
+ dual_chunk_attention_config = getattr(config,
177
+ "dual_chunk_attention_config",
178
+ None)
179
180
# By default, Qwen3 uses causal attention as it is a decoder-only model.
181
# You can override the HF config with `is_causal=False` to enable
@@ -185,6 +200,7 @@ class Qwen3DecoderLayer(nn.Module):
200
rope_scaling=rope_scaling,
201
prefix=f"{prefix}.self_attn",
202
attn_type=attn_type,
203
+ dual_chunk_attention_config=dual_chunk_attention_config,
204
)
205
self.mlp = Qwen3MLP(
206
hidden_size=self.hidden_size,
vllm/model_executor/models/qwen3_moe.py
+20
-7
@@ -185,6 +185,7 @@ class Qwen3MoeAttention(nn.Module):
185
cache_config: Optional[CacheConfig] = None,
186
quant_config: Optional[QuantizationConfig] = None,
187
prefix: str = "",
188
+ dual_chunk_attention_config: Optional[dict[str, Any]] = None,
189
) -> None:
190
super().__init__()
191
self.hidden_size = hidden_size
@@ -208,6 +209,7 @@ class Qwen3MoeAttention(nn.Module):
209
self.scaling = self.head_dim**-0.5
210
self.rope_theta = rope_theta
211
self.max_position_embeddings = max_position_embeddings
212
+ self.dual_chunk_attention_config = dual_chunk_attention_config
213
214
self.qkv_proj = QKVParallelLinear(hidden_size,
215
self.head_dim,
@@ -229,14 +231,21 @@ class Qwen3MoeAttention(nn.Module):
231
max_position=max_position_embeddings,
232
base=rope_theta,
233
rope_scaling=rope_scaling,
234
+ dual_chunk_attention_config=dual_chunk_attention_config,
235
+ )
236
+ self.attn = Attention(
237
+ self.num_heads,
238
+ self.head_dim,
239
+ self.scaling,
240
+ num_kv_heads=self.num_kv_heads,
241
+ cache_config=cache_config,
242
+ quant_config=quant_config,
243
+ prefix=f"{prefix}.attn",
244
+ **{
245
+ "layer_idx": extract_layer_index(prefix),
246
+ "dual_chunk_attention_config": dual_chunk_attention_config,
247
+ } if dual_chunk_attention_config else {},
248
)
233
- self.attn = Attention(self.num_heads,
234
- self.head_dim,
235
- self.scaling,
236
- num_kv_heads=self.num_kv_heads,
237
- cache_config=cache_config,
238
- quant_config=quant_config,
239
- prefix=f"{prefix}.attn")
249
250
self.q_norm = RMSNorm(self.head_dim, eps=rms_norm_eps)
251
self.k_norm = RMSNorm(self.head_dim, eps=rms_norm_eps)
@@ -280,6 +289,9 @@ class Qwen3MoeDecoderLayer(nn.Module):
289
rope_scaling = getattr(config, "rope_scaling", None)
290
max_position_embeddings = getattr(config, "max_position_embeddings",
291
8192)
292
+ dual_chunk_attention_config = getattr(config,
293
+ "dual_chunk_attention_config",
294
+ None)
295
self.self_attn = Qwen3MoeAttention(
296
hidden_size=self.hidden_size,
297
num_heads=config.num_attention_heads,
@@ -293,6 +305,7 @@ class Qwen3MoeDecoderLayer(nn.Module):
305
cache_config=cache_config,
306
quant_config=quant_config,
307
prefix=f"{prefix}.self_attn",
308
+ dual_chunk_attention_config=dual_chunk_attention_config,
309
)
310
311
# `mlp_only_layers` in the config.