Commit Graph

31 Commits

Author SHA1 Message Date
mrhaoxx
07fd9328fa refactor(sft): move SFT logic into kt_kernel.sft submodule
- Create python/sft/ with 11 modules: base, amx, arch, autograd, layer,
  lora, weights, wrapper, dist_utils, config, __init__
- Move BaseSFTMoEWrapper + buffer management into sft/base.py (template
  method pattern: subclass provides _make_forward/backward_task)
- Move AMXSFTMoEWrapper into sft/amx.py (thinner, no buffer logic)
- Move from accelerate kt_moe.py: KTMoEFunction, KTMoELayerWrapper,
  MOEArchConfig, PEFT LoRA adaptation, weight extraction, wrapping
- Add KTConfig dataclass (DeepSpeed pattern: opaque config passthrough)
- Add _get_kt_config() with old→new field name compat conversion
- Rename forward_sft→forward, submit_forward_sft→submit_forward,
  sync_forward_sft→sync_forward (Python only, C++ binding names unchanged)
- Delete dump utilities from sft_moe.hpp (-526) and moe-sft-tp.hpp (-78)
- Delete experts_sft.py and utils/amx_sft.py (moved to sft/)
- Remove SFT stubs from BaseMoEWrapper (experts_base.py)
- Lazy SFT import in __init__.py and experts.py (inference isolation)
- Delete all lifecycle/debug logging (~500 lines)

Verified: Qwen3-235B 4GPU AMXBF16 training, 3 steps loss converges.
2026-04-08 23:07:41 +08:00
mrhaoxx
0f85b2a744 fix seqlen buffer size int32 overflow 2026-03-18 06:05:09 +00:00
mrhaoxx
b1a9691201 online repack 2026-02-24 11:01:53 +00:00
mrhaoxx
950f052f96 ckpt 2026-02-24 03:03:29 +00:00
mrhaoxx
8aea4b3e2a ckpt 2026-02-22 09:18:26 +00:00
mrhaoxx
cf09b3e1ca [fix]: skiplora 2026-02-17 03:06:19 +00:00
mrhaoxx
b0740436b5 [fix]: perf 2026-02-05 14:52:59 +00:00
mrhaoxx
3b2de00593 [fix]: wip 2026-02-04 05:40:47 +00:00
mrhaoxx
391cb6f79d [fix]: fix memory footprint 2026-02-02 08:34:01 +00:00
mrhaoxx
06fb3b5dbf [fix]: prequant weight load 2026-02-01 15:17:47 +00:00
mrhaoxx
9efe1317b1 [fix]: fix forward cache (maybe) 2026-01-31 21:19:46 +00:00
mrhaoxx
e1e64f7948 [fix]: fix lora grad compute 2026-01-31 15:41:55 +00:00
mrhaoxx
6a1e7c48cb [fix]: direct accumulation 2026-01-27 04:27:59 +00:00
mrhaoxx
7b62d826e4 [fix]: pinned memory causes numa issue 2026-01-26 12:04:46 +00:00
mrhaoxx
7b432f4b5a [fix]: avoid unnecessary memcpy 2026-01-26 05:34:19 +00:00
mrhaoxx
57580016ea [feat]: support async sft forward 2026-01-24 14:45:08 +00:00
mrhaoxx
451c91dce1 [fix]: nan 2026-01-23 15:30:39 +00:00
mrhaoxx
8f44a64a7a [fix]: fix memory pool 2026-01-22 19:41:07 +00:00
mrhaoxx
5c89cec5e3 [fix]: optimize job sched 2026-01-22 07:04:52 +00:00
mrhaoxx
5f6482ff50 [feat]: support skip lora 2026-01-20 05:38:11 +00:00
JimmyPeilinLi
e60f199510 [feat](kt-sft-refactor): load from huggingface safetensor file 2026-01-16 03:36:13 +00:00
JimmyPeilinLi
18ab0cb943 [feat](kt-sft-refactor): add KT-SFT to KTMoEWrapper 2026-01-15 12:29:52 +00:00
ErvinXie
d8046e1bb4 Kt minimax (#1742)
[feat]: fp8 kernel and kt-cli support
2025-12-24 15:39:44 +08:00
SCDESPERTATE
008de19e16 [fix](kt-kernel): drop the weights held in Python for loading weights operation in C++ (#1695) 2025-12-12 11:42:33 +08:00
Oql
8139c092bf Reduce CPU memory usage during large chunk prefill (Fixes #1676) (#1683)
* fix(amx): add BufferASmallKGroupImpl to fix buffer overflow in from_mat

The original BufferAKGroupImpl::from_mat writes 64 bytes per K_STEP iteration
but when K_STEP=32 (for GemmKernel224Int4SmallKGroup), this causes buffer overflow.

BufferASmallKGroupImpl overrides from_mat to write only 32 bytes per iteration.

* perf(k2-moe): optimize memory allocation with pooled buffers

- Replace per-expert buffer allocation with shared memory pools
- Dynamically assign buffer slices based on activated experts
- Add group_size inference from scale tensor shape in amx.py

* delete kimi k2 forward test

* add TODO comment for pool_count_ calculation
2025-12-08 20:19:07 +08:00
ErvinXie
71f683acec Support Native Kimi K2 Thinking (#1663)
* [feat]: fix k2 prefill

* Update Kimi-K2-Thinking.md

* Create Kimi-K2-Thinking-Native.md

* Update Kimi-K2-Thinking.md

* Update Kimi-K2-Thinking.md

* Update Kimi-K2-Thinking-Native.md

* [perf] optimize K2 MoE weight loading with per-expert pointers

- Avoid expensive torch.stack().contiguous() in Python (was ~6.6s)
- Use per-expert pointer arrays (gate_projs) instead of contiguous memory
- C++ worker pool performs parallel memcpy for TP slicing
- Add LOAD_TIME_PROFILE for load_weights timing analysis

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <noreply@anthropic.com>

---------

Co-authored-by: ouqingliang <1692110604@qq.com>
Co-authored-by: Claude <noreply@anthropic.com>
2025-12-05 21:53:05 +08:00
Jiaqi Liao
fcf8882075 [Feature] Add avx-based kimi-k2 support (#1656)
* 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>
2025-12-02 16:01:07 +08:00
ZiWei Yuan
1374b98ee5 [feat](moe_kernel): add amd blis support (int8) (#1600)
* [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>
2025-11-27 12:08:53 +08:00
Jiaqi Liao
d483147307 Fix kt-kernel compile issue (#1595)
* update install.sh

* fix import issue

* update README
2025-11-11 19:30:27 +08:00
Jiaqi Liao
94c25626dc Fix kt-kernel for new wrapper (#1588)
* update README for kt-kernel

* style: format C++ and Python code in kt-kernel

  - Format C++ files: task_queue, ext_bindings, and MoE operators
  - Format Python utility modules: amx, llamafile, and loader
  - Improve code readability and consistency
2025-11-10 21:47:34 +08:00
Jiaqi Liao
9bc00e587b Refactor KTMoEWrapper backend (#1587)
* universal backend for cpu inference
* expert defer
2025-11-10 20:26:15 +08:00