vrr/kvcache-ai-ktransformers
2
0
Fork 0
mirror of https://github.com/kvcache-ai/ktransformers.git synced 2026-04-28 03:39:48 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 4

[feat](moe_kernel): add amd blis support (int8) (#1600)
Some checks are pending
Book-CI / test (push) Waiting to run
Details
Book-CI / test-1 (push) Waiting to run
Details
Book-CI / test-2 (push) Waiting to run
Details
Deploy / deploy (macos-latest) (push) Waiting to run
Details
Deploy / deploy (ubuntu-latest) (push) Waiting to run
Details
Deploy / deploy (windows-latest) (push) Waiting to run
Details

Browse source 
* [feat]: init amd adaption

* [feat]: add blis support

* [fix]: fix setup and moe kernel warpper

* [fix](setup.py): support rebuild with cache and import kt_kernel works
fine

* [feat]: add moe_kernel converter for amd and implement the load
method(haven't tested yet)

* [feat](moe_kernel/moe.hpp): delete unused memory when using save

* [fix](moe_kernel): update PLAIN for pack

* [fix](moe_kernel): rm printf debug

* [fix](moe_kernel): skip gpu experts

* [fix](moe_kernel/moe.hpp): update include memory path

* [feat](moe_kernel/moe.hpp): support expert deferral

* [feat]: finish amd

---------

Co-authored-by: mrhaoxx <mr.haoxx@gmail.com>
This commit is contained in:
ZiWei Yuan 2025-11-27 12:08:53 +08:00 • committed by GitHub
parent fef6dd98a8
commit 1374b98ee5
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
14 changed files with 655 additions and 238 deletions
Download patch file Download diff file Expand all files Collapse all files

5
kt-kernel/python/experts.py
View file

@ -19,6 +19,7 @@ from .experts_base import BaseMoEWrapper, KExpertsCPUBuffer
# Import backend implementations
from .utils.amx import AMXMoEWrapper
from .utils.llamafile import LlamafileMoEWrapper
from .utils.moe_kernel import GeneralMoEWrapper
class KTMoEWrapper:
@ -76,7 +77,7 @@ class KTMoEWrapper:
chunked_prefill_size: Maximum prefill chunk size
cpu_save: Whether to save weights to CPU memory
max_deferred_experts_per_token: Number of experts per token to defer. Defaults to 0.
method: Backend method ("AMXINT4", "AMXINT8", "LLAMAFILE")
method: Backend method ("AMXINT4", "AMXINT8", "LLAMAFILE", "MOE_INT4", "MOE_INT8")
Returns:
An instance of the appropriate backend implementation (e.g., AMXMoEWrapper)
@ -86,6 +87,8 @@ class KTMoEWrapper:
backend_cls = AMXMoEWrapper
elif method == "LLAMAFILE":
backend_cls = LlamafileMoEWrapper
elif method in ["MOE_INT4", "MOE_INT8"]:
backend_cls = GeneralMoEWrapper
else:
raise NotImplementedError(f"Unsupported method: {method}")