87a19d2d00
来源:
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
98 lines
3.7 KiB
Python
98 lines
3.7 KiB
Python
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
|