mirror of
https://github.com/kvcache-ai/ktransformers.git
synced 2026-04-28 03:39:48 +00:00
fcf8882075
Some checks are pending
Book-CI / test-2 (push) Waiting to run
Book-CI / test (push) Waiting to run
Book-CI / test-1 (push) Waiting to run
Deploy / deploy (macos-latest) (push) Waiting to run
Deploy / deploy (ubuntu-latest) (push) Waiting to run
Deploy / deploy (windows-latest) (push) Waiting to run
* support Kimi-K2-Thinking original weight fix amx kernel bug * update k2 avx kernel. * feat: add CPUInfer write buffer task * [feat]: add kimi k2 cpu write buffer support - Implement write_weights_to_buffer function in k2-moe.hpp for extracting GPU expert weights - Fix down (w2) weight column-wise slicing for different TP configurations - Support three TP scenarios: cpu_tp == gpu_tp, cpu_tp > gpu_tp, cpu_tp < gpu_tp - Add comprehensive test cases for weight extraction validation - Ensure compatibility with Kimi model's MoE architecture * [fix]: correct write_weight_scale_to_buffer expert offset calculation Fixed the bug in write_weight_scale_to_buffer_task where expert offsets in GPU buffers were incorrectly calculated. Changed from using per_expert_gpu sizes to using full gpu_tp sizes, ensuring correct memory layout for multi-expert scenarios. Also added benchmark scripts for k2 moe and write buffer operations, and cleaned up debug output in test files. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com> * [feat]: add write buffer wrapper * [fix] fix comment --------- Co-authored-by: ouqingliang <1692110604@qq.com> Co-authored-by: Claude <noreply@anthropic.com> |
||
|---|---|---|
| .. | ||
| .gitignore | ||
| bench_attention.py | ||
| bench_attention_torch.py | ||
| bench_k2_moe_amx.py | ||
| bench_k2_write_buffer.py | ||
| bench_linear.py | ||
| bench_linear_torch.py | ||
| bench_mla.py | ||
| bench_mlp.py | ||
| bench_mlp_torch.py | ||
| bench_moe.py | ||
| bench_moe_amx.py | ||
| bench_moe_amx_k.py | ||
| bench_moe_kernel.py | ||
| bench_moe_kernel_tiling.py | ||
| bench_moe_kml.py | ||
| bench_moe_torch.py | ||
| compare_moe_performance.py | ||
| Makefile | ||
| multi_bench_moe.py | ||
| upload-bench-json.py | ||