feat: support Ktransformers, CPU-GPU MoE offload via FusedMoE layer

[whjthu]

Summary

• Integrate kt-kernel (KTransformers kernels, unmodified PyPI package) for heterogeneous MoE inference: routed experts run on CPU (INT4 Q4_K_M GGUF) while attention / router / shared-experts stay on GPU, enabling models larger than VRAM (e.g. 80B on a 48GB L20).
• csrc/layers/moe/kt_moe_callback.hpp (new): per-layer callback registry; lock-free invocation (callback copied out under mutex, executed GIL-safe), TOCTOU-free get().
• csrc/layers/moe/fused_moe.{hpp,cpp}: KT dispatch at forward() entry; dispatcher/runner construction skipped under use_kt_moe; explicit error when the flag is set but no callback is registered for the layer.
• csrc/layers/moe/experts/fused_moe_experts.cpp: skip w13/w2 GPU weight allocation under use_kt_moe.
• csrc/models/deepseek_v2/deepseek_v2_moe.{hpp,cpp} (+ both deepseek decoder layers): model-level KT branch for the dedicated deepseek_moe kernel path (not built on FusedMoE), with a tensor_parallel_size > 1 guard.
• csrc/pybind11/bindings.cc: _infinilm.set_kt_moe_callback / clear_kt_moe_callbacks, Python exception translation at the GIL boundary, capsule cleanup at interpreter shutdown.
• python/infinilm/kt_integration.py (new): performance-tuned callback setup — persistent staging buffers, cudaMemcpyAsync via ctypes on the default stream, double-buffered output snapshot (KT reuses its internal output buffer across calls), prefill capacity guard, rollback on partial registration, enable_graph compatibility check.
• python/infinilm/modeling_utils.py: skip routed-expert weight keys and scope check_parameters to non-expert keys under use_kt_moe (safetensors and .bin paths).

MoE models built on FusedMoE (qwen3_moe, qwen3_next) require zero model-level changes for KT support.

Motivation

MoE models whose routed-expert weights exceed GPU VRAM cannot be served at all today (e.g. Qwen3-Next-80B-A3B ≈ 45GB in INT4 vs 46GB L20, plus attention/KV overhead). KT-style CPU offload is the established deployment pattern for such models (KTransformers, sglang-kt). This PR brings that capability to InfiniLM with a minimal, non-invasive hook at the common FusedMoE layer, keeping parity with SGLang+KT throughput on the same hardware (see Benchmark section).

Type of Change

• feat — new feature / new model
• fix — bug fix
• perf — performance improvement (no behavioral change)
• refactor — code restructuring without behavior change
• test — adding or fixing tests only
• docs — documentation only
• build / ci — build system or CI configuration
• chore — tooling, formatting, or other non-code changes
• Breaking change

use_kt_moe defaults to false; all pre-existing paths are unchanged when the flag is unset.

Test Results of Involved Models on Supported Platforms (Please attach screenshots)

KT offload requires an external kt-kernel install plus an INT4 GGUF, which the four standard tests do not exercise (they run the native GPU path). Verified instead with dedicated end-to-end scripts on L20:

Model	Platform	Test	Result
Qwen3-30B-A3B (BF16 hub + INT4 GGUF experts, all-CPU)	L20 48GB, 8-core Xeon, AVX512_BF16	E2E generate, B=1/8/32, 3-round soak	Output correct , 210–217 tok/s @b=32, stability max/min = 1.003
Qwen3-Next-80B-A3B (BF16 hub + INT4 GGUF experts, all-CPU)	same	E2E generate, B=1/8/32	Correct after the companion layer_idx fix (will be fixed in another PR); 30–32 / 74 / 185–189 tok/s
Qwen3-30B-A3B, KT disabled (native path)	same	Regression	Identical output & throughput to main baseline

Benchmark / Performance Impact

Stress matrix vs SGLang+KT (sglang-kt 0.6.4 server, same machine, same INT4 GGUF, same 8 CPU threads, experts all on CPU), Qwen3-Next-80B-A3B:

N (gen len)	B	SGLang+KT (tok/s)	InfiniLM+KT (tok/s)	Delta
128	1	37.4	31.5	−15.8%
128	8	83.2	72.8	−12.5%
128	32	179.1	176.0	−1.7%
512	1	39.7	30.8	−22.4%
512	8	86.1	72.0	−16.4%
512	32	196.5	178.7	−9.1%

Soak (B=32, N=512, ×3): SGLang 196.1 ± 0.3 vs InfiniLM 177.3 ± 0.3 (both stability 1.003). Remaining gap concentrates at low concurrency (SGLang CUDA-graph + overlap-scheduler advantage); at production-level concurrency (B≥32) the two engines are within 2–9%.

Methodology: 1 warmup request, then wall-clock over concurrent batch; aggregate tokens/s counted from actual completion_tokens.

Notes for Reviewers

• Dependency: qwen3_next correctness requires the companion PR fix/qwen3-next-layer-idx (decoder layer drops layer_idx, all 48 layers construct as layer 0). Please merge that first; qwen3_moe and deepseek_v2 are unaffected by that bug.
• KT itself is unmodified — stock kt-kernel==0.6.4 from PyPI; integration is via its public KTMoEWrapper API only (same API surface SGLang uses).
• Known intentional trade-offs: enable_graph=True is rejected by setup_kt_moe (CUDA-graph capture executes the Python callback); single GPU / tensor_parallel_size=1 only (DSV2 path throws otherwise); staging buffers bound InfiniLM-side max prefill batch (guarded with a clear RuntimeError, not silent corruption).
• Follow-ups intentionally out of scope: M2 (weight-naming/layout protocol inside MoeQuantMethod), GPU/CPU expert parallel execution inside one layer (SGLang-style cpu_stream overlap), ernie4_5_vl adaptation.

CI / ChatOps

CI will be triggered manually from the Actions tab on this branch after the PR is opened.

---

Checklist

Every contributor must verify every item below before requesting
review. Tick each box only after the check has actually been performed —
do not tick speculatively. If an item truly does not apply, replace the
checkbox with N/A and briefly explain why in an inline comment.

Title, Branch, and Commits

• PR title follows Conventional Commits (e.g. feat(nvidia): …, fix(cuda/gemm): …).
• Branch name follows <type>/xxx-yyyy-zzzz where <type> matches the PR title's Conventional Commits type and words are joined with hyphens (see CONTRIBUTING.md §Branches).
• Each commit message follows Conventional Commits.
• Small PR is a single squashable commit; or, for a large PR, every commit is meaningful, well-formed, and independently reviewable (see CONTRIBUTING.md §Pull Requests).
• No stray merge commits from main — the branch is rebased cleanly on top of the current main.
• No fixup! / squash! / wip commits remain.
• Existing PR/branch/commit that followed the legacy issue format.

Scope and Design

• Changes are minimal — nothing unrelated to the stated motivation was added (CONTRIBUTING.md §Code/General).
• No dead code, commented-out blocks, debug prints, printf/std::cout/print(...) left behind, or TODO without an owner and issue link.
• No unrelated formatting churn that would obscure the diff.
• Public API changes (if any) are intentional, documented, and reflected in affected callers/tests.

General Code Hygiene (applies to all languages)

• The code is self-explanatory; comments were added only where the why is non-obvious (CONTRIBUTING.md §Code/General).
• Every modified or added file ends with a single trailing newline (CONTRIBUTING.md §Code/General).
• No trailing whitespace, tab/space mixing, or stray BOMs.
• Identifiers in comments and error messages are wrapped in backticks (e.g. the `seqlens_k` tensor) (CONTRIBUTING.md §Code/General).
• All comments and error messages are in English (CONTRIBUTING.md §Code/General).
• Comments and error messages are complete sentences — capitalized first letter, terminal punctuation — unless the language/framework convention says otherwise (CONTRIBUTING.md §Code/General; §Python).

C++ Specific (if C++ files changed)

• Code follows the Google C++ Style Guide strictly.
• Error and warning message wording follows the LLVM Coding Standards (CONTRIBUTING.md §C++).
• Constructor initializer list order matches member declaration order (CONTRIBUTING.md §C++).
• No raw new/delete; RAII / smart pointers / existing allocators are used.
• Changed files are formatted by scripts/format.py (clang-format-16, same version as CI; re-verified build + 30B E2E regression after formatting).
• No changes/reference to csrc/models/llama_legacy/.

Python Specific (if Python files changed)

• Code is PEP 8 compliant.
• Comments are complete English sentences, starting with a capital letter and ending with punctuation; Markdown backticks are used for code references (CONTRIBUTING.md §Python).
• Docstrings (if any) follow PEP 257 (CONTRIBUTING.md §Python).
• Changed files are formatted by scripts/format.py (black; re-verified syntax + E2E regression after formatting).
• No changes/reference to python/infinilm/auto_config.py.

Testing

• For any platform that could not be tested, an explicit reason is given in the table and a reviewer with access has been tagged.
• Passed single request test (examples/test_infer.py), or specify the reason for skipping.
• Passed offline performance test (examples/bench.py), or specify the reason for skipping.
• Passed sanity test (test/bench/test_benchmark.py), or specify the reason for skipping.
• Passed service test (python/infinilm/server/inference_server.py + scripts/test_perf.py), or specify the reason for skipping.

Build, CI, and Tooling

• The project builds cleanly from a fresh directory on at least one affected platform.
• CI has been triggered manually (Actions → CI on this branch), or /retest was requested.

Documentation

• README.md, CONTRIBUTING.md, or inline docs updated when behavior, build flags, or developer workflow changed.
• Any user-visible breaking change is called out explicitly under "Motivation" and in the commit/PR title with a ! or BREAKING CHANGE: footer.

Security and Safety

• No secrets, access tokens, internal URLs, customer data, or personal hardware identifiers have been committed.
• Third-party code is license-compatible and attributed.
• No unsafe pointer arithmetic, uninitialized reads, or missing bounds checks were introduced.

[wooway777]

Codex review

结论：Request changes。当前 KT 路径能在部分模型上完成初始化和生成，但存在跨模型污染、未校验的并行配置、dtype 静默错误，以及已实测的模型兼容性/正确性失败，暂不建议合入。

Blocking findings

1. [P1] 全局 KT callback 会劫持同进程中的非 KT 模型。
   csrc/layers/moe/kt_moe_callback.hpp:24 的 registry 只按 layer_idx 保存 callback；csrc/layers/moe/fused_moe.cpp:55 和 DeepSeek V2 路径在当前模型未启用 use_kt_moe 时也会查询并执行它。第二个普通模型、draft model 或并行模型可能错误调用第一个模型的 KT experts；全局 clear() 也会影响所有活跃模型。建议 callback 按 engine/model ownership 隔离，并仅在该层 skip_experts_ 为 true 时执行，注册 API 返回可管理生命周期的 handle。

2. [P1] 文档声明只支持单卡 TP1，但 API 没有校验 TP/EP/device。
   python/infinilm/kt_integration.py:92 只检查 graph；staging buffer 固定分配在 cuda:0，Qwen/FusedMoE 路径也没有像 DeepSeek 路径那样拒绝 TP>1。多 rank 会共享并竞争同一组 buffer，可能跨设备复制失败或静默算错。应在注册 callback 前明确拒绝 TP != 1、EP != 1 和非 device 0。

3. [P1] staging buffer 固定 BF16，未验证模型 dtype。
   python/infinilm/kt_integration.py:124 强制创建 torch.bfloat16 buffer，python/infinilm/kt_integration.py:151 又按每元素 2 字节裸拷贝。FP16 hidden states 不会触发 dtype 错误，而会被按 BF16 位模式解释，造成静默错误。应拒绝非 BF16 模型，或根据实际 tensor dtype 分配和拷贝。

4. [P1] DeepSeek-V2-Lite 的 LLAMAFILE/GGUF 路径不可用。
   在 kt-kernel==0.6.4、DeepSeek-V2-Lite-Chat Q4_K_M GGUF 下，setup_kt_moe() 注册阶段失败：

   ValueError: intermediate_size (1408) must be divisible by QK_K (256) for Llamafile backend
   

   PR 修改了 DeepSeek V2 MoE 且模块文档将 LLAMAFILE 描述为 verified path，因此需要提供兼容实现，或明确从支持范围排除该组合。

5. [P1] Qwen3-Next-80B KT 链路可运行，但推理正确性失败。
   使用 Qwen 官方 Qwen3-Next-80B-A3B-Instruct-Q4_K_M.gguf（48,410,988,384 bytes，GGUF v3，807 tensors），配置 TP1/EP1/BF16、graph/prefix caching 关闭、512 experts、top-k 10、intermediate 512。48 个 KT callbacks 均注册成功，但：

   • raw prompt 128 tokens：输出为 , 加连续空格；
   • chat template 128 tokens：输出循环 . 1. 2. ... 9.，没有回答提示词；
   • 相同原生模型不使用 KT 时可正常生成连贯文本。
     这说明 wrapper 初始化成功不等于数值正确，Qwen3-Next 当前不能视为受支持。

Validation matrix

环境：172.22.162.61 / container pepe / branch feat/kt-support / head e73c49d6 / kt-kernel==0.6.4。_infinilm 已重新编译安装，测试后 GPU 无残留进程，仓库工作区干净。

Mode	Model	Result
KT	Qwen3-30B-A3B-Instruct-2507 Q4_K_M	PASS，注册 48 层并生成正常文本
KT	DeepSeek-V2-Lite-Chat Q4_K_M	FAIL，1408/QK_K 对齐限制
KT	Qwen3-Next-80B-A3B-Instruct Q4_K_M	FAIL，生成退化为标点/数字/空格循环
KT	Qwen3-235B、Qwen3.5-122B	未找到匹配 GGUF，按测试约束未下载
Native	DeepSeek-V2-Lite-Chat，batch16 1024/1024 TP4 MLA	PASS，45.48 s
Native	Qwen3-30B-A3B，batch16 512/512 TP4 legacy MoE	PASS，24.10 s
Native	Qwen3-235B-A22B，batch16 512/512 TP8 EP8 graph	PASS，53.99 s
Native	Qwen3-Next-80B-A3B，batch16 1024/1024 TP8 graph	PASS，22.19 s
Native	Qwen3.5-122B-A10B，batch16 1024/1024 TP8 EP8 graph	PASS，37.99 s
Multimodal native	Qwen3.5-0.8B	PASS，9.85 s
Multimodal native	Qwen3.6-27B	PASS，32.44 s
Multimodal native	Qwen3.5-122B-A10B	PASS，49.50 s

三组多模态均通过 examples/test_infer.py 使用指定图片、batch 16 和 1024-token 输出完成，输出能正确识别 Batman 形象及相关标识。

Current usage constraints

当前没有 CLI 开关。使用者需要在构造 LLM 前通过 shadow config.json 设置 "use_kt_moe": true，以 TP1、EP1、BF16、无 graph、无 prefix caching 构造模型，再调用 setup_kt_moe(...) 注册每个 MoE 层；max_tokens 必须覆盖最大 prefill batch。结束时必须调用 _infinilm.clear_kt_moe_callbacks() 后关闭模型。

建议合入前至少补充：callback ownership/lifecycle、TP/EP/device/dtype 强校验、DeepSeek/Qwen3-Next 正确性修复，以及覆盖双模型同进程、非 KT 回归和各支持架构的自动化测试。

— Codex