feat(CRITICAL): 从 GitHub 扫描搬运 ixformer SDK + xllm 完整 GDN/MoE 代码
来源:
1. Chranos/ixformer (GitHub) → ixformer_sdk/ (230 files, 70K lines)
- inference/functions/vllm.py: vllm_moe_topk_softmax 完整实现 (2033 lines)
- inference/functions/moe.py: MoE ops 完整实现 (1380 lines)
- contrib/vllm_flash_attn/: FA2 Python 接口 (1018 lines)
- contrib/tgi/fused_moe.py: TGI fused MoE (429 lines)
- csrc/include/ixformer/: C++ kernel headers + cmake
2. Deep-Spark/xllm (GitHub) → upstream_ref/xllm_latest/ (+15 files)
- npu_torch/qwen3_5_decoder_layer_impl.cpp/.h
- npu_torch/qwen3_5_gated_delta_net.cpp/.h
- npu_torch/qwen3_next_*.cpp/.h (6 files)
- npu_torch/attention.cpp/.h + fused_moe.cpp/.h + CMakeLists.txt
- models/llm/qwen3_5.h + qwen3_5_mtp.h + qwen3_next.h
- models/vlm/qwen3_5.h
调用链完整性:
ixformer_sdk/inference/functions/vllm.py
→ ops.infer.moe_topk_softmax() (C++ 层)
→ 这就是 base 镜像 libixformer.so 里的实现
upstream_ref/xllm_latest/core/layers/ilu/fused_moe.cpp
→ ixformer::infer::topk_softmax() (直接 C++ 调用)
→ ixformer::infer::group_gemm() → 完整 7-step MoE pipeline
This commit is contained in:
project6-dev
186
ixformer_sdk/train/speedformer/layers/llama/attention.py
Normal file
186
ixformer_sdk/train/speedformer/layers/llama/attention.py
Normal file
@@ -0,0 +1,186 @@
|
||||
import math
|
||||
import warnings
|
||||
from typing import List, Optional, Tuple, Union
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from ixformer.train.speedformer.models.llama.configuration_llama import LlamaConfig
|
||||
from ixformer.train.speedformer.models.llama.modeling_llama import LlamaFlashAttention2
|
||||
from transformers import Cache
|
||||
from transformers.utils import logging
|
||||
|
||||
from flash_attn import flash_attn_func, flash_attn_varlen_func
|
||||
from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input
|
||||
|
||||
from ixformer.train.functions.fused_rope import fused_apply_rotary_pos_emb
|
||||
from ixformer.train.speedformer.layers.rotary_pos_embedding import RotaryEmbedding
|
||||
|
||||
|
||||
class BaseLlamaAttention(LlamaFlashAttention2):
|
||||
"""
|
||||
加这个层的原因：1.当原模型中使用的是torch nvtive的attention，强制替换成flash_attn; 2.优化rope
|
||||
"""
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
if self.config.rope_scaling is None:
|
||||
self.rotary_emb = RotaryEmbedding(self.head_dim)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
attention_mask: Optional[torch.LongTensor] = None,
|
||||
position_ids: Optional[torch.LongTensor] = None,
|
||||
past_key_value: Optional[Cache] = None,
|
||||
output_attentions: bool = False,
|
||||
use_cache: bool = False,
|
||||
**kwargs,
|
||||
) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
|
||||
output_attentions = False
|
||||
bsz, q_len, _ = hidden_states.size()
|
||||
query_states = self.q_proj(hidden_states)
|
||||
key_states = self.k_proj(hidden_states)
|
||||
value_states = self.v_proj(hidden_states)
|
||||
|
||||
# fused_apply_rotary_pos_emb need qk to be in "sbhd"
|
||||
query_states = query_states.view(
|
||||
bsz, q_len, self.num_heads, self.head_dim).transpose(1, 0).contiguous()
|
||||
key_states = key_states.view(
|
||||
bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 0).contiguous()
|
||||
value_states = value_states.view(
|
||||
bsz, q_len, self.num_heads, self.head_dim)
|
||||
|
||||
kv_seq_len = key_states.shape[0]
|
||||
if past_key_value is not None:
|
||||
kv_seq_len += past_key_value[0].shape[0]
|
||||
|
||||
emb = self.rotary_emb(kv_seq_len).to(dtype=torch.float32)
|
||||
query_states = fused_apply_rotary_pos_emb(query_states, emb)
|
||||
key_states = fused_apply_rotary_pos_emb(key_states, emb)
|
||||
|
||||
# kv cache staff
|
||||
if past_key_value is not None:
|
||||
# reuse k, v, self_attention
|
||||
key_states = torch.cat([past_key_value[0], key_states], dim=0)
|
||||
value_states = torch.cat([past_key_value[1], value_states], dim=0)
|
||||
past_key_value = (key_states, value_states) if use_cache else None
|
||||
|
||||
dropout_rate = self.attention_dropout if self.training else 0.0
|
||||
|
||||
# after fused_apply_rotary_pos_emb, qk change to "bshd" for flashattn or "bhsd" for sdpa
|
||||
if attention_mask is None: # flash-attn
|
||||
query_states = query_states.transpose(0, 1).contiguous()
|
||||
key_states = key_states.transpose(0, 1).contiguous()
|
||||
else: # sdpa
|
||||
query_states = query_states.permute(1, 2, 0, 3).contiguous()
|
||||
key_states = key_states.permute(1, 2, 0, 3).contiguous()
|
||||
value_states = value_states.transpose(1, 2).contiguous()
|
||||
# In PEFT, usually we cast the layer norms in float32 for training stability reasons
|
||||
# therefore the input hidden states gets silently casted in float32. Hence, we need
|
||||
# cast them back in the correct dtype just to be sure everything works as expected.
|
||||
# This might slowdown training & inference so it is recommended to not cast the LayerNorms
|
||||
# in fp32. (LlamaRMSNorm handles it correctly)
|
||||
|
||||
input_dtype = query_states.dtype
|
||||
if input_dtype == torch.float32:
|
||||
# Handle the case where the model is quantized
|
||||
if hasattr(self.config, "_pre_quantization_dtype"):
|
||||
target_dtype = self.config._pre_quantization_dtype
|
||||
else:
|
||||
target_dtype = self.q_proj.weight.dtype
|
||||
|
||||
query_states = query_states.to(target_dtype)
|
||||
key_states = key_states.to(target_dtype)
|
||||
value_states = value_states.to(target_dtype)
|
||||
|
||||
attn_output = self._flash_attention_forward(
|
||||
query_states, key_states, value_states, attention_mask, q_len, dropout=dropout_rate
|
||||
)
|
||||
|
||||
attn_output = attn_output.reshape(
|
||||
bsz, q_len, self.hidden_size).contiguous()
|
||||
attn_output = self.o_proj(attn_output)
|
||||
|
||||
if not output_attentions:
|
||||
attn_weights = None
|
||||
|
||||
return attn_output, attn_weights, past_key_value
|
||||
|
||||
def _flash_attention_forward(
|
||||
self, query_states, key_states, value_states, attention_mask, query_length, dropout=0.0, softmax_scale=None
|
||||
):
|
||||
"""
|
||||
for now, if attention_mask is none, flash-attn has better performance than torch.nn.functional.scaled_dot_product_attention;
|
||||
if attention_mask is not none, torch.nn.functional.scaled_dot_product_attention works better
|
||||
so sdpa and flash-attn is perfered according to attention_mask
|
||||
|
||||
Args:
|
||||
query_states (`torch.Tensor`):
|
||||
Input query states to be passed to Flash Attention API
|
||||
key_states (`torch.Tensor`):
|
||||
Input key states to be passed to Flash Attention API
|
||||
value_states (`torch.Tensor`):
|
||||
Input value states to be passed to Flash Attention API
|
||||
attention_mask (`torch.Tensor`):
|
||||
The padding mask - corresponds to a tensor of size `(batch_size, seq_len)` where 0 stands for the
|
||||
position of padding tokens and 1 for the position of non-padding tokens.
|
||||
dropout (`int`, *optional*):
|
||||
Attention dropout
|
||||
softmax_scale (`float`, *optional*):
|
||||
The scaling of QK^T before applying softmax. Default to 1 / sqrt(head_dim)
|
||||
"""
|
||||
# Contains at least one padding token in the sequence
|
||||
# if attention_mask is not None:
|
||||
if attention_mask is not None:
|
||||
attn_output = torch.nn.functional.scaled_dot_product_attention(
|
||||
query_states,
|
||||
key_states,
|
||||
value_states,
|
||||
attn_mask=attention_mask,
|
||||
dropout_p=self.attention_dropout if self.training else 0.0,
|
||||
# The q_len > 1 is necessary to match with AttentionMaskConverter.to_causal_4d that does not create a causal mask in case q_len == 1.
|
||||
is_causal=self.is_causal and attention_mask is None and query_length > 1,
|
||||
)
|
||||
attn_output = attn_output.transpose(1, 2).contiguous()
|
||||
|
||||
else:
|
||||
attn_output = flash_attn_func(
|
||||
query_states, key_states, value_states, dropout, softmax_scale=softmax_scale, causal=self.is_causal
|
||||
)
|
||||
|
||||
return attn_output
|
||||
|
||||
|
||||
class LlamaAttention(BaseLlamaAttention):
|
||||
def __init__(self) -> None:
|
||||
raise NotImplementedError(
|
||||
"LlamaAttention is not implemented as a physical class. "
|
||||
"It is meant to be used only with the from_native_module interface to Convert a native LlamaAttention module to LlamaAttention module provided above."
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def from_native_module(module: nn.Module, *args, **kwargs) -> nn.Module:
|
||||
|
||||
# LazyInitContext.materialize(module)
|
||||
|
||||
# try to get normalized_shape, eps, elementwise_affine from the module
|
||||
config = getattr(module, "config")
|
||||
layer_idx = getattr(module, "layer_idx", None)
|
||||
|
||||
attention = BaseLlamaAttention(
|
||||
config=config,
|
||||
layer_idx=layer_idx,
|
||||
)
|
||||
|
||||
attention.q_proj.weight = module.q_proj.weight
|
||||
attention.k_proj.weight = module.k_proj.weight
|
||||
attention.v_proj.weight = module.v_proj.weight
|
||||
attention.o_proj.weight = module.o_proj.weight
|
||||
|
||||
if config.attention_bias:
|
||||
attention.q_proj.bias = module.q_proj.bias
|
||||
attention.k_proj.bias = module.k_proj.bias
|
||||
attention.v_proj.bias = module.v_proj.bias
|
||||
attention.o_proj.bias = module.o_proj.bias
|
||||
return attention
|
||||
224
ixformer_sdk/train/speedformer/layers/llama/llama_method.py
Normal file
224
ixformer_sdk/train/speedformer/layers/llama/llama_method.py
Normal file
@@ -0,0 +1,224 @@
|
||||
import math
|
||||
import warnings
|
||||
from typing import List, Optional, Tuple, Union
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from ixformer.train.speedformer.models.llama.modeling_llama import LlamaModel
|
||||
from ixformer.train.speedformer.models.llama.modeling_attn_mask_utils import _prepare_4d_causal_attention_mask_for_sdpa
|
||||
from ixformer.train.speedformer.layers.cross_entropy_loss import fast_cross_entropy_loss as CrossEntropyLoss
|
||||
from transformers.utils import logging
|
||||
from transformers.cache_utils import Cache, DynamicCache
|
||||
|
||||
logger = logging.get_logger(__name__)
|
||||
|
||||
|
||||
def LlamaModel_forward():
|
||||
from transformers.modeling_outputs import BaseModelOutputWithPast
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids: torch.LongTensor = None,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
position_ids: Optional[torch.LongTensor] = None,
|
||||
past_key_values: Optional[List[torch.FloatTensor]] = None,
|
||||
inputs_embeds: Optional[torch.FloatTensor] = None,
|
||||
use_cache: Optional[bool] = None,
|
||||
output_attentions: Optional[bool] = None,
|
||||
output_hidden_states: Optional[bool] = None,
|
||||
return_dict: Optional[bool] = None,
|
||||
) -> Union[Tuple, BaseModelOutputWithPast]:
|
||||
output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
|
||||
output_hidden_states = (
|
||||
output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
|
||||
)
|
||||
use_cache = use_cache if use_cache is not None else self.config.use_cache
|
||||
|
||||
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
|
||||
|
||||
# retrieve input_ids and inputs_embeds
|
||||
if input_ids is not None and inputs_embeds is not None:
|
||||
raise ValueError(
|
||||
"You cannot specify both input_ids and inputs_embeds at the same time")
|
||||
elif input_ids is not None:
|
||||
batch_size, seq_length = input_ids.shape[:2]
|
||||
elif inputs_embeds is not None:
|
||||
batch_size, seq_length = inputs_embeds.shape[:2]
|
||||
else:
|
||||
raise ValueError(
|
||||
"You have to specify either input_ids or inputs_embeds")
|
||||
|
||||
if self.gradient_checkpointing and self.training:
|
||||
if use_cache:
|
||||
logger.warning_once(
|
||||
"`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
|
||||
)
|
||||
use_cache = False
|
||||
|
||||
past_key_values_length = 0
|
||||
if use_cache:
|
||||
use_legacy_cache = not isinstance(past_key_values, Cache)
|
||||
if use_legacy_cache:
|
||||
past_key_values = DynamicCache.from_legacy_cache(
|
||||
past_key_values)
|
||||
past_key_values_length = past_key_values.get_usable_length(
|
||||
seq_length)
|
||||
|
||||
if position_ids is None:
|
||||
device = input_ids.device if input_ids is not None else inputs_embeds.device
|
||||
position_ids = torch.arange(
|
||||
past_key_values_length, seq_length + past_key_values_length, dtype=torch.long, device=device
|
||||
)
|
||||
position_ids = position_ids.unsqueeze(0)
|
||||
|
||||
if inputs_embeds is None:
|
||||
inputs_embeds = self.embed_tokens(input_ids)
|
||||
|
||||
if attention_mask is not None:
|
||||
# output_attentions=True can not be supported when using SDPA, and we fall back on
|
||||
# the manual implementation that requires a 4D causal mask in all cases.
|
||||
attention_mask = _prepare_4d_causal_attention_mask_for_sdpa(
|
||||
attention_mask,
|
||||
(batch_size, seq_length),
|
||||
inputs_embeds,
|
||||
past_key_values_length,
|
||||
)
|
||||
|
||||
# embed positions
|
||||
hidden_states = inputs_embeds
|
||||
|
||||
# decoder layers
|
||||
all_hidden_states = () if output_hidden_states else None
|
||||
all_self_attns = () if output_attentions else None
|
||||
next_decoder_cache = None
|
||||
|
||||
for decoder_layer in self.layers:
|
||||
if output_hidden_states:
|
||||
all_hidden_states += (hidden_states,)
|
||||
|
||||
if self.gradient_checkpointing and self.training:
|
||||
layer_outputs = self._gradient_checkpointing_func(
|
||||
decoder_layer.__call__,
|
||||
hidden_states,
|
||||
attention_mask,
|
||||
position_ids,
|
||||
past_key_values,
|
||||
output_attentions,
|
||||
use_cache,
|
||||
)
|
||||
else:
|
||||
layer_outputs = decoder_layer(
|
||||
hidden_states,
|
||||
attention_mask=attention_mask,
|
||||
position_ids=position_ids,
|
||||
past_key_value=past_key_values,
|
||||
output_attentions=output_attentions,
|
||||
use_cache=use_cache,
|
||||
)
|
||||
|
||||
hidden_states = layer_outputs[0]
|
||||
|
||||
if use_cache:
|
||||
next_decoder_cache = layer_outputs[2 if output_attentions else 1]
|
||||
|
||||
if output_attentions:
|
||||
all_self_attns += (layer_outputs[1],)
|
||||
|
||||
hidden_states = self.norm(hidden_states)
|
||||
|
||||
# add hidden states from the last decoder layer
|
||||
if output_hidden_states:
|
||||
all_hidden_states += (hidden_states,)
|
||||
|
||||
next_cache = None
|
||||
if use_cache:
|
||||
next_cache = next_decoder_cache.to_legacy_cache(
|
||||
) if use_legacy_cache else next_decoder_cache
|
||||
if not return_dict:
|
||||
return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None)
|
||||
return BaseModelOutputWithPast(
|
||||
last_hidden_state=hidden_states,
|
||||
past_key_values=next_cache,
|
||||
hidden_states=all_hidden_states,
|
||||
attentions=all_self_attns,
|
||||
)
|
||||
|
||||
return forward
|
||||
|
||||
|
||||
def LlamaForCausalLM_forward():
|
||||
from transformers.utils import add_start_docstrings_to_model_forward, replace_return_docstrings
|
||||
from transformers.models.llama.modeling_llama import LLAMA_INPUTS_DOCSTRING, CausalLMOutputWithPast, _CONFIG_FOR_DOC
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids: torch.LongTensor = None,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
position_ids: Optional[torch.LongTensor] = None,
|
||||
past_key_values: Optional[List[torch.FloatTensor]] = None,
|
||||
inputs_embeds: Optional[torch.FloatTensor] = None,
|
||||
labels: Optional[torch.LongTensor] = None,
|
||||
use_cache: Optional[bool] = None,
|
||||
output_attentions: Optional[bool] = None,
|
||||
output_hidden_states: Optional[bool] = None,
|
||||
return_dict: Optional[bool] = None,
|
||||
) -> Union[Tuple, CausalLMOutputWithPast]:
|
||||
|
||||
output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
|
||||
output_hidden_states = (
|
||||
output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
|
||||
)
|
||||
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
|
||||
|
||||
# decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
|
||||
outputs = self.model(
|
||||
input_ids=input_ids,
|
||||
attention_mask=attention_mask,
|
||||
position_ids=position_ids,
|
||||
past_key_values=past_key_values,
|
||||
inputs_embeds=inputs_embeds,
|
||||
use_cache=use_cache,
|
||||
output_attentions=output_attentions,
|
||||
output_hidden_states=output_hidden_states,
|
||||
return_dict=return_dict,
|
||||
)
|
||||
|
||||
hidden_states = outputs[0]
|
||||
if self.config.pretraining_tp > 1:
|
||||
lm_head_slices = self.lm_head.weight.split(
|
||||
self.vocab_size // self.config.pretraining_tp, dim=0)
|
||||
logits = [F.linear(hidden_states, lm_head_slices[i])
|
||||
for i in range(self.config.pretraining_tp)]
|
||||
logits = torch.cat(logits, dim=-1)
|
||||
else:
|
||||
logits = self.lm_head(hidden_states)
|
||||
logits = logits.float()
|
||||
|
||||
loss = None
|
||||
if labels is not None:
|
||||
# Shift so that tokens < n predict n
|
||||
shift_logits = logits[..., :-1, :].contiguous()
|
||||
shift_labels = labels[..., 1:].contiguous()
|
||||
# Flatten the tokens
|
||||
loss_fct = CrossEntropyLoss
|
||||
shift_logits = shift_logits.view(-1, self.config.vocab_size)
|
||||
shift_labels = shift_labels.view(-1)
|
||||
# Enable model parallelism
|
||||
shift_labels = shift_labels.to(shift_logits.device)
|
||||
loss = loss_fct(shift_logits, shift_labels)
|
||||
|
||||
if not return_dict:
|
||||
output = (logits,) + outputs[1:]
|
||||
return (loss,) + output if loss is not None else output
|
||||
|
||||
return CausalLMOutputWithPast(
|
||||
loss=loss,
|
||||
logits=logits,
|
||||
past_key_values=outputs.past_key_values,
|
||||
hidden_states=outputs.hidden_states,
|
||||
attentions=outputs.attentions,
|
||||
)
|
||||
|
||||
return forward
|
||||
55
ixformer_sdk/train/speedformer/layers/llama/mlp.py
Normal file
55
ixformer_sdk/train/speedformer/layers/llama/mlp.py
Normal file
@@ -0,0 +1,55 @@
|
||||
import math
|
||||
import warnings
|
||||
from typing import List, Optional, Tuple, Union
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
import ixformer.train.functions as F
|
||||
from ixformer.train.speedformer.models.llama.configuration_llama import LlamaConfig
|
||||
from ixformer.train.speedformer.models.llama.modeling_llama import LlamaMLP
|
||||
from transformers import Cache
|
||||
from transformers.utils import logging
|
||||
|
||||
from ixformer.train.speedformer.layers.lazy import LazyInitContext
|
||||
|
||||
|
||||
class BaseLlamaMLP(LlamaMLP):
|
||||
"""
|
||||
这个层主要的优化点是：将linear1(act(cat(linear2(x), linear3(x))))的结构变成 linear1(act(linear23(x)))
|
||||
"""
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
self.gate_up = nn.Linear(
|
||||
self.hidden_size, self.intermediate_size * 2, bias=False)
|
||||
del self.gate_proj, self.up_proj
|
||||
del self.act_fn
|
||||
|
||||
def forward(self, x):
|
||||
res = self.gate_up(x)
|
||||
down_proj = self.down_proj(F.swiglu(res))
|
||||
return down_proj
|
||||
|
||||
|
||||
class IXFLlamaMLP(BaseLlamaMLP):
|
||||
def __init__(self) -> None:
|
||||
raise NotImplementedError(
|
||||
"IXFLlamaMLP is not implemented as a physical class. "
|
||||
"It is meant to be used only with the from_native_module interface to Convert a native LlamaAttention module to IXFLlamaMLP module provided above."
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def from_native_module(module: nn.Module, *args, **kwargs) -> nn.Module:
|
||||
|
||||
LazyInitContext.materialize(module)
|
||||
|
||||
config = getattr(module, "config")
|
||||
|
||||
mlp = BaseLlamaMLP(config=config)
|
||||
|
||||
mlp.gate_up.weight.data = torch.concat(
|
||||
(module.gate_proj.weight.data, module.up_proj.weight.data), dim=0)
|
||||
mlp.down_proj.weight.data = module.down_proj.weight.data
|
||||
|
||||
return mlp
|
||||
Reference in New Issue
Block a user