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
199
ixformer_sdk/train/speedformer/layers/chatglm/attention.py
Normal file
199
ixformer_sdk/train/speedformer/layers/chatglm/attention.py
Normal file
@@ -0,0 +1,199 @@
|
||||
import math
|
||||
import os
|
||||
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.chatglm.modeling_chatglm import (
|
||||
CoreAttention,
|
||||
SelfAttention,
|
||||
split_tensor_along_last_dim,
|
||||
apply_rotary_pos_emb
|
||||
)
|
||||
from ixformer.train.speedformer.models.chatglm.configuration_chatglm import ChatGLMConfig
|
||||
|
||||
from transformers.utils import is_flash_attn_2_available
|
||||
|
||||
if is_flash_attn_2_available():
|
||||
from flash_attn import flash_attn_func, flash_attn_varlen_func
|
||||
from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input # noqa
|
||||
|
||||
|
||||
class FlashCoreAttention(CoreAttention):
|
||||
|
||||
def forward(self, query_layer, key_layer, value_layer, attention_mask):
|
||||
if int(os.environ.get("USE_FLASH_ATTN", 0)):
|
||||
query_layer, key_layer, value_layer = [
|
||||
k.permute(1, 0, 2, 3) for k in [query_layer, key_layer, value_layer]]
|
||||
batch_size, query_length, _, _ = query_layer.shape
|
||||
|
||||
if attention_mask is not None:
|
||||
batch_size = query_layer.shape[0]
|
||||
query_layer, key_layer, value_layer, indices_q, cu_seq_lens, max_seq_lens = self._upad_input(
|
||||
query_layer, key_layer, value_layer, attention_mask, query_length
|
||||
)
|
||||
|
||||
cu_seqlens_q, cu_seqlens_k = cu_seq_lens
|
||||
max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens
|
||||
|
||||
attn_output_unpad = flash_attn_varlen_func(
|
||||
query_layer,
|
||||
key_layer,
|
||||
value_layer,
|
||||
cu_seqlens_q=cu_seqlens_q,
|
||||
cu_seqlens_k=cu_seqlens_k,
|
||||
max_seqlen_q=max_seqlen_in_batch_q,
|
||||
max_seqlen_k=max_seqlen_in_batch_k,
|
||||
dropout_p=0.0,
|
||||
softmax_scale=None,
|
||||
causal=True,
|
||||
)
|
||||
attn_output = pad_input(
|
||||
attn_output_unpad, indices_q, batch_size, query_length)
|
||||
context_layer = attn_output.permute(1, 0, 2, 3)
|
||||
else:
|
||||
attn_output = flash_attn_func(
|
||||
query_layer, key_layer, value_layer, 0.0, softmax_scale=None, causal=True
|
||||
)
|
||||
context_layer = attn_output.permute(1, 0, 2, 3)
|
||||
|
||||
if attention_mask is not None:
|
||||
if query_layer.shape[2] != key_layer.shape[2]:
|
||||
num_group = query_layer.shape[2] // key_layer.shape[2]
|
||||
final_shape = (*key_layer.shape[:2], *query_layer.shape[2:])
|
||||
key_layer = key_layer.unsqueeze(-2)
|
||||
key_layer = key_layer.expand(
|
||||
-1, -1, -1, num_group, -1
|
||||
)
|
||||
key_layer = key_layer.contiguous().view(
|
||||
final_shape
|
||||
)
|
||||
value_layer = value_layer.unsqueeze(-2)
|
||||
value_layer = value_layer.expand(
|
||||
-1, -1, -1, num_group, -1
|
||||
)
|
||||
value_layer = value_layer.contiguous().view(
|
||||
final_shape
|
||||
)
|
||||
|
||||
query_layer, key_layer, value_layer = [
|
||||
k.permute(1, 2, 0, 3) for k in [query_layer, key_layer, value_layer]] # bhsd
|
||||
attention_mask = ~attention_mask
|
||||
context_layer = torch.nn.functional.scaled_dot_product_attention(query_layer, key_layer, value_layer,
|
||||
attention_mask)
|
||||
context_layer = context_layer.permute(2, 0, 1, 3)
|
||||
|
||||
else:
|
||||
query_layer, key_layer, value_layer = [
|
||||
k.permute(1, 0, 2, 3) for k in [query_layer, key_layer, value_layer]] # bshd
|
||||
context_layer = flash_attn_func(
|
||||
query_layer, key_layer, value_layer, 0, softmax_scale=None, causal=True
|
||||
) # bshd
|
||||
context_layer = context_layer.permute(1, 0, 2, 3)
|
||||
|
||||
context_layer = context_layer.reshape(
|
||||
context_layer.size(0), context_layer.size(1), -1)
|
||||
|
||||
return context_layer
|
||||
|
||||
|
||||
class FlashSelfAttention(SelfAttention):
|
||||
|
||||
def __init__(self, config: ChatGLMConfig, layer_number, device=None):
|
||||
super().__init__(config, layer_number, device=device)
|
||||
self.core_attention = FlashCoreAttention(config, self.layer_number)
|
||||
|
||||
def forward(self, hidden_states, attention_mask, rotary_pos_emb, kv_cache=None, use_cache=True):
|
||||
mixed_x_layer = self.query_key_value(hidden_states)
|
||||
if self.multi_query_attention:
|
||||
(query_layer, key_layer, value_layer) = mixed_x_layer.split(
|
||||
[
|
||||
self.num_attention_heads_per_partition * self.hidden_size_per_attention_head,
|
||||
self.num_multi_query_groups_per_partition * self.hidden_size_per_attention_head,
|
||||
self.num_multi_query_groups_per_partition * self.hidden_size_per_attention_head,
|
||||
],
|
||||
dim=-1,
|
||||
)
|
||||
query_layer = query_layer.view(
|
||||
query_layer.size()[
|
||||
:-1] + (self.num_attention_heads_per_partition, self.hidden_size_per_attention_head)
|
||||
)
|
||||
key_layer = key_layer.view(
|
||||
key_layer.size()[
|
||||
:-1] + (self.num_multi_query_groups_per_partition, self.hidden_size_per_attention_head)
|
||||
)
|
||||
value_layer = value_layer.view(
|
||||
value_layer.size()[:-1]
|
||||
+ (self.num_multi_query_groups_per_partition,
|
||||
self.hidden_size_per_attention_head)
|
||||
)
|
||||
else:
|
||||
new_tensor_shape = mixed_x_layer.size()[:-1] + \
|
||||
(self.num_attention_heads_per_partition,
|
||||
3 * self.hidden_size_per_attention_head)
|
||||
mixed_x_layer = mixed_x_layer.view(*new_tensor_shape)
|
||||
|
||||
# [sq, b, np, 3 * hn] --> 3 [sq, b, np, hn]
|
||||
(query_layer, key_layer, value_layer) = split_tensor_along_last_dim(
|
||||
mixed_x_layer, 3)
|
||||
|
||||
if rotary_pos_emb is not None:
|
||||
query_layer = apply_rotary_pos_emb(query_layer, rotary_pos_emb)
|
||||
key_layer = apply_rotary_pos_emb(key_layer, rotary_pos_emb)
|
||||
|
||||
# adjust key and value for inference
|
||||
if kv_cache is not None:
|
||||
cache_k, cache_v = kv_cache
|
||||
key_layer = torch.cat((cache_k, key_layer), dim=0)
|
||||
value_layer = torch.cat((cache_v, value_layer), dim=0)
|
||||
if use_cache:
|
||||
kv_cache = (key_layer, value_layer)
|
||||
else:
|
||||
kv_cache = None
|
||||
|
||||
# 这里省略了 kv "sbhd" -> "sb(h*num_multi-group)d" 的过程，因为flash-attn支持 MGA
|
||||
# ==================================
|
||||
# core attention computation
|
||||
# ==================================
|
||||
|
||||
context_layer = self.core_attention(
|
||||
query_layer, key_layer, value_layer, attention_mask)
|
||||
|
||||
# =================
|
||||
# Output. [sq, b, h]
|
||||
# =================
|
||||
|
||||
output = self.dense(context_layer)
|
||||
|
||||
return output, kv_cache
|
||||
|
||||
|
||||
class ChatglmFlashAttention(FlashSelfAttention):
|
||||
|
||||
def __init__(self) -> None:
|
||||
raise NotImplementedError(
|
||||
"BloomAttention is not implemented as a physical class. "
|
||||
"It is meant to be used only with the from_native_module interface to Convert a native BloomAttention module to FlashAttention module provided above."
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def from_native_module(module: nn.Module, *args, **kwargs) -> nn.Module:
|
||||
# 这个原实现没有在类中保存config，所以需要初始化一个config
|
||||
layer_number = getattr(module, "layer_number")
|
||||
config = getattr(module, "config")
|
||||
attention = FlashSelfAttention(
|
||||
config=config,
|
||||
layer_number=layer_number,
|
||||
)
|
||||
|
||||
attention.query_key_value.weight.data = module.query_key_value.weight.data
|
||||
attention.dense.weight.data = module.dense.weight.data
|
||||
if getattr(attention.query_key_value, "bias") is not None:
|
||||
attention.query_key_value.bias.data = module.query_key_value.bias.data
|
||||
if getattr(attention.dense, "bias") is not None:
|
||||
attention.dense.bias.data = module.dense.bias.data
|
||||
|
||||
return attention
|
||||
@@ -0,0 +1,9 @@
|
||||
from ixformer.train.speedformer.models.chatglm.modeling_chatglm import RotaryEmbedding
|
||||
from ixformer.train.speedformer.layers.rotary_pos_embedding import RotaryEmbedding
|
||||
|
||||
|
||||
class ChatglmRotaryEmbedding(RotaryEmbedding):
|
||||
def from_native_attr(attr_class, *args, **kwargs):
|
||||
dim = attr_class.dim
|
||||
rote = RotaryEmbedding(dim=dim)
|
||||
return rote
|
||||
127
ixformer_sdk/train/speedformer/layers/chatglm/methods.py
Normal file
127
ixformer_sdk/train/speedformer/layers/chatglm/methods.py
Normal file
@@ -0,0 +1,127 @@
|
||||
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.chatglm.modeling_chatglm import ChatGLMModel
|
||||
|
||||
|
||||
def ChatGLMModel_forward():
|
||||
from transformers.modeling_outputs import BaseModelOutputWithPast
|
||||
from transformers.utils import logging, is_flash_attn_2_available
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids,
|
||||
position_ids: Optional[torch.Tensor] = None,
|
||||
attention_mask: Optional[torch.BoolTensor] = None,
|
||||
full_attention_mask: Optional[torch.BoolTensor] = None,
|
||||
past_key_values: Optional[Tuple[Tuple[torch.Tensor, torch.Tensor], ...]] = None,
|
||||
inputs_embeds: Optional[torch.Tensor] = None,
|
||||
use_cache: Optional[bool] = None,
|
||||
output_hidden_states: Optional[bool] = None,
|
||||
return_dict: Optional[bool] = None,
|
||||
):
|
||||
def is_lower_triangular(mask):
|
||||
"""
|
||||
ixdnn 虽然支持2种causal mask, 如下图:
|
||||
mode0:
|
||||
if seqlen_q < seqlen_k
|
||||
1 0 0 0 0
|
||||
1 1 0 0 0
|
||||
if seqlen_k < seqlen_q
|
||||
1 0
|
||||
1 1
|
||||
1 1
|
||||
1 1
|
||||
1 1
|
||||
mode1:
|
||||
if seqlen_q < seqlen_k
|
||||
1 1 1 1 0
|
||||
1 1 1 1 1
|
||||
if seqlen_k < seqlen_q
|
||||
0 0
|
||||
0 0
|
||||
0 0
|
||||
1 0
|
||||
1 1
|
||||
|
||||
但 flash-attn 目前只支持 mode1, 所以下面需要判断一下传入的mask是不是mode1这种模式
|
||||
"""
|
||||
batch_size, _, rows, cols = mask.shape
|
||||
|
||||
# 创建一个mode1的下三角矩阵
|
||||
if rows <= cols:
|
||||
part = torch.ones(rows, cols - rows,
|
||||
dtype=torch.bool, device=mask.device)
|
||||
gt = ~torch.triu(torch.ones(
|
||||
rows, rows, dtype=torch.bool, device=mask.device), diagonal=1)
|
||||
gt = torch.cat((part, gt), dim=1)
|
||||
else:
|
||||
part = torch.zeros(
|
||||
rows-cols, cols, dtype=torch.bool, device=mask.device)
|
||||
gt = ~torch.triu(torch.ones(
|
||||
cols, cols, dtype=torch.bool, device=mask.device), diagonal=1)
|
||||
gt = torch.cat((part, gt), dim=0)
|
||||
gt = gt[None, None, :, :].expand(batch_size, -1, -1, -1)
|
||||
|
||||
# 检查所有的元素是不是都一样
|
||||
check = (gt == mask).all()
|
||||
|
||||
return check
|
||||
|
||||
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
|
||||
|
||||
batch_size, seq_length = input_ids.shape
|
||||
|
||||
if inputs_embeds is None:
|
||||
inputs_embeds = self.embedding(input_ids)
|
||||
|
||||
if self.pre_seq_len is not None:
|
||||
if past_key_values is None:
|
||||
past_key_values = self.get_prompt(batch_size=batch_size, device=input_ids.device,
|
||||
dtype=inputs_embeds.dtype)
|
||||
if attention_mask is not None:
|
||||
attention_mask = torch.cat([attention_mask.new_ones((batch_size, self.pre_seq_len)),
|
||||
attention_mask], dim=-1)
|
||||
|
||||
if full_attention_mask is None:
|
||||
if (attention_mask is not None and not attention_mask.all()) or (past_key_values and seq_length != 1):
|
||||
full_attention_mask = self.get_masks(
|
||||
input_ids, past_key_values, padding_mask=attention_mask)
|
||||
|
||||
# Rotary positional embeddings
|
||||
rotary_pos_emb = self.rotary_pos_emb(self.seq_length)
|
||||
if position_ids is not None:
|
||||
rotary_pos_emb = rotary_pos_emb[position_ids]
|
||||
else:
|
||||
rotary_pos_emb = rotary_pos_emb[None, :seq_length]
|
||||
rotary_pos_emb = rotary_pos_emb.transpose(0, 1).contiguous()
|
||||
|
||||
# Run encoder.
|
||||
attn_mask = None
|
||||
if full_attention_mask is not None:
|
||||
if not is_lower_triangular(full_attention_mask):
|
||||
attn_mask = full_attention_mask
|
||||
hidden_states, presents, all_hidden_states, all_self_attentions = self.encoder(
|
||||
inputs_embeds, attn_mask, rotary_pos_emb=rotary_pos_emb,
|
||||
kv_caches=past_key_values, use_cache=use_cache, output_hidden_states=output_hidden_states
|
||||
)
|
||||
|
||||
if not return_dict:
|
||||
return tuple(v for v in [hidden_states, presents, all_hidden_states, all_self_attentions] if v is not None)
|
||||
|
||||
return BaseModelOutputWithPast(
|
||||
last_hidden_state=hidden_states,
|
||||
past_key_values=presents,
|
||||
hidden_states=all_hidden_states,
|
||||
attentions=all_self_attentions,
|
||||
)
|
||||
return forward
|
||||
Reference in New Issue
Block a user