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:
61
ixformer_sdk/train/functions/group_norm.py
Normal file
61
ixformer_sdk/train/functions/group_norm.py
Normal file
@@ -0,0 +1,61 @@
|
||||
import ixformer._C as ops
|
||||
import torch
|
||||
from torch.nn import init
|
||||
from torch.nn.parameter import Parameter
|
||||
|
||||
|
||||
class GN_NHWC_Func(torch.autograd.Function):
|
||||
@staticmethod
|
||||
def forward(ctx, X: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor, G: int, eps: float, activation: str):
|
||||
X_out, means, rstds = ops.train.gn_nhwc_fwd(X, weight, bias, G, eps, activation)
|
||||
ctx.save_for_backward(X, weight, bias, means, rstds)
|
||||
ctx.G = G
|
||||
ctx.activation = activation
|
||||
return X_out
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, dy: torch.Tensor):
|
||||
dy = dy.contiguous(memory_format=torch.channels_last)
|
||||
X, weight, bias, means, rstds = ctx.saved_tensors
|
||||
dx, dgamma, dbeta = ops.train.gn_nhwc_bwd(dy, X, weight, bias, means, rstds, ctx.G, ctx.activation)
|
||||
return dx, dgamma, dbeta, None, None, None
|
||||
|
||||
|
||||
class GroupNorm_nhwc(torch.nn.GroupNorm):
|
||||
def __init__(self, num_groups: int, nc: int, activation='identity', **kwargs):
|
||||
super().__init__(num_groups, nc, **kwargs)
|
||||
assert activation in {'identity', 'silu', 'relu', 'gelu', 'gelu_tanh'}
|
||||
if activation == 'identity':
|
||||
self.activation = 0
|
||||
if activation == 'relu':
|
||||
self.activation = 1
|
||||
if activation == 'silu':
|
||||
self.activation = 2
|
||||
if activation == 'gelu':
|
||||
self.activation = 3
|
||||
if activation == 'gelu_tanh':
|
||||
self.activation = 4
|
||||
|
||||
@torch._dynamo.disable
|
||||
def forward(self, x):
|
||||
#print(x.shape, self.num_channels)
|
||||
if len(x.size()) == 3:
|
||||
N, C, L = x.shape
|
||||
elif len(x.size()) == 4:
|
||||
N, C, H, W = x.shape
|
||||
else:
|
||||
raise ValueError
|
||||
G = self.num_groups
|
||||
|
||||
#if C // G > 512:
|
||||
# raise ValueError(f'Error in fwd for X.shape={x.shape}, G={G}: C // G = {C // G} which is greater than 512. This input is not supported.')
|
||||
|
||||
#if H * W % 8 != 0:
|
||||
# raise ValueError(f'Error in fwd for X.shape={x.shape}, G={G}: H * W is not a multiple of 8. This input is not supported.')
|
||||
|
||||
if self.affine:
|
||||
return GN_NHWC_Func.apply(x, self.weight, self.bias, self.num_groups, self.eps, self.activation)
|
||||
else:
|
||||
w = torch.ones((self.num_channels,), device=x.device, dtype=x.dtype)
|
||||
b = torch.zeros((self.num_channels,), device=x.device, dtype=x.dtype)
|
||||
return GN_NHWC_Func.apply(x, w, b, self.num_groups, self.eps, self.activation)
|
||||
Reference in New Issue
Block a user