Merge pull request #2436 from Datta0/qwen3_support
Qwen3 inference fixes
This commit is contained in:
Michael Han
GitHub
commit
a8f4c7dac0
2 changed files with 368 additions and 95 deletions
|
|
@ -911,98 +911,104 @@ pass
|
|||
|
||||
|
||||
# https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py#L825
|
||||
def LlamaModel_fast_forward_inference(
|
||||
self,
|
||||
input_ids,
|
||||
past_key_values,
|
||||
position_ids,
|
||||
attention_mask = None,
|
||||
):
|
||||
input_ids = input_ids[:,:self.max_seq_length]
|
||||
bsz, q_len = input_ids.shape
|
||||
hd = self.config.hidden_size
|
||||
mlp_size = self.config.intermediate_size
|
||||
def _LlamaModel_fast_forward_inference(attention_fast_forward_inference=LlamaAttention_fast_forward_inference, mlp_fast_forward_inference=fast_swiglu_inference):
|
||||
# This makes the attention and MLP customisable.
|
||||
# Now for models like qwen3 or cohere which use custom attention operations, we can use this function
|
||||
def LlamaModel_fast_forward_inference_custom(
|
||||
self,
|
||||
input_ids,
|
||||
past_key_values,
|
||||
position_ids,
|
||||
attention_mask = None,
|
||||
):
|
||||
input_ids = input_ids[:,:self.max_seq_length]
|
||||
bsz, q_len = input_ids.shape
|
||||
hd = self.config.hidden_size
|
||||
mlp_size = self.config.intermediate_size
|
||||
|
||||
X = self.model.embed_tokens(input_ids)
|
||||
X = X.to(_get_dtype(self.config.torch_dtype))
|
||||
bsz, q_len, hd = X.shape
|
||||
assert(q_len == 1)
|
||||
# Get saved buffers to reduce memory movement
|
||||
residual = torch.empty((bsz, q_len, hd), dtype = torch.float32, device = "cuda:0")
|
||||
_XX = torch.empty((2, bsz, q_len, hd), dtype = torch.float32, device = "cuda:0")
|
||||
XX, XX2 = _XX[0], _XX[1]
|
||||
variance = torch.empty((bsz, q_len, 1), dtype = torch.float32, device = "cuda:0")
|
||||
temp_mlp = torch.empty((2, bsz, 1, mlp_size), dtype = X.dtype, device = "cuda:0")
|
||||
temp_gate, temp_up = temp_mlp[0], temp_mlp[1]
|
||||
X = self.model.embed_tokens(input_ids)
|
||||
X = X.to(_get_dtype(self.config.torch_dtype))
|
||||
bsz, q_len, hd = X.shape
|
||||
assert(q_len == 1)
|
||||
# Get saved buffers to reduce memory movement
|
||||
residual = torch.empty((bsz, q_len, hd), dtype = torch.float32, device = "cuda:0")
|
||||
_XX = torch.empty((2, bsz, q_len, hd), dtype = torch.float32, device = "cuda:0")
|
||||
XX, XX2 = _XX[0], _XX[1]
|
||||
variance = torch.empty((bsz, q_len, 1), dtype = torch.float32, device = "cuda:0")
|
||||
temp_mlp = torch.empty((2, bsz, 1, mlp_size), dtype = X.dtype, device = "cuda:0")
|
||||
temp_gate, temp_up = temp_mlp[0], temp_mlp[1]
|
||||
|
||||
seq_len = past_key_values[0][0].shape[-2]
|
||||
if bsz != 1:
|
||||
attention_mask = _prepare_4d_causal_attention_mask_for_sdpa(
|
||||
attention_mask,
|
||||
(bsz, q_len),
|
||||
X,
|
||||
seq_len,
|
||||
sliding_window = getattr(self.config, "sliding_window", None),
|
||||
)
|
||||
else:
|
||||
attention_mask = None
|
||||
pass
|
||||
seq_len = past_key_values[0][0].shape[-2]
|
||||
if bsz != 1:
|
||||
attention_mask = _prepare_4d_causal_attention_mask_for_sdpa(
|
||||
attention_mask,
|
||||
(bsz, q_len),
|
||||
X,
|
||||
seq_len,
|
||||
sliding_window = getattr(self.config, "sliding_window", None),
|
||||
)
|
||||
else:
|
||||
attention_mask = None
|
||||
pass
|
||||
|
||||
next_decoder_cache = []
|
||||
next_decoder_cache = []
|
||||
|
||||
for idx, decoder_layer in enumerate(self.model.layers):
|
||||
residual.copy_(X) # residual = X
|
||||
for idx, decoder_layer in enumerate(self.model.layers):
|
||||
residual.copy_(X) # residual = X
|
||||
X = fast_rms_layernorm_inference(
|
||||
decoder_layer.input_layernorm,
|
||||
X,
|
||||
XX = XX,
|
||||
XX2 = XX2,
|
||||
variance = variance,
|
||||
)
|
||||
X, present_key_value = attention_fast_forward_inference(
|
||||
decoder_layer.self_attn,
|
||||
hidden_states = X,
|
||||
past_key_value = past_key_values[idx],
|
||||
position_ids = position_ids,
|
||||
attention_mask = attention_mask,
|
||||
do_prefill = not hasattr(decoder_layer.self_attn, "paged_attention"),
|
||||
)
|
||||
X += residual
|
||||
|
||||
residual.copy_(X) # residual = X
|
||||
X = fast_rms_layernorm_inference(
|
||||
decoder_layer.post_attention_layernorm,
|
||||
X,
|
||||
XX = XX,
|
||||
XX2 = XX2,
|
||||
variance = variance,
|
||||
)
|
||||
X = mlp_fast_forward_inference(
|
||||
decoder_layer.mlp,
|
||||
X,
|
||||
temp_gate = temp_gate,
|
||||
temp_up = temp_up,
|
||||
)
|
||||
X += residual
|
||||
|
||||
next_decoder_cache.append(present_key_value)
|
||||
pass
|
||||
X = fast_rms_layernorm_inference(
|
||||
decoder_layer.input_layernorm,
|
||||
self.model.norm,
|
||||
X,
|
||||
XX = XX,
|
||||
XX2 = XX2,
|
||||
variance = variance,
|
||||
)
|
||||
X, present_key_value = LlamaAttention_fast_forward_inference(
|
||||
decoder_layer.self_attn,
|
||||
hidden_states = X,
|
||||
past_key_value = past_key_values[idx],
|
||||
position_ids = position_ids,
|
||||
attention_mask = attention_mask,
|
||||
do_prefill = not hasattr(decoder_layer.self_attn, "paged_attention"),
|
||||
)
|
||||
X += residual
|
||||
|
||||
residual.copy_(X) # residual = X
|
||||
X = fast_rms_layernorm_inference(
|
||||
decoder_layer.post_attention_layernorm,
|
||||
X,
|
||||
XX = XX,
|
||||
XX2 = XX2,
|
||||
variance = variance,
|
||||
return BaseModelOutputWithPast(
|
||||
last_hidden_state = X,
|
||||
past_key_values = next_decoder_cache,
|
||||
hidden_states = [],
|
||||
attentions = [],
|
||||
)
|
||||
X = fast_swiglu_inference(
|
||||
decoder_layer.mlp,
|
||||
X,
|
||||
temp_gate = temp_gate,
|
||||
temp_up = temp_up,
|
||||
)
|
||||
X += residual
|
||||
|
||||
next_decoder_cache.append(present_key_value)
|
||||
pass
|
||||
X = fast_rms_layernorm_inference(
|
||||
self.model.norm,
|
||||
X,
|
||||
XX = XX,
|
||||
XX2 = XX2,
|
||||
variance = variance,
|
||||
)
|
||||
|
||||
return BaseModelOutputWithPast(
|
||||
last_hidden_state = X,
|
||||
past_key_values = next_decoder_cache,
|
||||
hidden_states = [],
|
||||
attentions = [],
|
||||
)
|
||||
pass
|
||||
return LlamaModel_fast_forward_inference_custom
|
||||
|
||||
# For ensuring backwards compatibility, we create LlamaModel_fast_forward_inference that is consumed by other models
|
||||
LlamaModel_fast_forward_inference = _LlamaModel_fast_forward_inference()
|
||||
|
||||
def CausalLM_fast_forward(fast_forward_inference):
|
||||
def _CausalLM_fast_forward(
|
||||
|
|
|
|||
|
|
@ -18,6 +18,7 @@ from ._utils import __version__
|
|||
from .llama import (
|
||||
LlamaRotaryEmbedding,
|
||||
LlamaLinearScalingRotaryEmbedding,
|
||||
_LlamaModel_fast_forward_inference,
|
||||
)
|
||||
try:
|
||||
from transformers.models.qwen3.modeling_qwen3 import (
|
||||
|
|
@ -37,7 +38,9 @@ except:
|
|||
f"to obtain the latest transformers build, then restart this session."\
|
||||
)
|
||||
pass
|
||||
|
||||
from transformers.modeling_attn_mask_utils import (
|
||||
_prepare_4d_causal_attention_mask_for_sdpa,
|
||||
)
|
||||
# For Pytorch 2.1.1
|
||||
try:
|
||||
from transformers.models.qwen3.modeling_qwen3 import (
|
||||
|
|
@ -103,17 +106,19 @@ def Qwen3Attention_fast_forward(
|
|||
if past_key_value is not None:
|
||||
kv_seq_len += past_key_value[0].shape[-2]
|
||||
|
||||
# Extend RoPE dynamically to fit in VRAM
|
||||
self.rotary_emb.extend_rope_embedding(V, seq_len = kv_seq_len)
|
||||
|
||||
if position_ids is None:
|
||||
cos = self.rotary_emb.cos_cached
|
||||
sin = self.rotary_emb.sin_cached
|
||||
Q, K = fast_rope_embedding(Q, K, cos, sin)
|
||||
if position_embeddings:
|
||||
cos, sin = position_embeddings
|
||||
else:
|
||||
cos, sin = self.rotary_emb(V, seq_len = kv_seq_len)
|
||||
Q, K = inplace_rope_embedding(Q, K, cos, sin, position_ids)
|
||||
pass
|
||||
# Extend RoPE dynamically to fit in VRA
|
||||
rotary_emb = self.rotary_emb
|
||||
rotary_emb.extend_rope_embedding(V, seq_len = kv_seq_len)
|
||||
|
||||
if position_ids is None:
|
||||
# Useful for LongRoPE
|
||||
cos, sin = rotary_emb.get_cached(kv_seq_len)
|
||||
else:
|
||||
cos, sin = rotary_emb(V, seq_len = kv_seq_len)
|
||||
Q, K = fast_rope_embedding(Q, K, cos, sin)
|
||||
|
||||
if past_key_value is not None:
|
||||
K = torch.cat([past_key_value[0], K], dim = 2)
|
||||
|
|
@ -164,8 +169,7 @@ def Qwen3Attention_fast_forward(
|
|||
Q = Q.transpose(1, 2)
|
||||
K = K.transpose(1, 2)
|
||||
V = V.transpose(1, 2)
|
||||
sw = getattr(self.config, "sliding_window", None)
|
||||
sw = kv_seq_len if (sw is None or sw == "null") else sw
|
||||
sw = kv_seq_len
|
||||
window = (-1, -1) if (kv_seq_len <= sw) else (sw, sw)
|
||||
A = flash_attn_func(Q, K, V, causal = True, window_size = window)
|
||||
else:
|
||||
|
|
@ -185,13 +189,276 @@ def Qwen3Attention_fast_forward(
|
|||
# Go back to (batch_size, seq_len, n_heads, head_dim)
|
||||
A = A.transpose(1, 2).contiguous()
|
||||
pass
|
||||
|
||||
|
||||
attn_output = A.reshape(bsz, q_len, n_heads*head_dim)
|
||||
attn_output = self.apply_o(self, attn_output)
|
||||
attn_weights = None
|
||||
return attn_output, attn_weights, past_key_value
|
||||
pass
|
||||
|
||||
torch_matmul = torch.matmul
|
||||
def Qwen3Attention_fast_forward_inference(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
past_key_value: Optional[Tuple[torch.Tensor]],
|
||||
position_ids,
|
||||
do_prefill = False,
|
||||
attention_mask = None,
|
||||
):
|
||||
"""
|
||||
https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py#L406
|
||||
Fast inference using KV cache.
|
||||
QK^T can be computed in 4 chunks
|
||||
|
||||
[Q, q] @ [K, k].T where q, k are the new tokens.
|
||||
[QK^T, Qk^T]
|
||||
[qK^T, qk^T]
|
||||
|
||||
Since the attention mask wipes Qk^T, we just get
|
||||
[QK^T, 0]
|
||||
[qK^T, qk^T]
|
||||
|
||||
Since softmax is row-wise, we get
|
||||
softmax([QK^T, 0])
|
||||
softmax([qK^T, qk^T])
|
||||
|
||||
We then multiply by [V]
|
||||
[v]
|
||||
softmax([QK^T, 0]) [softmax(QK^T)V] *
|
||||
softmax([qK^T, qk^T]) [softmax([qK^T, qk^T]) @ [V, v]]
|
||||
|
||||
But notice * [softmax(QK^T)V] is just the last attention.
|
||||
We just need to compute the last final row.
|
||||
|
||||
This means we can pass in a row of Q, but we need to
|
||||
remember K and V, which are called the KV cache.
|
||||
"""
|
||||
Xn = hidden_states
|
||||
bsz, _, hd = hidden_states.size()
|
||||
K1, V1 = past_key_value
|
||||
dtype = Xn.dtype
|
||||
|
||||
n_heads = self.config.num_attention_heads
|
||||
n_groups = self.num_key_value_groups
|
||||
n_kv_heads = self.config.num_key_value_heads
|
||||
head_dim = self.head_dim
|
||||
# assert(n_kv_heads * n_groups == n_heads)
|
||||
|
||||
hidden_size = self.config.hidden_size
|
||||
attention_size = n_heads*head_dim
|
||||
seq_len = K1.shape[-2]
|
||||
kv_seq_len = seq_len + 1
|
||||
|
||||
# Prefill phase
|
||||
# if not hasattr(self, "paged_attention"):
|
||||
device = hidden_states.device
|
||||
if do_prefill:
|
||||
self.paged_attention = torch.empty((KV_CACHE_INCREMENT+seq_len+1, 2, bsz, n_kv_heads, head_dim), dtype = dtype, device = device)
|
||||
self.paged_attention_K = self.paged_attention[:,0]
|
||||
self.paged_attention_V = self.paged_attention[:,1]
|
||||
self.paged_attention_K[:seq_len] = K1.permute(2, 0, 1, 3)
|
||||
self.paged_attention_V[:seq_len] = V1.permute(2, 0, 1, 3)
|
||||
self.temp_QA = torch.empty((2, bsz, 1, attention_size), dtype = dtype, device = device)
|
||||
self.temp_KV = torch.empty((2, bsz, 1, n_kv_heads*head_dim), dtype = dtype, device = device)
|
||||
self.RH_Q = torch.empty((bsz, n_heads, 1, head_dim), dtype = dtype, device = device)
|
||||
|
||||
# Mistral Nemo 12b has weird dimensions
|
||||
if attention_size != hidden_size:
|
||||
self.temp_O = torch.empty((1, bsz, hidden_size), dtype = dtype, device = device)
|
||||
else:
|
||||
self.temp_O = self.temp_QA[1][:,:,:hidden_size]
|
||||
pass
|
||||
|
||||
self.attention = torch.empty((bsz, n_heads, 1, KV_CACHE_INCREMENT+seq_len), dtype = dtype, device = device)
|
||||
self.scalar = 1.0 / math_sqrt(self.head_dim)
|
||||
self.half_head_dim = head_dim // 2
|
||||
elif kv_seq_len >= self.paged_attention.shape[0]:
|
||||
self.paged_attention.resize_((self.paged_attention.shape[0]+KV_CACHE_INCREMENT, 2, bsz, n_kv_heads, head_dim))
|
||||
self.paged_attention_K = self.paged_attention[:,0]
|
||||
self.paged_attention_V = self.paged_attention[:,1]
|
||||
self.attention.resize_((bsz, n_heads, 1, self.attention.shape[-1]+KV_CACHE_INCREMENT))
|
||||
pass
|
||||
|
||||
Qn = fast_linear_forward(self.q_proj, Xn, out = self.temp_QA[0])
|
||||
Kn = fast_linear_forward(self.k_proj, Xn, out = self.temp_KV[0])
|
||||
Vn = fast_linear_forward(self.v_proj, Xn, out = self.temp_KV[1])
|
||||
Qn = Qn.view(bsz, 1, n_heads, head_dim)#.transpose(1, 2) # we will transpose after normalisation
|
||||
Kn = Kn.view(bsz, 1, n_kv_heads, head_dim)#.transpose(1, 2) # we will transpose after normalisation
|
||||
Vn = Vn.view(bsz, 1, n_kv_heads, head_dim).transpose(1, 2)
|
||||
|
||||
Qn = fast_rms_layernorm(self.q_norm, Qn)
|
||||
Kn = fast_rms_layernorm(self.k_norm, Kn)
|
||||
|
||||
Qn = Qn.transpose(1, 2)
|
||||
Kn = Kn.transpose(1, 2)
|
||||
|
||||
# cos, sin = self.rotary_emb(Vn, seq_len = kv_seq_len)
|
||||
# Qn, Kn = inplace_rope_embedding(Qn, Kn, cos, sin, position_ids)
|
||||
|
||||
# Need to do it prior 2 steps before hitting full on short KV cache
|
||||
# or else error
|
||||
self.rotary_emb.extend_rope_embedding(Vn, seq_len + 2)
|
||||
cos, sin = self.rotary_emb.get_cached(kv_seq_len)
|
||||
cos = cos[position_ids].unsqueeze(1)
|
||||
sin = sin[position_ids].unsqueeze(1)
|
||||
h = self.half_head_dim
|
||||
|
||||
RH_Q = self.RH_Q
|
||||
RH_Q[:,:,:,:h] = Qn[:,:,:,h:]
|
||||
RH_Q[:,:,:,h:] = Qn[:,:,:,:h]
|
||||
RH_Q[:,:,:,:h].neg_() # torch.neg(RH_Q[:,:,:,:h], out = RH_Q[:,:,:,:h])
|
||||
Qn *= cos
|
||||
Qn.addcmul_(RH_Q, sin)
|
||||
|
||||
RH_K = RH_Q[:,:n_kv_heads,:,:] # torch.empty((n_kv_heads, 1, head_dim), dtype = dtype, device = "cuda:0")
|
||||
RH_K[:,:,:,:h] = Kn[:,:,:,h:]
|
||||
RH_K[:,:,:,h:] = Kn[:,:,:,:h]
|
||||
RH_K[:,:,:,:h].neg_() #torch.neg(RH_K[:,:,:,:h], out = RH_K[:,:,:,:h])
|
||||
Kn *= cos
|
||||
Kn.addcmul_(RH_K, sin)
|
||||
|
||||
# New KV cache
|
||||
# Kn = torch.cat([K1, Kn], dim = 2)
|
||||
# Vn = torch.cat([V1, Vn], dim = 2)
|
||||
self.paged_attention_K[seq_len] = Kn.permute(2, 0, 1, 3)
|
||||
self.paged_attention_V[seq_len] = Vn.permute(2, 0, 1, 3)
|
||||
Kn = self.paged_attention_K[:kv_seq_len].permute(1, 2, 0, 3)
|
||||
Vn = self.paged_attention_V[:kv_seq_len].permute(1, 2, 0, 3)
|
||||
|
||||
# Handle sliding windows
|
||||
sliding_window = getattr(self.config, "sliding_window", None)
|
||||
if sliding_window is not None and kv_seq_len > sliding_window:
|
||||
# From https://github.com/huggingface/transformers/blob/main/src/transformers/models/mistral/modeling_mistral.py#L193
|
||||
slicing_tokens = 1 - sliding_window
|
||||
Knn = Kn[:, :, slicing_tokens:, :]#.contiguous()
|
||||
Vnn = Vn[:, :, slicing_tokens:, :]#.contiguous()
|
||||
else:
|
||||
Knn, Vnn = Kn, Vn
|
||||
pass
|
||||
|
||||
# Grouped query attention
|
||||
_, _, cached_len, _ = Knn.shape
|
||||
if bsz == 1 or not SDPA_HAS_GQA and n_groups != 1:
|
||||
Knn = Knn[:, :, None, :, :].expand(bsz, n_kv_heads, n_groups, cached_len, head_dim)
|
||||
Vnn = Vnn[:, :, None, :, :].expand(bsz, n_kv_heads, n_groups, cached_len, head_dim)
|
||||
Knn = Knn.reshape(bsz, n_heads, cached_len, head_dim)
|
||||
Vnn = Vnn.reshape(bsz, n_heads, cached_len, head_dim)
|
||||
pass
|
||||
# else:
|
||||
# Knn, Vnn = Knn, Vnn
|
||||
# pass
|
||||
|
||||
# Attention
|
||||
if bsz == 1:
|
||||
Qn *= self.scalar # See https://github.com/ggerganov/llama.cpp/issues/7805#issuecomment-2153349963
|
||||
# It seems like doing (Q * scalar) @ K is better than (Q @ K) * scalar to stop overflows
|
||||
A = torch_matmul(Qn, Knn.transpose(2, 3), out = self.attention[:,:,:,:cached_len])
|
||||
# if attention_mask is not None: A += attention_mask # Must add attention_mask for batched
|
||||
A[:] = torch_nn_functional_softmax(A, dim = -1, dtype = torch.float32)#.to(A.dtype)
|
||||
A = torch_matmul(A, Vnn, out = Qn)
|
||||
else:
|
||||
if SDPA_HAS_GQA:
|
||||
A = scaled_dot_product_attention(Qn, Knn, Vnn, attn_mask = attention_mask, is_causal = False, enable_gqa = True)
|
||||
else:
|
||||
A = scaled_dot_product_attention(Qn, Knn, Vnn, attn_mask = attention_mask, is_causal = False)
|
||||
pass
|
||||
A = A.transpose(1, 2)
|
||||
A = A.reshape(bsz, 1, attention_size)
|
||||
A = fast_linear_forward(self.o_proj, A, out = self.temp_O)
|
||||
return A, (Kn, Vn)
|
||||
pass
|
||||
|
||||
# def Qwen3Model_fast_forward_inference(
|
||||
# self,
|
||||
# input_ids,
|
||||
# past_key_values,
|
||||
# position_ids,
|
||||
# attention_mask = None,
|
||||
# ):
|
||||
# input_ids = input_ids[:,:self.max_seq_length]
|
||||
# bsz, q_len = input_ids.shape
|
||||
# hd = self.config.hidden_size
|
||||
# mlp_size = self.config.intermediate_size
|
||||
|
||||
# X = self.model.embed_tokens(input_ids)
|
||||
# X = X.to(_get_dtype(self.config.torch_dtype))
|
||||
# bsz, q_len, hd = X.shape
|
||||
# assert(q_len == 1)
|
||||
# # Get saved buffers to reduce memory movement
|
||||
# residual = torch.empty((bsz, q_len, hd), dtype = torch.float32, device = "cuda:0")
|
||||
# _XX = torch.empty((2, bsz, q_len, hd), dtype = torch.float32, device = "cuda:0")
|
||||
# XX, XX2 = _XX[0], _XX[1]
|
||||
# variance = torch.empty((bsz, q_len, 1), dtype = torch.float32, device = "cuda:0")
|
||||
# temp_mlp = torch.empty((2, bsz, 1, mlp_size), dtype = X.dtype, device = "cuda:0")
|
||||
# temp_gate, temp_up = temp_mlp[0], temp_mlp[1]
|
||||
|
||||
# seq_len = past_key_values[0][0].shape[-2]
|
||||
# if bsz != 1:
|
||||
# attention_mask = _prepare_4d_causal_attention_mask_for_sdpa(
|
||||
# attention_mask,
|
||||
# (bsz, q_len),
|
||||
# X,
|
||||
# seq_len,
|
||||
# sliding_window = getattr(self.config, "sliding_window", None),
|
||||
# )
|
||||
# else:
|
||||
# attention_mask = None
|
||||
# pass
|
||||
|
||||
# next_decoder_cache = []
|
||||
|
||||
# for idx, decoder_layer in enumerate(self.model.layers):
|
||||
# residual.copy_(X) # residual = X
|
||||
# X = fast_rms_layernorm_inference(
|
||||
# decoder_layer.input_layernorm,
|
||||
# X,
|
||||
# XX = XX,
|
||||
# XX2 = XX2,
|
||||
# variance = variance,
|
||||
# )
|
||||
# X, present_key_value = Qwen3Attention_fast_forward_inference(
|
||||
# decoder_layer.self_attn,
|
||||
# hidden_states = X,
|
||||
# past_key_value = past_key_values[idx],
|
||||
# position_ids = position_ids,
|
||||
# attention_mask = attention_mask,
|
||||
# do_prefill = not hasattr(decoder_layer.self_attn, "paged_attention"),
|
||||
# )
|
||||
# X += residual
|
||||
|
||||
# residual.copy_(X) # residual = X
|
||||
# X = fast_rms_layernorm_inference(
|
||||
# decoder_layer.post_attention_layernorm,
|
||||
# X,
|
||||
# XX = XX,
|
||||
# XX2 = XX2,
|
||||
# variance = variance,
|
||||
# )
|
||||
# X = fast_swiglu_inference(
|
||||
# decoder_layer.mlp,
|
||||
# X,
|
||||
# temp_gate = temp_gate,
|
||||
# temp_up = temp_up,
|
||||
# )
|
||||
# X += residual
|
||||
|
||||
# next_decoder_cache.append(present_key_value)
|
||||
# pass
|
||||
# X = fast_rms_layernorm_inference(
|
||||
# self.model.norm,
|
||||
# X,
|
||||
# XX = XX,
|
||||
# XX2 = XX2,
|
||||
# variance = variance,
|
||||
# )
|
||||
|
||||
# return BaseModelOutputWithPast(
|
||||
# last_hidden_state = X,
|
||||
# past_key_values = next_decoder_cache,
|
||||
# hidden_states = [],
|
||||
# attentions = [],
|
||||
# )
|
||||
# pass
|
||||
|
||||
class FastQwen3Model(FastLlamaModel):
|
||||
|
||||
|
|
@ -212,7 +479,7 @@ class FastQwen3Model(FastLlamaModel):
|
|||
Qwen3FlashAttention2.forward = Qwen3Attention_fast_forward
|
||||
Qwen3DecoderLayer .forward = LlamaDecoderLayer_fast_forward
|
||||
Qwen3Model .forward = LlamaModel_fast_forward
|
||||
Qwen3ForCausalLM .forward = CausalLM_fast_forward(LlamaModel_fast_forward_inference)
|
||||
Qwen3ForCausalLM .forward = CausalLM_fast_forward(_LlamaModel_fast_forward_inference(Qwen3Attention_fast_forward_inference))
|
||||
PeftModelForCausalLM.forward = PeftModelForCausalLM_fast_forward
|
||||
fix_prepare_inputs_for_generation(Qwen3ForCausalLM)
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue