255 lines
11 KiB
Python
255 lines
11 KiB
Python
################################################################################
|
|
# Copyright(c)2020-2025 Shanghai Biren Technology Co., Ltd. All rights reserved.
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
#
|
|
################################################################################
|
|
from collections.abc import Iterable
|
|
from typing import Optional
|
|
|
|
import torch
|
|
import torch_br
|
|
from fastcore.basics import patch_to
|
|
from transformers import Qwen3Config
|
|
|
|
import vllm.model_executor.models.qwen3
|
|
from vllm.attention import AttentionType
|
|
from vllm.config import CacheConfig
|
|
from vllm.forward_context import ForwardContext, get_forward_context
|
|
from vllm.logger import logger
|
|
from vllm.model_executor.layers.layernorm import RMSNorm
|
|
from vllm.model_executor.layers.quantization import QuantizationConfig
|
|
from vllm.model_executor.layers.quantization.base_config import (
|
|
QuantizationConfig)
|
|
from vllm.model_executor.layers.rotary_embedding import MRotaryEmbedding
|
|
from vllm.model_executor.model_loader.weight_utils import (
|
|
default_weight_loader, maybe_remap_kv_scale_name)
|
|
from vllm.model_executor.models.qwen3 import (Qwen3Attention,
|
|
Qwen3DecoderLayer, Qwen3Model)
|
|
from vllm.model_executor.models.utils import is_pp_missing_parameter
|
|
from vllm_br.v1.attention.backends.attention_v1 import (
|
|
SUPAFlashAttentionMetadata)
|
|
from .qwen2 import model_forward
|
|
from .supa_module import MergedGateUpMLPSiluL2
|
|
|
|
|
|
@patch_to(vllm.model_executor.models.qwen3.Qwen3Attention)
|
|
def forward(
|
|
self,
|
|
positions: torch.Tensor,
|
|
hidden_states: torch.Tensor,
|
|
) -> torch.Tensor:
|
|
forward_context: ForwardContext = get_forward_context()
|
|
attn_metadata: SUPAFlashAttentionMetadata = forward_context.attn_metadata
|
|
if attn_metadata is None:
|
|
## for dummy run
|
|
return hidden_states
|
|
|
|
seq_len = hidden_states.shape[-2]
|
|
decode_seql = 512
|
|
|
|
if isinstance(attn_metadata, dict):
|
|
attn_metadata = attn_metadata[self.attn.layer_name]
|
|
kv_cache = self.attn.kv_cache[forward_context.virtual_engine]
|
|
if kv_cache is not None:
|
|
if seq_len <= decode_seql:
|
|
if hasattr(self.qkv_proj, "qweight"):
|
|
qkv_weight = self.qkv_proj.qweight.data
|
|
qkv_scales = self.qkv_proj.scales.data
|
|
elif hasattr(self.qkv_proj, "weight_packed"):
|
|
qkv_weight = self.qkv_proj.weight_packed.data
|
|
qkv_scales = self.qkv_proj.weight_scale.data
|
|
else:
|
|
qkv_weight = self.qkv_proj.weight
|
|
qkv_scales = None
|
|
if isinstance(self.rotary_emb, MRotaryEmbedding):
|
|
assert len(
|
|
self.rotary_emb.mrope_section
|
|
) == 3 and self.rotary_emb.mrope_section[
|
|
1] == self.rotary_emb.mrope_section[
|
|
2], "current only support mrope_section width and height are equal!"
|
|
q, k, v = torch_br.br_qwen3_vl_prefix_attn_infer(
|
|
hidden_states,
|
|
qkv_weight, [self.q_size, self.kv_size, self.kv_size],
|
|
self.head_dim,
|
|
self.q_norm.variance_epsilon,
|
|
self.q_norm.weight,
|
|
self.k_norm.weight,
|
|
self.rotary_emb.cos_sin_cache,
|
|
kv_cache,
|
|
positions,
|
|
attn_metadata.slot_mapping,
|
|
self.rotary_emb.mrope_section[1],
|
|
bias=self.qkv_proj.bias,
|
|
scales=qkv_scales)
|
|
else:
|
|
q, k, v = torch_br.br_qwen3_prefix_attn_infer(
|
|
hidden_states,
|
|
qkv_weight, [self.q_size, self.kv_size, self.kv_size],
|
|
self.head_dim,
|
|
self.q_norm.variance_epsilon,
|
|
self.q_norm.weight,
|
|
self.k_norm.weight,
|
|
self.rotary_emb.sin_cache,
|
|
self.rotary_emb.cos_cache,
|
|
kv_cache,
|
|
positions,
|
|
attn_metadata.slot_mapping,
|
|
bias=self.qkv_proj.bias,
|
|
scales=qkv_scales)
|
|
else:
|
|
qkv, _ = self.qkv_proj(hidden_states)
|
|
if isinstance(self.rotary_emb, MRotaryEmbedding):
|
|
assert len(
|
|
self.rotary_emb.mrope_section
|
|
) == 3 and self.rotary_emb.mrope_section[
|
|
1] == self.rotary_emb.mrope_section[
|
|
2], "current only support mrope_section width and height are equal!"
|
|
q, k, v = torch_br.br_fused_rms_mrope_kvstore_infer(
|
|
qkv, [self.q_size, self.kv_size, self.kv_size],
|
|
self.head_dim, self.q_norm.variance_epsilon,
|
|
self.q_norm.weight, self.k_norm.weight,
|
|
self.rotary_emb.cos_sin_cache, kv_cache, positions,
|
|
attn_metadata.slot_mapping, attn_metadata.block_table,
|
|
attn_metadata.query_start_loc, attn_metadata.context_lens,
|
|
self.rotary_emb.mrope_section[1])
|
|
else:
|
|
q, k, v = torch_br.br_fused_rms_rope_kvstore_infer(
|
|
qkv, [self.q_size, self.kv_size, self.kv_size],
|
|
self.head_dim, self.q_norm.variance_epsilon,
|
|
self.q_norm.weight, self.k_norm.weight,
|
|
self.rotary_emb.sin_cache, self.rotary_emb.cos_cache,
|
|
kv_cache, positions, attn_metadata.slot_mapping,
|
|
attn_metadata.block_table, attn_metadata.query_start_loc,
|
|
attn_metadata.context_lens)
|
|
if hasattr(attn_metadata, 'do_cache'):
|
|
attn_metadata.do_cache = False
|
|
attn_output = self.attn(q, k, v)
|
|
output, _ = self.o_proj(attn_output)
|
|
return output
|
|
else:
|
|
return hidden_states
|
|
|
|
|
|
def Qwen3DecoderLayer__init__(
|
|
self,
|
|
config: Qwen3Config,
|
|
cache_config: Optional[CacheConfig] = None,
|
|
quant_config: Optional[QuantizationConfig] = None,
|
|
prefix: str = "",
|
|
) -> None:
|
|
super(Qwen3DecoderLayer, self).__init__()
|
|
self.hidden_size = config.hidden_size
|
|
# Requires transformers > 4.32.0
|
|
rope_theta = getattr(config, "rope_theta", 1000000)
|
|
rope_scaling = getattr(config, "rope_scaling", None)
|
|
|
|
# By default, Qwen3 uses causal attention as it is a decoder-only model.
|
|
# You can override the HF config with `is_causal=False` to enable
|
|
# bidirectional attention, which is used in some embedding models
|
|
# (e.g. Alibaba-NLP/gte-Qwen3-7B-instruct)
|
|
if getattr(config, "is_causal", True):
|
|
attn_type = AttentionType.DECODER
|
|
else:
|
|
attn_type = AttentionType.ENCODER_ONLY
|
|
|
|
self.self_attn = Qwen3Attention(
|
|
hidden_size=self.hidden_size,
|
|
num_heads=config.num_attention_heads,
|
|
max_position=config.max_position_embeddings,
|
|
num_kv_heads=config.num_key_value_heads,
|
|
rope_theta=rope_theta,
|
|
rms_norm_eps=config.rms_norm_eps,
|
|
qkv_bias=getattr(config, 'attention_bias', False),
|
|
head_dim=getattr(config, 'head_dim', None),
|
|
cache_config=cache_config,
|
|
quant_config=quant_config,
|
|
rope_scaling=rope_scaling,
|
|
prefix=f"{prefix}.self_attn",
|
|
attn_type=attn_type,
|
|
)
|
|
self.mlp = MergedGateUpMLPSiluL2(
|
|
hidden_size=self.hidden_size,
|
|
intermediate_size=config.intermediate_size,
|
|
hidden_act=config.hidden_act,
|
|
quant_config=quant_config,
|
|
prefix=f"{prefix}.mlp",
|
|
)
|
|
self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
|
|
self.post_attention_layernorm = RMSNorm(config.hidden_size,
|
|
eps=config.rms_norm_eps)
|
|
|
|
|
|
def load_weights(self, weights: Iterable[tuple[str,
|
|
torch.Tensor]]) -> set[str]:
|
|
stacked_params_mapping = [
|
|
# (param_name, shard_name, shard_id)
|
|
("qkv_proj", "q_proj", "q"),
|
|
("qkv_proj", "k_proj", "k"),
|
|
("qkv_proj", "v_proj", "v"),
|
|
("gate_up_proj", "gate_proj", 0),
|
|
("gate_up_proj", "up_proj", 1),
|
|
]
|
|
params_dict = dict(self.named_parameters(remove_duplicate=False))
|
|
loaded_params: set[str] = set()
|
|
for name, loaded_weight in weights:
|
|
if "rotary_emb.inv_freq" in name:
|
|
continue
|
|
if (self.quant_config is not None
|
|
and (scale_name := self.quant_config.get_cache_scale(name))):
|
|
# Loading kv cache quantization scales
|
|
param = params_dict[scale_name]
|
|
weight_loader = getattr(param, "weight_loader",
|
|
default_weight_loader)
|
|
loaded_weight = (loaded_weight
|
|
if loaded_weight.dim() == 0 else loaded_weight[0])
|
|
weight_loader(param, loaded_weight)
|
|
loaded_params.add(scale_name)
|
|
continue
|
|
for (param_name, weight_name, shard_id) in stacked_params_mapping:
|
|
if weight_name not in name:
|
|
continue
|
|
name = name.replace(weight_name, param_name)
|
|
# Skip loading extra bias for GPTQ models.
|
|
if name.endswith(".bias") and name not in params_dict:
|
|
continue
|
|
if is_pp_missing_parameter(name, self):
|
|
continue
|
|
param = params_dict[name]
|
|
weight_loader = param.weight_loader
|
|
weight_loader(param, loaded_weight, shard_id)
|
|
break
|
|
else:
|
|
# Skip loading extra bias for GPTQ models.
|
|
if name.endswith(".bias") and name not in params_dict:
|
|
continue
|
|
# Remapping the name of FP8 kv-scale.
|
|
name = maybe_remap_kv_scale_name(name, params_dict)
|
|
if name is None:
|
|
continue
|
|
if is_pp_missing_parameter(name, self):
|
|
continue
|
|
param = params_dict[name]
|
|
weight_loader = getattr(param, "weight_loader",
|
|
default_weight_loader)
|
|
weight_loader(param, loaded_weight)
|
|
if name.find("norm.weight") != -1:
|
|
param.data = param.data.to(torch.float32)
|
|
loaded_params.add(name)
|
|
return loaded_params
|
|
|
|
|
|
vllm.model_executor.models.qwen3.Qwen3DecoderLayer.__init__ = Qwen3DecoderLayer__init__
|
|
logger.debug('[Patch] patch Qwen3 MLP with MergedGateUpMLPSiluL2')
|
|
Qwen3Model.load_weights = load_weights
|
|
Qwen3Model.forward = model_forward
|