* [fix(loader)]: correct off-by-one expert-count guard in SafeTensorLoader.load_experts
After the discovery loop, max_experts_count is the highest expert index found
(expert count - 1), and is -1 only when the key has no experts. The guard
checked == 0, which falsely rejected single-expert layers and silently returned
empty weight lists for the zero-expert case. Check == -1 instead.
Adds a CPU regression test covering the single-, zero-, and multi-expert cases.
* [test(loader)]: import loader as a top-level module in expert-count guard test
Per review feedback: add python/utils to sys.path and import loader directly
instead of the importlib.util boilerplate. Still bypasses utils/__init__.py
(and the compiled kt_kernel_ext) while keeping the import idiomatic.
* [feat](kt-kernel): AVX2 MXFP4 MoE MXFP4 dispatch
- Add AVX2 MXFP4 MoE kernel (mxfp4-moe.hpp) with 4-token M-blocking
- Wire AVX2MXFP4_MOE binding in ext_bindings.cpp
- Support TP_MOE down_proj slicing and multi-pool per-expert loading
- Add test_fp4_moe_avx2.py integration test
* [fix](kt-kernel): address PR #2010 review — memory leaks, alignment, dynamic expert update
- Track aligned_alloc pointers in AVX2_MOE_BASE::owned_aligned_allocs_ and
free them in the destructor (fixes BufferB backing memory leak on destroy).
- Track per-TP down_buf allocations in TP_MOE::tp_owned_down_bufs_ with
nullptr checks and size rounding to alignment boundary.
- Add nibble-alignment runtime check for per_tp_interm in MXFP4 TP K-split.
- Add write_weight_scale_to_buffer override to TP_MOE<AVX2_MXFP4_MOE_TP>,
enabling dynamic expert update with kt-threadpool-count>=2.
- Guard against ZeroDivisionError in test_fp4_moe_avx2.py.
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* [fix](kt-kernel): add intermediate_size parity check in MXFP4 TP flat-buffer path
The per-expert path validates that intermediate_size is even (required for
nibble-aligned FP4 addressing), but the flat-buffer path was missing this
check — an odd value would silently truncate /2 divisions, corrupting
memcpy sizes and offsets.
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* fix(avx2-moe): fix TP offset calculation and add safety checks
C1-C4: Fix incorrect TP offset calculations in load_weights()
- Per-expert mode used per_tp_interm instead of full_interm for offsets
- This caused segfault when TP > 1 due to invalid pointer arithmetic
H1-H3: Add safety checks
- H1: Validate source weight pointers are not null
- H2: Check lid index is within bounds
- H3: Check BufferB.b is not null in gemm_mxfp4
Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
* fix(avx2-moe): revert incorrect C2/C4 offset changes, keep safety checks
Reverts the incorrect offset calculation changes from previous commit.
The original per_tp_interm-based offsets were correct:
- gate/up weights are N-split (along intermediate dim)
- Each TP partition handles per_tp_interm rows
- Offset = i * per_tp_interm * hidden / 2 (not full_interm)
Keeps H1-H3 safety checks:
- H1: Validate source weight pointers are not null
- H2: Check lid index is within bounds
- H3: Check BufferB.b is not null in gemm_mxfp4
Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
* fix(avx2): copy weights to owned buffers in per-expert mode
Previously, AVX2 MXFP4 MoE per-expert mode directly pointed BufferB.b
into mmap'd safetensor data. This caused use-after-free when Python
layer releases the mmap after load_weights() returns.
Now AVX2 copies weights into owned buffers via memcpy/from_raw_mat(),
matching AMX behavior. This decouples the MoE weights from mmap lifecycle.
Changes:
- buffer_b_required_size_impl: always allocate full buffer (weights + scales)
- make_buffer_b_impl: always create full BufferB with owned storage
- Single-TP per-expert: use from_raw_mat() instead of direct pointer
- TP_MOE per-expert: add gate/up owned buffers with memcpy
- Destructor: free gate/up buffers alongside down
Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
* Revert "[fix] Add runtime AMX BF16 check to prevent SIGILL on pre-Sapphire Rapids CPUs (#2018)"
This reverts commit f1e2b82c74.
* Remove AMX tile MXFP4 kernel (GemmKernel224MXFP4)
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
---------
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
Release the SafeTensor mmap loader singleton after each layer's
load_weights() completes. The C++ engine already holds a deep copy
(cpu_infer.sync() guarantees this), so releasing the mmap handles is
safe. The next layer recreates the loader on demand.
This halves peak memory usage during model loading (e.g. DSv3.2:
1.2T -> 613G).
Based on #1966 by @poryfly — adapted to v0.6.2.post3 codebase
(adds MXFP4 support missing from the original PR).
Co-authored-by: xiongchenhui <xiongchenhui@hisense.com>
* [feat](kt-kernel): add MXFP4 MoE operator with E2M1 weights × BF16 activations
Implements AMX_FP4_MOE_TP based on the RAWINT4 (k2-moe) CRTP pattern.
FP4 E2M1 weights are nibble-packed and decoded via PSHUFB LUT, then
computed with BF16 activations using _mm512_dpbf16_ps. Supports weight-only
per-kgroup scaling (group_size=32) and tensor parallelism.
Includes a Python validation test covering uniform, alternating, ramp,
and random weight patterns.
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
* [feat](kt-kernel): adapt MXFP4 MoE backend for DeepSeek-V4-Flash (#1950)
V4-Flash routed experts ship as native MXFP4 (E2M1 nibble + ue8m0 group
scale). Expose AMXFP4_KGroup_MOE through NativeMoEWrapper, add a loader
that handles V4's `layers.{L}.ffn.experts.{i}.{w1,w3,w2}.{weight,scale}`
naming and converts ue8m0 → bf16 via a lossless bit-cast, register the
model entry, and ship an end-to-end numerical validation script.
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
* [perf](kt-kernel): MXFP4 MoE add mat-mat 4×4 tile, refine mat-vec reduce (#1957)
mat_mul_kgroup previously aliased to fp4_mat_vec_kgroup, leaving large
batches stuck on the per-token path. Implement fp4_mat_mat_kgroup as a
4×4 register tile (MB=NB=4, 16 zmm accumulators) so each PSHUFB decode
of four weight rows is reused across four tokens.
Refactor fp4_mat_vec_kgroup to accumulate four N-rows in parallel and
flush them with a new reduce4 helper, removing per-row reduce_add_ps
calls from the hot loop. Mark mxfp4_to_bf16_32 always_inline.
Add bench/bench_fp4_moe.py with --routing {balanced,concentrated} and
a backend registry so future kernels can be added without changing the
runner.
Dispatch thresholds, derived_init, GeneralMOEConfig handling,
load_weights, write_weights_to_buffer and the TP_MOE specialization are
unchanged.
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
* fix(loader): avoid uint16 lshift in ue8m0->bf16 conversion
PyTorch CPU has no lshift kernel for UInt16, so the previous
`(scale_t.to(torch.uint16) << 7)` raised NotImplementedError when
loading any V4-Flash MXFP4 routed-expert scale tensor on the host.
Switch to int32 for the shift (kernel exists) and narrow to int16
afterwards. The shifted value max is 255<<7 = 32640, well within
int16 range, so the narrow is lossless. The .view(bfloat16) bit
pattern is identical (bf16 sign bit is always 0 for ue8m0 values).
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
* docs(v4-flash): hybrid CPU/GPU recipe + bump kt-sglang submodule
Bumps third_party/sglang to kvcache-ai/sglang main (3cbd49c29) which now
contains DeepSeek V4 Flash model support + consumer-GPU (SM_120) portable
Triton/TileLang fallbacks (kt-sglang PR #38).
Adds doc/en/DeepSeek-V4-Flash.md tutorial: 8x RTX 5090 hybrid recipe with
the full launch command, OpenAI-compatible /generate + /v1/chat/completions
examples, and the kt chat CLI client.
---------
Co-authored-by: ouqingliang <1692110604@qq.com>
Co-authored-by: Claude Opus 4.7 <noreply@anthropic.com>
Add numa_nodes parameter to BaseMoEWrapper and all subclasses, allowing
users to explicitly specify which NUMA node IDs to use for subpool
mapping instead of always defaulting to sequential [0, 1, ..., N-1].
This enables running multiple KTransformers instances on different NUMA
nodes of the same machine, e.g. --kt-threadpool-count 1 --kt-numa-nodes 1
to bind to NUMA node 1. Previously this required external numactl
workarounds since subpool_numa_map was hardcoded to start from 0.
* 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
* 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>
* 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
