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
0
ixformer_sdk/core/__init__.py
Normal file
0
ixformer_sdk/core/__init__.py
Normal file
184
ixformer_sdk/core/config.py
Normal file
184
ixformer_sdk/core/config.py
Normal file
@@ -0,0 +1,184 @@
|
||||
import os
|
||||
from typing import Callable, Optional
|
||||
|
||||
# =========================================================
|
||||
# Utils
|
||||
# =========================================================
|
||||
|
||||
|
||||
def number_type(scalar_type):
|
||||
def wrap(val: Optional[str]):
|
||||
if val is None:
|
||||
return None
|
||||
|
||||
return scalar_type(val)
|
||||
|
||||
return wrap
|
||||
|
||||
|
||||
def bool_type(val: Optional[str]):
|
||||
if val is None:
|
||||
return False
|
||||
|
||||
if isinstance(val, str):
|
||||
return val.lower() in ["1", "t", "true"]
|
||||
|
||||
if isinstance(val, int):
|
||||
return val != 0
|
||||
|
||||
raise RuntimeError(f"Invalid bool type, got {type(val), val}")
|
||||
|
||||
|
||||
def list_type(scalar_type=str):
|
||||
def wrap(val: Optional[str]):
|
||||
if val is None:
|
||||
return []
|
||||
|
||||
if not isinstance(val, str):
|
||||
raise RuntimeError(
|
||||
f"list_type: Got invalid type, expect str, but got {val}."
|
||||
)
|
||||
|
||||
return [scalar_type(v) for v in val.split(",")]
|
||||
|
||||
return wrap
|
||||
|
||||
|
||||
def Field(
|
||||
name: str,
|
||||
static: bool = True,
|
||||
type: Callable = str,
|
||||
choices: Optional[list] = None,
|
||||
help: Optional[str] = None,
|
||||
**kwargs,
|
||||
):
|
||||
"""
|
||||
Define environment variable field
|
||||
|
||||
Example:
|
||||
Static mode:
|
||||
# define
|
||||
ENABLE_XX = Field("ENABLE_XX", type=bool, help="ENABLE_XX")
|
||||
|
||||
# use
|
||||
config.ENABLE_XX
|
||||
|
||||
Dynamic mode:
|
||||
# Please use lowercase naming to differentiate it with static mode.
|
||||
|
||||
# define
|
||||
enable_cc = Field("ENABLE_CC", type=bool, static=False, help="enable_cc")
|
||||
|
||||
# use
|
||||
config.enable_cc()
|
||||
|
||||
Set default value:
|
||||
# define
|
||||
ENABLE_TT = Field("ENABLE_TT", type=bool, default=False, help="ENABLE_TT")
|
||||
|
||||
# use
|
||||
config.ENABLE_TT
|
||||
|
||||
Use list:
|
||||
# define
|
||||
CUDA_VISIBLE_DEVICES = Field("CUDA_VISIBLE_DEVICES", type=list_type(int), help="CUDA_VISIBLE_DEVICES")
|
||||
|
||||
# use
|
||||
# the CUDA_VISIBLE_DEVICES is parsed to list, and it's value is int type.
|
||||
for device_id in CUDA_VISIBLE_DEVICES:
|
||||
...
|
||||
|
||||
"""
|
||||
|
||||
if type == bool:
|
||||
type = bool_type
|
||||
|
||||
elif type in [list, tuple]:
|
||||
type = list_type(scalar_type=str)
|
||||
|
||||
elif type in [int, float]:
|
||||
type = number_type(type)
|
||||
|
||||
if static:
|
||||
env_val = type(os.environ.get(name, **kwargs))
|
||||
if choices is not None and env_val is not None and env_val not in choices:
|
||||
raise RuntimeError(
|
||||
f"Got invalid value, expect {choices}, but got {env_val}."
|
||||
)
|
||||
return env_val
|
||||
|
||||
def _get():
|
||||
env_val = type(os.environ.get(name, **kwargs))
|
||||
if choices is not None and env_val is not None and env_val not in choices:
|
||||
raise RuntimeError(
|
||||
f"Got invalid value, expect {choices}, but got {env_val}."
|
||||
)
|
||||
return env_val
|
||||
|
||||
return _get
|
||||
|
||||
|
||||
# =========================================================
|
||||
# Functions Config
|
||||
# =========================================================
|
||||
|
||||
IXFORMER_GEMV_THRESHOLD = Field(
|
||||
"IXFORMER_GEMV_THRESHOLD",
|
||||
type=int,
|
||||
default=1,
|
||||
help="Set the threshold for using gemv.",
|
||||
)
|
||||
|
||||
|
||||
# =========================================================
|
||||
# Distributed Config
|
||||
# =========================================================
|
||||
|
||||
IXFORMER_COMM_SHM_SIZE = Field(
|
||||
"IXFORMER_COMM_SHM_SIZE",
|
||||
type=int,
|
||||
default=None,
|
||||
help="set shared memory size of ipc comm.",
|
||||
)
|
||||
|
||||
IXFORMER_ENABLE_OVERLAP_COMM = Field(
|
||||
"IXFORMER_ENABLE_OVERLAP_COMM",
|
||||
type=bool,
|
||||
default=False,
|
||||
help="enable overlap communcation and compute.",
|
||||
)
|
||||
|
||||
IXFORMER_OVERLAP_GEMM_METHOD = Field(
|
||||
"IXFORMER_OVERLAP_GEMM_METHOD",
|
||||
type=int,
|
||||
default=None,
|
||||
choices=[0, 1],
|
||||
help="set gemm backend, 0: ixinfer, 1: cublas.",
|
||||
)
|
||||
|
||||
IXFORMER_OVERLAP_CHUNKS = Field(
|
||||
"IXFORMER_OVERLAP_CHUNKS", type=int, default=2, help="set split chunks."
|
||||
)
|
||||
|
||||
IXFORMER_OVERLAP_SPLIT_RATIO = Field(
|
||||
"IXFORMER_OVERLAP_SPLIT_RATIO",
|
||||
type=float,
|
||||
default=None,
|
||||
help="set split chunks ratio.",
|
||||
)
|
||||
|
||||
IXFORMER_PAGED_ATTENTION_ALGO = Field(
|
||||
"IXFORMER_PAGED_ATTENTION_ALGO",
|
||||
type=str,
|
||||
default="ixinfer",
|
||||
choices=["ixinfer", "ixformer"],
|
||||
help="set paged attention algo.",
|
||||
)
|
||||
|
||||
IXFORMER_UNPAD_ATTENTION_ALGO = Field(
|
||||
"IXFORMER_UNPAD_ATTENTION_ALGO",
|
||||
type=str,
|
||||
default="ixinfer",
|
||||
choices=["ixinfer", "ixinfer-ex"],
|
||||
help="set enpad attention algo.",
|
||||
)
|
||||
20
ixformer_sdk/core/dispatcher.py
Normal file
20
ixformer_sdk/core/dispatcher.py
Normal file
@@ -0,0 +1,20 @@
|
||||
class Dispatcher(object):
|
||||
"""
|
||||
create object by dispatcher to reuse object.
|
||||
"""
|
||||
|
||||
_dispatcher = dict()
|
||||
|
||||
@classmethod
|
||||
def dispatcher(cls, *args, **kwargs):
|
||||
key = cls.dispatcher_key(*args, **kwargs)
|
||||
obj = cls._dispatcher.get(key, None)
|
||||
if obj is None:
|
||||
obj = cls(*args, **kwargs)
|
||||
cls._dispatcher[key] = obj
|
||||
|
||||
return obj
|
||||
|
||||
@classmethod
|
||||
def dispatcher_key(cls, *args, **kwargs):
|
||||
raise NotImplementedError()
|
||||
54
ixformer_sdk/core/multi_level_cache.py
Normal file
54
ixformer_sdk/core/multi_level_cache.py
Normal file
@@ -0,0 +1,54 @@
|
||||
class MultiLevelCache(object):
|
||||
def __init__(self):
|
||||
self._l1_key = None
|
||||
self._l1_value = None
|
||||
|
||||
self._l2_size = 3
|
||||
self._l2 = [(None, None) for _ in range(self._l2_size)]
|
||||
self._l2_ptr = 0
|
||||
|
||||
self._l3 = dict()
|
||||
|
||||
def set(self, key, value):
|
||||
self._l1_key = key
|
||||
self._l1_value = value
|
||||
|
||||
self._l2[self._l2_ptr] = (key, value)
|
||||
self._l2_ptr = (self._l2_ptr + 1) % 3 # l2_size: 3
|
||||
|
||||
self._l3[key] = value
|
||||
|
||||
def get(self, key, *args):
|
||||
if key == self._l1_key:
|
||||
return self._l1_value
|
||||
|
||||
l2 = self._l2
|
||||
if key == l2[0][0]:
|
||||
return l2[0][1]
|
||||
|
||||
if key == l2[1][0]:
|
||||
return l2[1][1]
|
||||
|
||||
if key == l2[2][0]:
|
||||
return l2[2][1]
|
||||
|
||||
return self._l3.get(key, *args)
|
||||
|
||||
def containe(self, key):
|
||||
return key in self._l3
|
||||
|
||||
def __getitem__(self, item):
|
||||
return self.get(item)
|
||||
|
||||
def __setitem__(self, key, value):
|
||||
self.set(key, value)
|
||||
|
||||
def __contains__(self, item):
|
||||
if item == self._l1_key:
|
||||
return True
|
||||
|
||||
l2 = self._l2
|
||||
if item == l2[0][0] or item == l2[1][0] or item == l2[2][0]:
|
||||
return True
|
||||
|
||||
return item in self._l3
|
||||
237
ixformer_sdk/core/operator_autotuning.py
Normal file
237
ixformer_sdk/core/operator_autotuning.py
Normal file
@@ -0,0 +1,237 @@
|
||||
import abc
|
||||
import bisect
|
||||
import functools
|
||||
import itertools
|
||||
import random
|
||||
from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple, Union
|
||||
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
|
||||
import ixformer.distributed as ixfd
|
||||
from ixformer.utils.benchmark.cuda_benchmark import Functor, cuda_benchmark
|
||||
|
||||
|
||||
def sync_ranks_metric(value, group=None):
|
||||
if not isinstance(value, (torch.Tensor, int, float)):
|
||||
raise RuntimeError(
|
||||
f"Invalid metric value, expect `Tensor`, `int`, or `float` type, but got {value}."
|
||||
)
|
||||
|
||||
if torch.is_tensor(value):
|
||||
value = value.to("cuda")
|
||||
else:
|
||||
value = torch.tensor([value], dtype=torch.float, device="cuda")
|
||||
|
||||
dist.broadcast(value, src=0, group=group)
|
||||
return value.cpu().item()
|
||||
|
||||
|
||||
class AutotuningFinder(object):
|
||||
def freeze(self):
|
||||
pass
|
||||
|
||||
@abc.abstractmethod
|
||||
def get(self, key) -> Callable:
|
||||
pass
|
||||
|
||||
@abc.abstractmethod
|
||||
def set(self, *args, **kwargs):
|
||||
pass
|
||||
|
||||
|
||||
class BasedKeyFinder(AutotuningFinder):
|
||||
def __init__(self):
|
||||
self._key_to_value: Dict[Any, Callable] = dict()
|
||||
|
||||
def get(self, key, **kwargs) -> Callable:
|
||||
if "default" in kwargs:
|
||||
return self._key_to_value.get(key, kwargs["default"])
|
||||
return self._key_to_value[key]
|
||||
|
||||
def set(self, key, value):
|
||||
self._key_to_value[key] = value
|
||||
|
||||
def containe(self, key):
|
||||
return key in self._key_to_value
|
||||
|
||||
|
||||
class TreeNode:
|
||||
def __init__(self):
|
||||
self.nodes: List[Union[Any, TreeNode]] = list()
|
||||
self.key_to_nodes: Dict[Any, TreeNode] = dict()
|
||||
|
||||
def add(self, key, value):
|
||||
if isinstance(key, (tuple, list)):
|
||||
if len(key) == 1:
|
||||
self.insert_value(key[0], value)
|
||||
else:
|
||||
self.recurse_add_node(key, value)
|
||||
else:
|
||||
self.insert_value(key, value)
|
||||
|
||||
def insert_value(self, key, value):
|
||||
self.nodes.append((key, value))
|
||||
self.key_to_nodes[key] = value
|
||||
|
||||
def recurse_add_node(self, key, value):
|
||||
if key[0] in self.key_to_nodes:
|
||||
node = self.key_to_nodes[key[0]]
|
||||
else:
|
||||
node = TreeNode()
|
||||
self.key_to_nodes[key[0]] = node
|
||||
self.insert_value(key[0], node)
|
||||
|
||||
node.add(key[1:], value)
|
||||
|
||||
def sort(self):
|
||||
self.nodes.sort(key=lambda x: x[0])
|
||||
for _, node in self.nodes:
|
||||
if isinstance(node, TreeNode):
|
||||
node.sort()
|
||||
|
||||
def find(self, key):
|
||||
is_list_key = isinstance(key, (tuple, list))
|
||||
if not is_list_key:
|
||||
key = (key,)
|
||||
|
||||
num_querys = len(key)
|
||||
node = self
|
||||
for key_idx in range(num_querys):
|
||||
query_key = key[key_idx]
|
||||
idx = bisect.bisect_left(node.nodes, (query_key,)) - 1
|
||||
if idx <= 0:
|
||||
node = node.nodes[0][1]
|
||||
elif idx >= len(node.nodes):
|
||||
node = node.nodes[-1][1]
|
||||
else:
|
||||
node = node.nodes[idx][1]
|
||||
|
||||
return node
|
||||
|
||||
def show(self, indent=0):
|
||||
for k, node in self.nodes:
|
||||
print(" " * indent, end="")
|
||||
if isinstance(node, TreeNode):
|
||||
print(f"key: {k}")
|
||||
node.show(indent=indent + 4)
|
||||
else:
|
||||
print(f"key: {k}, node: {node}")
|
||||
|
||||
|
||||
class BasedRangeFinder(AutotuningFinder):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
|
||||
self.tree = TreeNode()
|
||||
self._found_cache: Dict[Any, Callable] = dict()
|
||||
|
||||
def freeze(self):
|
||||
self.tree.sort()
|
||||
|
||||
def get(self, key) -> Callable:
|
||||
value = self._found_cache.get(key, None)
|
||||
if value is not None:
|
||||
return value
|
||||
|
||||
value = self.tree.find(key)
|
||||
self._found_cache[key] = value
|
||||
return value
|
||||
|
||||
def set(self, key, value):
|
||||
self.tree.add(key, value)
|
||||
|
||||
|
||||
class OperatorAutotuning(object):
|
||||
def __init__(self, num_repeated=5, num_warmup=3, dist_barrier=False):
|
||||
self.num_repeated = num_repeated
|
||||
self.num_warmup = num_warmup
|
||||
self.dist_barrier = dist_barrier
|
||||
|
||||
@abc.abstractmethod
|
||||
def operators(self):
|
||||
raise NotImplementedError()
|
||||
|
||||
def __call__(self, *args, **kwargs):
|
||||
return self.exec_best_operator(args, kwargs)
|
||||
|
||||
@abc.abstractmethod
|
||||
def exec_best_operator(self, args, kwargs):
|
||||
raise NotImplementedError()
|
||||
|
||||
@abc.abstractmethod
|
||||
def autotuning(self, *args, **kwargs):
|
||||
raise NotImplementedError()
|
||||
|
||||
def perf_best_operator(self, *args, **kwargs) -> Callable:
|
||||
best_operator = None
|
||||
best_operator_time = float("inf")
|
||||
|
||||
for idx, operator in enumerate(self.operators()):
|
||||
op_time = self.perf_operator_time(operator, *args, **kwargs)
|
||||
if op_time < best_operator_time:
|
||||
best_operator = operator
|
||||
best_operator_time = op_time
|
||||
|
||||
# print(operator, op_time)
|
||||
|
||||
return best_operator
|
||||
|
||||
def perf_operator_time(self, op: Callable, *args, **kwargs) -> float:
|
||||
fn = Functor(op, *args, **kwargs)
|
||||
time = cuda_benchmark(fn, self.num_repeated, self.num_warmup, self.dist_barrier)
|
||||
return time.gpu
|
||||
|
||||
|
||||
class OperatorRuntimeAutotuning(OperatorAutotuning):
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
|
||||
self.operator_finder = BasedKeyFinder()
|
||||
|
||||
@abc.abstractmethod
|
||||
def get_operator_key(self, *args, **kwargs):
|
||||
raise NotImplementedError()
|
||||
|
||||
def exec_best_operator(self, args, kwargs):
|
||||
key = self.get_operator_key(*args, **kwargs)
|
||||
operator = self.operator_finder.get(key, default=None)
|
||||
if operator is None:
|
||||
operator = self.perf_best_operator(*args, **kwargs)
|
||||
self.operator_finder.set(key, operator)
|
||||
|
||||
return operator(*args, **kwargs)
|
||||
|
||||
|
||||
class OperatorPreBaseRangeAutotuning(OperatorAutotuning):
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
|
||||
self.operator_finder = BasedRangeFinder()
|
||||
self._finished_autotuning = False
|
||||
|
||||
@abc.abstractmethod
|
||||
def get_operator_key(self, *args, **kwargs):
|
||||
raise NotImplementedError()
|
||||
|
||||
@abc.abstractmethod
|
||||
def generate_operator_inputs(self) -> Iterable[Tuple[Tuple, Dict]]:
|
||||
raise NotImplementedError()
|
||||
|
||||
def exec_best_operator(self, args, kwargs):
|
||||
if not self._finished_autotuning:
|
||||
self.autotuning()
|
||||
|
||||
best_op = self.operator_finder.get(self.get_operator_key(*args, **kwargs))
|
||||
return best_op(*args, **kwargs)
|
||||
|
||||
def autotuning(self):
|
||||
for op_args, op_kwargs in self.generate_operator_inputs():
|
||||
best_op = self.perf_best_operator(*op_args, **op_kwargs)
|
||||
self.operator_finder.set(
|
||||
self.get_operator_key(*op_args, **op_kwargs), best_op
|
||||
)
|
||||
|
||||
self.operator_finder.freeze()
|
||||
self._finished_autotuning = True
|
||||
# self.operator_finder.tree.show()
|
||||
Reference in New Issue
Block a user