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
[Qwen] Add fp8 checkpoint support for qwen3-next. (#25079)
Signed-off-by: Tao He <linzhu.ht@alibaba-inc.com>
Tao He committed
Sep 18, 2025 at 16:16 UTC
ef7eefe17a7dc212ddb8a8aabd7760218a10e25e
2 files changed
+22
-21
vllm/model_executor/models/qwen3_next.py
+17
-18
@@ -30,7 +30,6 @@ from vllm.model_executor.layers.layernorm import (
30
GemmaRMSNorm as Qwen3NextRMSNorm)
31
# yapf: enable
32
from vllm.model_executor.layers.linear import (ColumnParallelLinear,
33
- MergedColumnParallelLinear,
33
QKVParallelLinear,
34
ReplicatedLinear,
35
RowParallelLinear)
@@ -254,12 +253,20 @@ class Qwen3NextGatedDeltaNet(nn.Module, MambaBase):
253
# projection of the input hidden states
254
self.projection_size_qkvz = self.key_dim * 2 + self.value_dim * 2
255
self.projection_size_ba = self.num_v_heads * 2
257
- self.in_proj = MergedColumnParallelLinear(
256
+ self.in_proj_qkvz = ColumnParallelLinear(
257
input_size=self.hidden_size,
259
- output_sizes=[self.projection_size_qkvz, self.projection_size_ba],
258
+ output_size=self.projection_size_qkvz,
259
bias=False,
260
quant_config=quant_config,
262
- prefix=f"{prefix}.in_proj",
261
+ prefix=f"{prefix}.in_proj_qkvz",
262
+ )
263
+ # ba_proj doesn't support blockwise fp8 quantization.
264
+ self.in_proj_ba = ColumnParallelLinear(
265
+ input_size=self.hidden_size,
266
+ output_size=self.projection_size_ba,
267
+ bias=False,
268
+ quant_config=quant_config,
269
+ prefix=f"{prefix}.in_proj_ba",
270
)
271
272
query_key_settings = (self.key_dim, 0, False)
@@ -420,19 +427,14 @@ class Qwen3NextGatedDeltaNet(nn.Module, MambaBase):
427
ssm_state = self_kv_cache[1]
428
num_actual_tokens = attn_metadata.num_actual_tokens
429
num_accepted_tokens = attn_metadata.num_accepted_tokens
423
-
424
- # 1. Set up dimensions for reshapes later
425
- projected_states, _ = self.in_proj(hidden_states[:num_actual_tokens])
430
if spec_token_masks is not None:
431
spec_token_masks = spec_token_masks[:num_actual_tokens]
428
- projected_states_qkvz, projected_states_ba = torch.split(
429
- projected_states,
430
- [
431
- self.projection_size_qkvz // self.tp_size,
432
- self.projection_size_ba // self.tp_size
433
- ],
434
- dim=-1,
435
- )
432
+
433
+ # 1. Set up dimensions for reshapes later
434
+ projected_states_qkvz, _ = self.in_proj_qkvz(
435
+ hidden_states[:num_actual_tokens])
436
+ projected_states_ba, _ = self.in_proj_ba(
437
+ hidden_states[:num_actual_tokens])
438
query, key, value, z, b, a = self.fix_query_key_value_ordering(
439
projected_states_qkvz, projected_states_ba)
440
query, key, value = map(lambda x: rearrange(x, 'l p d -> l (p d)'),
@@ -976,8 +978,6 @@ class Qwen3NextModel(nn.Module):
978
("qkv_proj", "v_proj", "v"),
979
("gate_up_proj", "gate_proj", 0),
980
("gate_up_proj", "up_proj", 1),
979
- ("in_proj", "in_proj_qkvz", 0),
980
- ("in_proj", "in_proj_ba", 1),
981
]
982
983
params_dict = dict(self.named_parameters())
@@ -1055,7 +1055,6 @@ class Qwen3NextForCausalLM(nn.Module, HasInnerState, SupportsLoRA, SupportsPP,
1055
"v_proj",
1056
],
1057
"gate_up_proj": ["gate_proj", "up_proj"],
1058
- "in_proj": ["in_proj_qkvz", "in_proj_ba"],
1058
}
1059
1060
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
vllm/model_executor/models/qwen3_next_mtp.py
+5
-3
@@ -63,7 +63,9 @@ class Qwen3NextMultiTokenPredictor(nn.Module):
63
self.config.hidden_size,
64
gather_output=True,
65
bias=False,
66
- return_bias=False)
66
+ return_bias=False,
67
+ quant_config=quant_config,
68
+ prefix=f'{prefix}.fc')
69
70
self.layers = torch.nn.ModuleList(
71
Qwen3NextDecoderLayer(
@@ -72,7 +74,7 @@ class Qwen3NextMultiTokenPredictor(nn.Module):
74
model_config=model_config,
75
cache_config=cache_config,
76
quant_config=quant_config,
75
- prefix=f'{prefix}.layers.{self.mtp_start_layer_idx + idx}',
77
+ prefix=f'{prefix}.layers.{idx}',
78
) for idx in range(self.num_mtp_layers))
79
80
self.make_empty_intermediate_tensors = (
@@ -233,7 +235,7 @@ class Qwen3NextMTP(nn.Module, SupportsPP):
235
self.config = config
236
self.model = Qwen3NextMultiTokenPredictor(vllm_config=vllm_config,
237
prefix=maybe_prefix(
236
- prefix, "model"))
238
+ prefix, "mtp"))
239
self.unpadded_vocab_size = config.vocab_size
240
self.lm_head = ParallelLMHead(self.unpadded_vocab_size,
241
config.hidden_size,