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:
97
ixformer_sdk/inference/functions/t5.py
Normal file
97
ixformer_sdk/inference/functions/t5.py
Normal file
@@ -0,0 +1,97 @@
|
||||
from typing import Union
|
||||
|
||||
import ixformer._C as ops
|
||||
import torch
|
||||
|
||||
__all__ = [
|
||||
"t5_split_qkv",
|
||||
"t5_split_qkv_update_kv_cache",
|
||||
"ref_t5_split_qkv_update_kv_cache",
|
||||
"ref_t5_split_qkv",
|
||||
]
|
||||
|
||||
|
||||
def reshape_query(query, head_num, head_dim):
|
||||
batch_size, seq_len, _ = query.shape
|
||||
query = query.view(batch_size, seq_len, head_num, head_dim)
|
||||
query = query.transpose(1, 2)
|
||||
return query
|
||||
|
||||
|
||||
def ref_t5_split_qkv(qkv: "torch.Tensor", head_num: int, head_dim: int):
|
||||
assert qkv.size(-1) == head_dim * head_num * 3
|
||||
batch_size, seq_len, _ = qkv.shape
|
||||
q, k, v = torch.chunk(qkv, 3, dim=-1)
|
||||
q = reshape_query(q, head_num, head_dim)
|
||||
k = reshape_query(k, head_num, head_dim)
|
||||
v = reshape_query(v, head_num, head_dim)
|
||||
return q, k, v
|
||||
|
||||
|
||||
def t5_split_qkv(qkv: "torch.Tensor", head_num: int, head_dim: int):
|
||||
|
||||
"""
|
||||
Args:
|
||||
qkv: (batch_size, seq_len, head_dim * head_num * 3) torch.half, torch.bfloat16
|
||||
head_num: int
|
||||
head_dim: int
|
||||
Returns:
|
||||
q: (batch_size, head_num, seq_len, head_dim) torch.half, torch.bfloat16
|
||||
k: (batch_size, head_num, seq_len, head_dim) torch.half, torch.bfloat16
|
||||
v: (batch_size, head_num, seq_len, head_dim) torch.half, torch.bfloat16
|
||||
"""
|
||||
batch_size, seq_len, _ = qkv.shape
|
||||
q = qkv.new_empty([batch_size, head_num, seq_len, head_dim])
|
||||
k = qkv.new_empty([batch_size, head_num, seq_len, head_dim])
|
||||
v = qkv.new_empty([batch_size, head_num, seq_len, head_dim])
|
||||
ops.infer.t5_split_qkv(qkv, q, k, v, head_num, head_dim)
|
||||
return q, k, v
|
||||
|
||||
|
||||
def ref_t5_split_qkv_update_kv_cache(
|
||||
qkv: "torch.Tensor",
|
||||
past_key: "torch.Tensor",
|
||||
past_value: "torch.Tensor",
|
||||
head_num: int,
|
||||
head_dim: int,
|
||||
):
|
||||
assert qkv.size(-1) == head_dim * head_num * 3
|
||||
batch_size, seq_len, _ = qkv.shape
|
||||
q, k, v = torch.chunk(qkv, 3, dim=-1)
|
||||
q = reshape_query(q, head_num, head_dim)
|
||||
k = reshape_query(k, head_num, head_dim)
|
||||
v = reshape_query(v, head_num, head_dim)
|
||||
k = torch.cat([past_key, k], dim=2)
|
||||
v = torch.cat([past_value, v], dim=2)
|
||||
return q, k, v
|
||||
|
||||
|
||||
def t5_split_qkv_update_kv_cache(
|
||||
qkv: "torch.Tensor",
|
||||
past_key: "torch.Tensor",
|
||||
past_value: "torch.Tensor",
|
||||
head_num: int,
|
||||
head_dim: int,
|
||||
):
|
||||
|
||||
"""
|
||||
Args:
|
||||
qkv: (batch_size, 1 , head_dim * head_num * 3) torch.half, torch.bfloat16
|
||||
past_key: (batch_size, head_num, seq_len - 1, head_dim) torch.half, torch.bfloat16
|
||||
past_value: (batch_size, head_num, seq_len - 1, head_dim) torch.half, torch.bfloat16
|
||||
head_num: int
|
||||
head_dim: int
|
||||
Returns:
|
||||
q: (batch_size, head_num, 1, head_dim) torch.half, torch.bfloat16
|
||||
k: (batch_size, head_num, seq_len, head_dim) torch.half, torch.bfloat16
|
||||
v: (batch_size, head_num, seq_len, head_dim) torch.half, torch.bfloat16
|
||||
"""
|
||||
batch_size, _, past_seq_len, _ = list(past_key.shape)
|
||||
seq_len = past_seq_len + 1
|
||||
q = qkv.new_empty([batch_size, head_num, 1, head_dim])
|
||||
k = qkv.new_empty([batch_size, head_num, seq_len, head_dim])
|
||||
v = qkv.new_empty([batch_size, head_num, seq_len, head_dim])
|
||||
ops.infer.t5_split_qkv_update_kv_cache(
|
||||
qkv, past_key, past_value, q, k, v, head_num, head_dim
|
||||
)
|
||||
return q, k, v
|
||||
Reference in New Issue
Block a user