[feat](kt-lora): add end-to-end Qwen3.5 MoE KT LoRA serving workflow (#2031)
Some checks are pending
Book-CI / test (push) Waiting to run
Book-CI / test-1 (push) Waiting to run
Book-CI / test-2 (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
Release sglang-kt to PyPI / Build sglang-kt wheel (push) Waiting to run
Release sglang-kt to PyPI / Publish sglang-kt to PyPI (push) Blocked by required conditions

* [feat](kt-lora): add KT expert LoRA adapter serving

* [feat]: pin Qwen3.5 non-expert LoRA support

* [feat](kt-lora): add merged SGLang adapter workflow

Document the KT SFT to SGLang serving loop and extend the converter with optional split outputs so users can serve one merged adapter while retaining debug-friendly expert/non-expert artifacts.

Co-authored-by: Cursor <cursoragent@cursor.com>

* [fix](kt-lora): validate adapter conversion

Co-authored-by: Cursor <cursoragent@cursor.com>

---------

Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
Jiaheng Dai 2026-06-05 16:57:14 +08:00 committed by GitHub
parent d41f569e84
commit c9a915e6ac
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
16 changed files with 2435 additions and 656 deletions

View file

@ -7,6 +7,8 @@
# Tutorial
- [kt-sft part](en/SFT/README.md)
- [KT-FT Fine-Tuning and Inference Loop](en/SFT/Qwen3.5-SGLang-LoRA-Serving.md)
- [KT-FT 微调推理闭环](zh/Qwen3.5-SGLang-LoRA-Serving_zh.md)
- [Injection Tutorial](en/SFT/injection_tutorial.md)
- [kt-sft developer tech notes](en/SFT/KTransformers-Fine-Tuning_Developer-Technical-Notes.md)
- [DPO tutorial](en/SFT/DPO_tutorial.md)

View file

@ -0,0 +1,221 @@
# KT-FT Fine-Tuning and Inference Loop
Last updated: 2026-06-01
This guide documents the current KT-FT loop for Qwen3.5 MoE: train with KT SFT, convert the output once, and serve the fine-tuned result through SGLang with a single merged adapter path.
```text
KT SFT raw output
-> convert_kt_to_sglang_adapter.py
-> <MERGED_ADAPTER_DIR>
-> sglang --lora-paths <name>=<MERGED_ADAPTER_DIR>
-> server auto-splits expert / non-expert internally
-> request model=<served_model>:<name>
```
Training-side KT SFT docs remain separate. This page focuses on the bridge from trained LoRA artifacts to online inference.
## 1. Scope
Current supported and validated workflow:
- Base model: Qwen3.5 MoE, for example `Qwen3.5-35B-A3B`
- KT expert weights: AMX/BF16 SFT-compatible KT CPU expert path
- User-facing serving input: one converted merged adapter directory
- Runtime split: expert LoRA goes to the KT CPU expert path; non-expert LoRA goes to SGLang's LoRA manager. This split happens automatically at server startup.
## 2. Artifacts At Each Stage
### Raw KT SFT output
After LLaMA-Factory + KT training, the output directory contains two LoRA artifacts:
```text
<KT_SFT_OUTPUT_DIR>/
adapter_model.safetensors # non-expert LoRA
fused_expert_lora.safetensors # expert LoRA in KT fused format
adapter_config.json
```
Do not pass this raw directory directly to SGLang serving.
### Converted merged adapter
Run the converter once to produce the serving input:
```text
<MERGED_ADAPTER_DIR>/
adapter_config.json
adapter_model.safetensors
```
This merged directory contains both expert and non-expert LoRA tensors in one PEFT-style adapter. Pass only this directory to `--lora-paths`.
## 3. Convert Once
```bash
python kt-kernel/scripts/convert_kt_to_sglang_adapter.py \
<KT_SFT_OUTPUT_DIR> \
<MERGED_ADAPTER_DIR> \
--base-model-name-or-path /path/to/Qwen3.5-35B-A3B \
--overwrite
```
Example:
```bash
python kt-kernel/scripts/convert_kt_to_sglang_adapter.py \
saves/KT_FT_qwen35B_Moe_nekoqa_eod_240 \
saves/KT_FT_qwen35B_Moe_nekoqa_eod_240_sglang \
--base-model-name-or-path /mnt/data3/models/Qwen3.5-35B-A3B \
--overwrite
```
The converter reads `fused_expert_lora.safetensors` and the existing non-expert `adapter_model.safetensors`, then writes one merged adapter directory.
Optional split outputs for debugging:
```bash
python kt-kernel/scripts/convert_kt_to_sglang_adapter.py \
<KT_SFT_OUTPUT_DIR> \
<MERGED_ADAPTER_DIR> \
--base-model-name-or-path /path/to/Qwen3.5-35B-A3B \
--expert-output-dir <EXPERT_ADAPTER_DIR> \
--nonexpert-output-dir <NONEXPERT_ADAPTER_DIR> \
--overwrite
```
For normal serving, only `<MERGED_ADAPTER_DIR>` is needed.
## 4. Launch SGLang
Use the KTransformers SGLang fork from this repository and point `PYTHONPATH` at both `kt-kernel/python` and `third_party/sglang/python`.
```bash
cd /path/to/ktransformers
PYTHONPATH=/path/to/ktransformers/kt-kernel/python:/path/to/ktransformers/third_party/sglang/python:$PYTHONPATH \
python -m sglang.launch_server \
--host 127.0.0.1 \
--port 30006 \
--model-path /path/to/Qwen3.5-35B-A3B \
--tokenizer-path /path/to/Qwen3.5-35B-A3B \
--kt-weight-path /path/to/Qwen3.5-35B-A3B-AMXINT4 \
--kt-method AMXINT4 \
--kt-cpuinfer 60 \
--kt-threadpool-count 2 \
--kt-numa-nodes 0 1 \
--kt-num-gpu-experts 0 \
--attention-backend flashinfer \
--trust-remote-code \
--mem-fraction-static 0.98 \
--chunked-prefill-size 4096 \
--max-running-requests 2 \
--max-total-tokens 32000 \
--served-model-name qwen3.5-kt-ft \
--enable-mixed-chunk \
--tensor-parallel-size 4 \
--enable-p2p-check \
--disable-cuda-graph \
--disable-custom-all-reduce \
--enable-lora \
--lora-backend triton \
--lora-paths qwen35b_neko=/path/to/KT_FT_qwen35B_Moe_nekoqa_eod_240_sglang \
--log-level info
```
Important points:
- Pass only one merged adapter through `--lora-paths`.
- Do not also pass `--kt-expert-lora-path` in the normal user workflow.
- At startup, the server detects the merged KT MoE adapter, splits it internally, and writes runtime cache directories under `$TMPDIR/sglang_kt_lora_cache/` (or `$SGLANG_KT_LORA_CACHE_DIR` if set).
- Prefer `--lora-backend triton` for Qwen3.5 full-LoRA generation.
Current constraints:
- single merged KT composite adapter only
- `--kt-num-gpu-experts 0`
- do not enable `--kt-enable-dynamic-expert-update`
- do not use `--kt-gpu-prefill-token-threshold`
- use an AMX/BF16 SFT-compatible KT method such as `AMXINT4`, `AMXINT8`, `AMXBF16`, or `BF16`
## 5. Request Semantics
The OpenAI-compatible request `model` field uses names, not paths.
```text
--served-model-name qwen3.5-kt-ft
--lora-paths qwen35b_neko=/path/to/merged_adapter
```
Request behavior in the current single-adapter implementation:
```text
model=qwen3.5-kt-ft
=> base + KT expert LoRA
model=qwen3.5-kt-ft:qwen35b_neko
=> base + KT expert LoRA + SGLang non-expert LoRA
```
The suffix after `:` must match the left-side name in `--lora-paths`.
## 6. Smoke Test
```bash
curl -sS http://127.0.0.1:30006/v1/chat/completions \
-H 'Content-Type: application/json' \
-d '{
"model": "qwen3.5-kt-ft:qwen35b_neko",
"messages": [{"role": "user", "content": "我回来了,你在干嘛?"}],
"temperature": 0.7,
"max_tokens": 160,
"chat_template_kwargs": {"enable_thinking": false}
}'
```
Startup logs should include lines similar to:
```text
Prepared merged KT LoRA adapter ... for runtime: expert=... nonexpert=...
Loaded KT expert LoRA for layer ...
Using triton as backend of LoRA kernels.
```
## 7. Advanced: Manual Split Serving
The older split-runtime contract is still available for debugging:
```bash
--kt-expert-lora-path <EXPERT_ADAPTER_DIR> \
--enable-lora \
--lora-paths <NONEXPERT_LORA_NAME>=<NONEXPERT_ADAPTER_DIR>
```
This is not the recommended user-facing path. Normal users should pass one merged adapter directory through `--lora-paths` only.
## 8. Troubleshooting
### `Got LoRA adapter that has never been loaded: lora0`
The adapter name in the request must match the left side of `--lora-paths`. If you launched with `qwen35b_neko=...`, request `model=qwen3.5-kt-ft:qwen35b_neko`, not `:lora0`.
### No visible adapter effect
Make sure you are serving the intended merged adapter directory. For example, use the Neko adapter at `..._nekoqa_eod_240_sglang`, not a generic sanity adapter such as `..._Moe_sglang`.
### `connection refused`
Check that the server is listening on the port you curl, and remember the example above binds to `127.0.0.1`, not `0.0.0.0`.
### Server resolves upstream SGLang instead of this checkout
```bash
python - <<'PY'
import inspect
import sglang.srt.models.qwen3_5 as qwen3_5
print(inspect.getfile(qwen3_5))
PY
```
The path should come from this repository's `third_party/sglang`.

View file

@ -2,6 +2,7 @@
- [v0.6.1 Quick Start](./KTransformers-Fine-Tuning_Quick-Start.md)
- [Fine-Tuning User Guide](./KTransformers-Fine-Tuning_User-Guide.md)
- [KT-FT Fine-Tuning and Inference Loop](./Qwen3.5-SGLang-LoRA-Serving.md)
- [Developer Technical Notes](./KTransformers-Fine-Tuning_Developer-Technical-Notes.md)
- [DPO Tutorial](./DPO_tutorial.md)
- [Injection Tutorial](./injection_tutorial.md)

View file

@ -0,0 +1,221 @@
# KT-FT 微调推理闭环
最后更新2026-06-01
本文档描述当前 Qwen3.5 MoE 的 KT-FT 闭环:用 KT SFT 完成微调,转换一次输出,再通过 SGLang 用单个 merged adapter path 把微调结果服务化。
```text
KT SFT 原始输出
-> convert_kt_to_sglang_adapter.py
-> <MERGED_ADAPTER_DIR>
-> sglang --lora-paths <name>=<MERGED_ADAPTER_DIR>
-> server 内部自动拆分 expert / non-expert
-> 请求 model=<served_model>:<name>
```
训练侧 KT SFT 文档仍然独立维护;本文重点说明从已训练 LoRA artifacts 到在线推理的连接部分。
## 1. 范围
当前已验证路径:
- 基座模型Qwen3.5 MoE例如 `Qwen3.5-35B-A3B`
- KT expert 权重AMX/BF16 SFT 兼容的 KT CPU expert 路径
- 用户侧 serving 输入:一个 converted merged adapter 目录
- Runtime 内部仍会 splitexpert LoRA 走 KT CPU expert pathnon-expert LoRA 走 SGLang LoRA manager但这一步对用户不可见
## 2. 各阶段产物
### 原始 KT SFT 输出
LLaMA-Factory + KT 训练完成后,输出目录里有两个 LoRA 文件:
```text
<KT_SFT_OUTPUT_DIR>/
adapter_model.safetensors # non-expert LoRA
fused_expert_lora.safetensors # KT fused expert LoRA
adapter_config.json
```
不要把 raw 训练目录直接传给 SGLang serving。
### Convert 后的 merged adapter
converter 一次性生成 serving 输入:
```text
<MERGED_ADAPTER_DIR>/
adapter_config.json
adapter_model.safetensors
```
这个 merged 目录同时包含 expert 和 non-expert LoRA。正常 serving 只需要传这一个目录。
## 3. 转换一次
```bash
python kt-kernel/scripts/convert_kt_to_sglang_adapter.py \
<KT_SFT_OUTPUT_DIR> \
<MERGED_ADAPTER_DIR> \
--base-model-name-or-path /path/to/Qwen3.5-35B-A3B \
--overwrite
```
示例:
```bash
python kt-kernel/scripts/convert_kt_to_sglang_adapter.py \
saves/KT_FT_qwen35B_Moe_nekoqa_eod_240 \
saves/KT_FT_qwen35B_Moe_nekoqa_eod_240_sglang \
--base-model-name-or-path /mnt/data3/models/Qwen3.5-35B-A3B \
--overwrite
```
converter 会读取 `fused_expert_lora.safetensors` 和已有的 non-expert `adapter_model.safetensors`,写出一个 merged adapter 目录。
如需调试,也可以额外输出 split 目录:
```bash
python kt-kernel/scripts/convert_kt_to_sglang_adapter.py \
<KT_SFT_OUTPUT_DIR> \
<MERGED_ADAPTER_DIR> \
--base-model-name-or-path /path/to/Qwen3.5-35B-A3B \
--expert-output-dir <EXPERT_ADAPTER_DIR> \
--nonexpert-output-dir <NONEXPERT_ADAPTER_DIR> \
--overwrite
```
正常 serving 只需要 `<MERGED_ADAPTER_DIR>`
## 4. 启动 SGLang
请使用本仓库的 KTransformers SGLang fork并把 `PYTHONPATH` 指向 `kt-kernel/python``third_party/sglang/python`
```bash
cd /path/to/ktransformers
PYTHONPATH=/path/to/ktransformers/kt-kernel/python:/path/to/ktransformers/third_party/sglang/python:$PYTHONPATH \
python -m sglang.launch_server \
--host 127.0.0.1 \
--port 30006 \
--model-path /path/to/Qwen3.5-35B-A3B \
--tokenizer-path /path/to/Qwen3.5-35B-A3B \
--kt-weight-path /path/to/Qwen3.5-35B-A3B-AMXINT4 \
--kt-method AMXINT4 \
--kt-cpuinfer 60 \
--kt-threadpool-count 2 \
--kt-numa-nodes 0 1 \
--kt-num-gpu-experts 0 \
--attention-backend flashinfer \
--trust-remote-code \
--mem-fraction-static 0.98 \
--chunked-prefill-size 4096 \
--max-running-requests 2 \
--max-total-tokens 32000 \
--served-model-name qwen3.5-kt-ft \
--enable-mixed-chunk \
--tensor-parallel-size 4 \
--enable-p2p-check \
--disable-cuda-graph \
--disable-custom-all-reduce \
--enable-lora \
--lora-backend triton \
--lora-paths qwen35b_neko=/path/to/KT_FT_qwen35B_Moe_nekoqa_eod_240_sglang \
--log-level info
```
要点:
- 用户只需要传一个 merged adapter`--lora-paths <name>=<MERGED_ADAPTER_DIR>`
- 正常 workflow 不要再额外传 `--kt-expert-lora-path`
- server 启动时会自动识别 merged KT MoE adapter并在 `$TMPDIR/sglang_kt_lora_cache/`(或 `$SGLANG_KT_LORA_CACHE_DIR`)下生成 runtime cache
- Qwen3.5 full LoRA 生成优先使用 `--lora-backend triton`
当前限制:
- 只支持单个 merged KT composite adapter
- `--kt-num-gpu-experts 0`
- 不启用 `--kt-enable-dynamic-expert-update`
- 不使用 `--kt-gpu-prefill-token-threshold`
- 使用 AMX/BF16 SFT 兼容 KT method例如 `AMXINT4``AMXINT8``AMXBF16``BF16`
## 5. 请求语义
OpenAI-compatible 请求里的 `model` 字段用 name不用 path。
```text
--served-model-name qwen3.5-kt-ft
--lora-paths qwen35b_neko=/path/to/merged_adapter
```
当前 single-adapter 实现的请求语义:
```text
model=qwen3.5-kt-ft
=> base + KT expert LoRA
model=qwen3.5-kt-ft:qwen35b_neko
=> base + KT expert LoRA + SGLang non-expert LoRA
```
冒号后的 adapter 名必须和 `--lora-paths` 左侧注册名一致。
## 6. Smoke Test
```bash
curl -sS http://127.0.0.1:30006/v1/chat/completions \
-H 'Content-Type: application/json' \
-d '{
"model": "qwen3.5-kt-ft:qwen35b_neko",
"messages": [{"role": "user", "content": "我回来了,你在干嘛?"}],
"temperature": 0.7,
"max_tokens": 160,
"chat_template_kwargs": {"enable_thinking": false}
}'
```
启动日志里应能看到类似输出:
```text
Prepared merged KT LoRA adapter ... for runtime: expert=... nonexpert=...
Loaded KT expert LoRA for layer ...
Using triton as backend of LoRA kernels.
```
## 7. 高级:手动 split serving
旧 split runtime 仍可用于调试:
```bash
--kt-expert-lora-path <EXPERT_ADAPTER_DIR> \
--enable-lora \
--lora-paths <NONEXPERT_LORA_NAME>=<NONEXPERT_ADAPTER_DIR>
```
这不是推荐的用户路径。正常用户只需要通过 `--lora-paths` 传一个 merged adapter 目录。
## 8. Troubleshooting
### `Got LoRA adapter that has never been loaded: lora0`
请求里的 adapter 名必须和 `--lora-paths` 左侧一致。如果启动时写的是 `qwen35b_neko=...`,请求应使用 `model=qwen3.5-kt-ft:qwen35b_neko`,而不是 `:lora0`
### 看不出 adapter 效果
确认 serving 用的是目标 merged adapter。例如 Neko 风格应使用 `..._nekoqa_eod_240_sglang`,而不是通用 sanity adapter `..._Moe_sglang`
### `connection refused`
确认 server 监听的端口与 curl 一致;上面的示例绑定的是 `127.0.0.1`,不是 `0.0.0.0`
### Server 解析到了上游 SGLang而不是当前 checkout
```bash
python - <<'PY'
import inspect
import sglang.srt.models.qwen3_5 as qwen3_5
print(inspect.getfile(qwen3_5))
PY
```
路径应来自本仓库的 `third_party/sglang`

File diff suppressed because it is too large Load diff

View file

@ -119,20 +119,20 @@ class AVX2_MOE_BASE {
down_ba_.push_back(make_buffer_a(config_.max_len, config_.intermediate_size, nullptr));
down_bc_.push_back(make_buffer_c(config_.max_len, config_.hidden_size, nullptr));
void* gate_bb_ptr =
std::aligned_alloc(64, (buffer_b_required_size(config_.intermediate_size, config_.hidden_size) + 63) & ~63ULL);
void* gate_bb_ptr = std::aligned_alloc(
64, (buffer_b_required_size(config_.intermediate_size, config_.hidden_size) + 63) & ~63ULL);
if (!gate_bb_ptr) throw std::runtime_error("aligned_alloc failed for gate BufferB");
owned_aligned_allocs_.push_back(gate_bb_ptr);
gate_bb_.push_back(make_buffer_b(config_.intermediate_size, config_.hidden_size, gate_bb_ptr));
void* up_bb_ptr =
std::aligned_alloc(64, (buffer_b_required_size(config_.intermediate_size, config_.hidden_size) + 63) & ~63ULL);
void* up_bb_ptr = std::aligned_alloc(
64, (buffer_b_required_size(config_.intermediate_size, config_.hidden_size) + 63) & ~63ULL);
if (!up_bb_ptr) throw std::runtime_error("aligned_alloc failed for up BufferB");
owned_aligned_allocs_.push_back(up_bb_ptr);
up_bb_.push_back(make_buffer_b(config_.intermediate_size, config_.hidden_size, up_bb_ptr));
void* down_bb_ptr =
std::aligned_alloc(64, (buffer_b_required_size(config_.hidden_size, config_.intermediate_size) + 63) & ~63ULL);
void* down_bb_ptr = std::aligned_alloc(
64, (buffer_b_required_size(config_.hidden_size, config_.intermediate_size) + 63) & ~63ULL);
if (!down_bb_ptr) throw std::runtime_error("aligned_alloc failed for down BufferB");
owned_aligned_allocs_.push_back(down_bb_ptr);
down_bb_.push_back(make_buffer_b(config_.hidden_size, config_.intermediate_size, down_bb_ptr));
@ -234,23 +234,28 @@ class AVX2_MOE_BASE {
size_t max_m = (m_local_num_[i] + M_STEP - 1) / M_STEP * M_STEP;
gate_up_ba_[i]->max_m = max_m;
gate_up_ba_[i]->set_data(gate_up_ba_pool_ptr);
gate_up_ba_pool_ptr = (void*)((uintptr_t)gate_up_ba_pool_ptr + align64(buffer_a_required_size(max_m, config_.hidden_size)));
gate_up_ba_pool_ptr =
(void*)((uintptr_t)gate_up_ba_pool_ptr + align64(buffer_a_required_size(max_m, config_.hidden_size)));
gate_bc_[i]->max_m = max_m;
gate_bc_[i]->set_data(gate_bc_pool_ptr);
gate_bc_pool_ptr = (void*)((uintptr_t)gate_bc_pool_ptr + align64(buffer_c_required_size(max_m, config_.intermediate_size)));
gate_bc_pool_ptr =
(void*)((uintptr_t)gate_bc_pool_ptr + align64(buffer_c_required_size(max_m, config_.intermediate_size)));
up_bc_[i]->max_m = max_m;
up_bc_[i]->set_data(up_bc_pool_ptr);
up_bc_pool_ptr = (void*)((uintptr_t)up_bc_pool_ptr + align64(buffer_c_required_size(max_m, config_.intermediate_size)));
up_bc_pool_ptr =
(void*)((uintptr_t)up_bc_pool_ptr + align64(buffer_c_required_size(max_m, config_.intermediate_size)));
down_ba_[i]->max_m = max_m;
down_ba_[i]->set_data(down_ba_pool_ptr);
down_ba_pool_ptr = (void*)((uintptr_t)down_ba_pool_ptr + align64(buffer_a_required_size(max_m, config_.intermediate_size)));
down_ba_pool_ptr =
(void*)((uintptr_t)down_ba_pool_ptr + align64(buffer_a_required_size(max_m, config_.intermediate_size)));
down_bc_[i]->max_m = max_m;
down_bc_[i]->set_data(down_bc_pool_ptr);
down_bc_pool_ptr = (void*)((uintptr_t)down_bc_pool_ptr + align64(buffer_c_required_size(max_m, config_.hidden_size)));
down_bc_pool_ptr =
(void*)((uintptr_t)down_bc_pool_ptr + align64(buffer_c_required_size(max_m, config_.hidden_size)));
}
auto direct_or_pool = [&](int count, auto&& fn) {
@ -330,9 +335,8 @@ class AVX2_MOE_BASE {
__m256 weight = _mm256_set1_ps(weights[i * k + j]);
__m256 d0, d1;
avx2::load_16xbf16_to_2x8xfp32(
m_local_down_output_ptr_[expert_ids[i * k + j]] +
m_local_pos_[i][j] * config_.hidden_size + e,
&d0, &d1);
m_local_down_output_ptr_[expert_ids[i * k + j]] + m_local_pos_[i][j] * config_.hidden_size + e, &d0,
&d1);
x0 = _mm256_fmadd_ps(d0, weight, x0);
x1 = _mm256_fmadd_ps(d1, weight, x1);
}
@ -381,19 +385,23 @@ class AVX2_MOE_BASE {
gate_bc_[expert_idx]->max_m = max_m;
gate_bc_[expert_idx]->set_data(gate_bc_pool_ptr);
gate_bc_pool_ptr = (void*)((uintptr_t)gate_bc_pool_ptr + align64(buffer_c_required_size(max_m, config_.intermediate_size)));
gate_bc_pool_ptr =
(void*)((uintptr_t)gate_bc_pool_ptr + align64(buffer_c_required_size(max_m, config_.intermediate_size)));
up_bc_[expert_idx]->max_m = max_m;
up_bc_[expert_idx]->set_data(up_bc_pool_ptr);
up_bc_pool_ptr = (void*)((uintptr_t)up_bc_pool_ptr + align64(buffer_c_required_size(max_m, config_.intermediate_size)));
up_bc_pool_ptr =
(void*)((uintptr_t)up_bc_pool_ptr + align64(buffer_c_required_size(max_m, config_.intermediate_size)));
down_ba_[expert_idx]->max_m = max_m;
down_ba_[expert_idx]->set_data(down_ba_pool_ptr);
down_ba_pool_ptr = (void*)((uintptr_t)down_ba_pool_ptr + align64(buffer_a_required_size(max_m, config_.intermediate_size)));
down_ba_pool_ptr =
(void*)((uintptr_t)down_ba_pool_ptr + align64(buffer_a_required_size(max_m, config_.intermediate_size)));
down_bc_[expert_idx]->max_m = max_m;
down_bc_[expert_idx]->set_data(down_bc_pool_ptr);
down_bc_pool_ptr = (void*)((uintptr_t)down_bc_pool_ptr + align64(buffer_c_required_size(max_m, config_.hidden_size)));
down_bc_pool_ptr =
(void*)((uintptr_t)down_bc_pool_ptr + align64(buffer_c_required_size(max_m, config_.hidden_size)));
}
// Pack input into BufferA for each activated expert
@ -403,7 +411,8 @@ class AVX2_MOE_BASE {
size_t max_m = (qlen + M_STEP - 1) / M_STEP * M_STEP;
gate_up_ba_[expert_idx]->max_m = max_m;
gate_up_ba_[expert_idx]->set_data(gate_up_ba_pool_ptr);
gate_up_ba_pool_ptr = (void*)((uintptr_t)gate_up_ba_pool_ptr + align64(buffer_a_required_size(max_m, config_.hidden_size)));
gate_up_ba_pool_ptr =
(void*)((uintptr_t)gate_up_ba_pool_ptr + align64(buffer_a_required_size(max_m, config_.hidden_size)));
gate_up_ba_[expert_idx]->from_mat(qlen, (ggml_bf16_t*)input, 0, 1);
}
@ -458,8 +467,7 @@ class AVX2_MOE_BASE {
__m256 weight = _mm256_set1_ps(weights[j]);
__m256 d0, d1;
avx2::load_16xbf16_to_2x8xfp32(
m_local_down_output_ptr_[expert_ids[j]] + m_local_pos_[0][j] * config_.hidden_size + e,
&d0, &d1);
m_local_down_output_ptr_[expert_ids[j]] + m_local_pos_[0][j] * config_.hidden_size + e, &d0, &d1);
x0 = _mm256_fmadd_ps(d0, weight, x0);
x1 = _mm256_fmadd_ps(d1, weight, x1);
}

View file

@ -444,7 +444,8 @@ class AVX2_MXFP4_MOE_TP : public AVX2_MOE_BASE<T, AVX2_MXFP4_MOE_TP<T>> {
// H2: Bounds check (already validated in weight loading, but be safe)
if (lid >= config_.gate_scales[0].size()) return;
if (config_.gate_scales[0][lid] == nullptr || config_.up_scales[0][lid] == nullptr ||
config_.down_scales[0][lid] == nullptr) return;
config_.down_scales[0][lid] == nullptr)
return;
size_t scale_elem_count = ((size_t)config_.hidden_size * config_.intermediate_size) / group_size;
// tp_part_idx == 0 guaranteed here, so offset is 0
convert_or_copy(gate_bb_[expert_idx]->d, (const ggml_bf16_t*)config_.gate_scales[0][lid], scale_elem_count);
@ -678,7 +679,8 @@ class TP_MOE<AVX2_MXFP4_MOE_TP<K>> : public TP_MOE<AVX2_MOE_BASE<K, AVX2_MXFP4_M
auto subpool = pool->get_subpool(i);
subpool->do_work_stealing_job(
tpc.expert_num, nullptr,
[&, i, per_tp_interm, full_interm, gate_buf, up_buf, down_buf, gate_up_wt_per_expert, down_wt_per_expert](int eid) {
[&, i, per_tp_interm, full_interm, gate_buf, up_buf, down_buf, gate_up_wt_per_expert,
down_wt_per_expert](int eid) {
if (tpc.should_skip_expert(eid)) return;
uint64_t lid = expert_map(physical_to_logical_map, eid);
@ -725,7 +727,8 @@ class TP_MOE<AVX2_MXFP4_MOE_TP<K>> : public TP_MOE<AVX2_MOE_BASE<K, AVX2_MXFP4_M
// down scales: K-split, non-contiguous → per-row convert
if (config.down_scales[0][lid] == nullptr) {
throw std::runtime_error("TP_MOE load_weights: null down_scale pointer for expert " + std::to_string(lid));
throw std::runtime_error("TP_MOE load_weights: null down_scale pointer for expert " +
std::to_string(lid));
}
int full_interm_groups = full_interm / group_size;
int per_tp_groups = per_tp_interm / group_size;
@ -743,8 +746,9 @@ class TP_MOE<AVX2_MXFP4_MOE_TP<K>> : public TP_MOE<AVX2_MOE_BASE<K, AVX2_MXFP4_M
pool->dispense_backend()->do_numa_job([&, this](int i) {
auto& tpc = tps[i]->config_;
if (tpc.intermediate_size % 2 != 0)
throw std::runtime_error("MXFP4 TP flat-buffer: intermediate_size must be even for nibble-aligned addressing, got " +
std::to_string(tpc.intermediate_size));
throw std::runtime_error(
"MXFP4 TP flat-buffer: intermediate_size must be even for nibble-aligned addressing, got " +
std::to_string(tpc.intermediate_size));
size_t weight_elem_count = (size_t)tpc.intermediate_size * tpc.hidden_size;
size_t scales_elem_count = ((size_t)tpc.hidden_size / group_size) * tpc.intermediate_size;
tpc.gate_proj = new uint8_t[(tpc.expert_num * weight_elem_count) / 2];
@ -819,8 +823,7 @@ class TP_MOE<AVX2_MXFP4_MOE_TP<K>> : public TP_MOE<AVX2_MOE_BASE<K, AVX2_MXFP4_M
this->weights_loaded = true;
}
void write_weight_scale_to_buffer(int gpu_tp_count, int expert_id,
const std::vector<uintptr_t>& w13_weight_ptrs,
void write_weight_scale_to_buffer(int gpu_tp_count, int expert_id, const std::vector<uintptr_t>& w13_weight_ptrs,
const std::vector<uintptr_t>& w13_scale_ptrs,
const std::vector<uintptr_t>& w2_weight_ptrs,
const std::vector<uintptr_t>& w2_scale_ptrs) {

View file

@ -8,7 +8,6 @@
#ifndef CPUINFER_OPERATOR_MOE_SFT_TP_HPP
#define CPUINFER_OPERATOR_MOE_SFT_TP_HPP
#include <immintrin.h>
#include <algorithm>
@ -266,6 +265,9 @@ class TP_MOE_SFT : public TP_MOE<T> {
if (!config.gate_projs.empty()) {
// Pre-quantized per-NUMA weights (INT8/INT4 with separate scales)
printf("TP_MOE_SFT: Pre-quantized per-NUMA mode (gate_projs path)\n");
for (int i = 0; i < tp_count; i++) {
tps[i]->set_physical_to_logical_map(config.physical_to_logical_map);
}
pool->dispense_backend()->do_numa_job([this](int numa_id) { tps[numa_id]->load_weights(); });
// Check if pre-quantized backward weights are available
@ -391,6 +393,7 @@ class TP_MOE_SFT : public TP_MOE<T> {
// Step 2: Set weight pointers BEFORE load_weights (Bug #24 fix)
for (int i = 0; i < tp_count; i++) {
tps[i]->set_physical_to_logical_map(config.physical_to_logical_map);
tps[i]->set_weight_pointers_for_forward(temp_gate[i], temp_up[i], temp_down[i]);
}
@ -401,7 +404,6 @@ class TP_MOE_SFT : public TP_MOE<T> {
if (!config.share_backward_bb) {
tps[i]->prepare_bwd(temp_gate[i], temp_up[i], temp_down[i]);
}
tps[i]->set_physical_to_logical_map(config.physical_to_logical_map);
}
for (int i = 0; i < tp_count; i++) {
@ -411,6 +413,9 @@ class TP_MOE_SFT : public TP_MOE<T> {
}
} else {
// Other loading methods (from loader or file)
for (int i = 0; i < tp_count; i++) {
tps[i]->set_physical_to_logical_map(config.physical_to_logical_map);
}
pool->dispense_backend()->do_numa_job([this](int numa_id) { tps[numa_id]->load_weights(); });
// Try loading backward weights from disk (.kt files) — parallel across NUMA nodes.
@ -489,7 +494,6 @@ class TP_MOE_SFT : public TP_MOE<T> {
throw std::runtime_error("Weights not loaded");
}
int qlen = *qlen_ptr;
auto pool = config.pool;
@ -504,7 +508,6 @@ class TP_MOE_SFT : public TP_MOE<T> {
save_for_backward);
});
// // Collect per-thread timing from all NUMA subpools
// for (int i = 0; i < tp_count; i++) {
// }
@ -514,9 +517,7 @@ class TP_MOE_SFT : public TP_MOE<T> {
// Merge results from all NUMA nodes
this->merge_results(qlen, output);
pool->dispense_backend()->do_numa_job([&](int numa_id) {
});
pool->dispense_backend()->do_numa_job([&](int numa_id) {});
}
/**
@ -551,7 +552,6 @@ class TP_MOE_SFT : public TP_MOE<T> {
void* grad_weights) {
auto pool = config.pool;
// Get full intermediate_size (before TP partitioning)
int full_intermediate_size = sft_config.intermediate_size;
int expert_num = config.expert_num;
@ -656,7 +656,6 @@ class TP_MOE_SFT : public TP_MOE<T> {
},
nullptr);
// Compute TP-slice pointers for copy-type direct writes
// Each TP writes to its own I-slice of the final output tensor
std::vector<ggml_bf16_t*> tp_gate_b_ptr(tp_count);
@ -882,8 +881,7 @@ class TP_MOE_SFT : public TP_MOE<T> {
nullptr);
}
pool->dispense_backend()->do_numa_job([&](int numa_id) {
});
pool->dispense_backend()->do_numa_job([&](int numa_id) {});
}
/**
@ -934,43 +932,38 @@ class TP_MOE_SFT : public TP_MOE<T> {
}
}
// Single do_numa_job: work-stealing memcpy + update_lora_weights
auto pool = config.pool;
pool->dispense_backend()->do_numa_job([this, gate_lora_a, gate_lora_b, up_lora_a, up_lora_b, down_lora_a,
down_lora_b, full_intermediate_size, expert_num, lora_rank,
pool](int numa_id) {
// LoRA weights are installed at load time. Keep the partitioning copy
// synchronous and serial here instead of nesting work-stealing jobs inside
// SGLang's scheduler process during model-loading barriers.
for (int numa_id = 0; numa_id < tp_count; numa_id++) {
int tp_inter = tp_configs[numa_id].intermediate_size;
size_t lora_b_slice = (size_t)tp_inter * lora_rank;
auto subpool = pool->get_subpool(numa_id);
// Work-stealing: copy all weights for this expert (gate + up + down)
subpool->do_work_stealing_job(
expert_num,
[&](int e) {
// gate_lora_b: [expert_num, intermediate_size, lora_rank]
memcpy(partitioned_gate_lora_b_[numa_id] + e * lora_b_slice,
(ggml_bf16_t*)gate_lora_b + e * full_intermediate_size * lora_rank + numa_id * lora_b_slice,
sizeof(ggml_bf16_t) * lora_b_slice);
for (int e = 0; e < expert_num; e++) {
// gate_lora_b: [expert_num, intermediate_size, lora_rank]
memcpy(partitioned_gate_lora_b_[numa_id] + e * lora_b_slice,
(ggml_bf16_t*)gate_lora_b + e * full_intermediate_size * lora_rank + numa_id * lora_b_slice,
sizeof(ggml_bf16_t) * lora_b_slice);
// up_lora_b: [expert_num, intermediate_size, lora_rank]
memcpy(partitioned_up_lora_b_[numa_id] + e * lora_b_slice,
(ggml_bf16_t*)up_lora_b + e * full_intermediate_size * lora_rank + numa_id * lora_b_slice,
sizeof(ggml_bf16_t) * lora_b_slice);
// up_lora_b: [expert_num, intermediate_size, lora_rank]
memcpy(partitioned_up_lora_b_[numa_id] + e * lora_b_slice,
(ggml_bf16_t*)up_lora_b + e * full_intermediate_size * lora_rank + numa_id * lora_b_slice,
sizeof(ggml_bf16_t) * lora_b_slice);
// down_lora_a: [expert_num, lora_rank, intermediate_size] - row-wise slice
for (int r = 0; r < lora_rank; r++) {
memcpy(partitioned_down_lora_a_[numa_id] + e * lora_rank * tp_inter + r * tp_inter,
(ggml_bf16_t*)down_lora_a + e * lora_rank * full_intermediate_size + r * full_intermediate_size +
numa_id * tp_inter,
sizeof(ggml_bf16_t) * tp_inter);
}
});
// down_lora_a: [expert_num, lora_rank, intermediate_size] - row-wise slice
for (int r = 0; r < lora_rank; r++) {
memcpy(partitioned_down_lora_a_[numa_id] + e * lora_rank * tp_inter + r * tp_inter,
(ggml_bf16_t*)down_lora_a + e * lora_rank * full_intermediate_size + r * full_intermediate_size +
numa_id * tp_inter,
sizeof(ggml_bf16_t) * tp_inter);
}
}
// Update weights after all memcpy complete
tps[numa_id]->update_lora_weights(gate_lora_a, partitioned_gate_lora_b_[numa_id], up_lora_a,
partitioned_up_lora_b_[numa_id], partitioned_down_lora_a_[numa_id],
down_lora_b);
});
}
}
/**

View file

@ -40,6 +40,8 @@ except (ImportError, AttributeError):
from .base import BaseSFTMoEWrapper, KExpertsSFTBuffer
_AMX_M_STEP = 32
# Mapping from method string to C++ SFT MOE class
_SFT_METHOD_TO_CLASS = {
@ -159,6 +161,17 @@ class AMXSFTMoEWrapper(BaseSFTMoEWrapper):
if self._weights_loaded:
return
if physical_to_logical_map_cpu is None:
physical_to_logical_map_cpu = torch.arange(self.num_experts, dtype=torch.int64)
self._physical_to_logical_map_cpu = physical_to_logical_map_cpu.to(
dtype=torch.int64, device="cpu"
).contiguous()
if self._physical_to_logical_map_cpu.numel() < self.num_experts:
raise ValueError(
"physical_to_logical_map_cpu must contain at least "
f"{self.num_experts} entries, got {self._physical_to_logical_map_cpu.numel()}."
)
if self.gate_proj is None and not getattr(self, "_use_projs_path", False):
self._load_base_weights_from_file()
@ -170,10 +183,11 @@ class AMXSFTMoEWrapper(BaseSFTMoEWrapper):
config.lora_rank = self.lora_rank
config.lora_alpha = self.lora_alpha
config.max_cache_depth = self.max_cache_depth
config.max_len = self.chunked_prefill_size
config.max_len = self._aligned_max_len()
config.layer_idx = self.layer_idx
config.share_backward_bb = getattr(self, "share_backward_bb", False)
config.share_cache_pool = getattr(self, "share_cache_pool", False)
config.physical_to_logical_map = self._physical_to_logical_map_cpu.data_ptr()
if getattr(self, "_use_kt_direct_load", False):
config.load = True
@ -220,8 +234,9 @@ class AMXSFTMoEWrapper(BaseSFTMoEWrapper):
self.cpu_infer.submit(self.moe.load_weights_task())
self.cpu_infer.sync()
self.cpu_infer.submit(self.moe.warm_up_task())
self.cpu_infer.sync()
if os.environ.get("KT_SFT_ENABLE_WARMUP", "0") == "1":
self.cpu_infer.submit(self.moe.warm_up_task())
self.cpu_infer.sync()
# Release Python-side weight tensors (C++ copied them)
self.gate_proj = None
@ -312,6 +327,7 @@ class AMXSFTMoEWrapper(BaseSFTMoEWrapper):
self._gate_scales_per_numa = experts_data["gate_scale"]
self._up_scales_per_numa = experts_data["up_scale"]
self._down_scales_per_numa = experts_data["down_scale"]
self._validate_prepartitioned_weights()
self._gate_projs_ptrs = _make_ptrs(gate_weights)
self._up_projs_ptrs = _make_ptrs(up_weights)
@ -345,6 +361,57 @@ class AMXSFTMoEWrapper(BaseSFTMoEWrapper):
loader.close_all_handles()
def _aligned_max_len(self) -> int:
return ((self.chunked_prefill_size + _AMX_M_STEP - 1) // _AMX_M_STEP) * _AMX_M_STEP
def _validate_prepartitioned_weights(self) -> None:
numa_count = len(self._gate_weights_per_numa)
if self.moe_intermediate_size % self.threadpool_count != 0:
raise ValueError(
f"moe_intermediate_size={self.moe_intermediate_size} must be divisible by "
f"threadpool_count={self.threadpool_count} for {self.method} SFT."
)
if numa_count != self.threadpool_count:
raise ValueError(
f"{self.method} SFT pre-partitioned expert weights have {numa_count} NUMA partitions, "
f"but CPUInfer was created with threadpool_count={self.threadpool_count}. "
f"Use --kt-threadpool-count {numa_count} for this weight directory, or convert weights "
"for the requested threadpool count."
)
collections = {
"gate": self._gate_weights_per_numa,
"up": self._up_weights_per_numa,
"down": self._down_weights_per_numa,
"gate_scale": self._gate_scales_per_numa,
"up_scale": self._up_scales_per_numa,
"down_scale": self._down_scales_per_numa,
}
for name, per_numa in collections.items():
if len(per_numa) != numa_count:
raise ValueError(f"{name} has {len(per_numa)} NUMA partitions, expected {numa_count}.")
for numa_id, entries in enumerate(per_numa):
if len(entries) != self.num_experts:
raise ValueError(
f"{name}[numa={numa_id}] has {len(entries)} experts, expected {self.num_experts}."
)
for numa_id in range(numa_count):
gate_scale_len = self._gate_scales_per_numa[numa_id][0].size
up_scale_len = self._up_scales_per_numa[numa_id][0].size
down_scale_len = self._down_scales_per_numa[numa_id][0].size
expected_intermediate = self.moe_intermediate_size // self.threadpool_count
if gate_scale_len != expected_intermediate or up_scale_len != expected_intermediate:
raise ValueError(
f"{self.method} gate/up scale length for NUMA {numa_id} is "
f"{gate_scale_len}/{up_scale_len}, expected {expected_intermediate}."
)
if down_scale_len != self.hidden_size:
raise ValueError(
f"{self.method} down scale length for NUMA {numa_id} is "
f"{down_scale_len}, expected {self.hidden_size}."
)
# ========== LoRA ==========
def init_lora_weights(
@ -373,6 +440,10 @@ class AMXSFTMoEWrapper(BaseSFTMoEWrapper):
expected = expected_shapes[name]
if tensor.shape != expected:
raise ValueError(f"{name} shape mismatch: expected {expected}, got {tuple(tensor.shape)}")
if tensor.device.type != "cpu":
raise ValueError(
f"{name} must be a CPU tensor for {self.method} SFT, got {tensor.device}."
)
self.gate_lora_a = gate_lora_a.contiguous()
self.gate_lora_b = gate_lora_b.contiguous()
@ -401,17 +472,17 @@ class AMXSFTMoEWrapper(BaseSFTMoEWrapper):
if not self._lora_initialized:
raise RuntimeError("LoRA weights not initialized. Call init_lora_weights() first.")
self.cpu_infer.submit(
self.moe.update_lora_weights_task(
self.gate_lora_a.data_ptr(),
self.gate_lora_b.data_ptr(),
self.up_lora_a.data_ptr(),
self.up_lora_b.data_ptr(),
self.down_lora_a.data_ptr(),
self.down_lora_b.data_ptr(),
)
# Weight pointer updates are load-time synchronous work. Calling the
# direct binding avoids nesting an update task inside CPUInfer's queue
# while SGLang is still in distributed model-loading barriers.
self.moe.update_lora_weights(
self.gate_lora_a.data_ptr(),
self.gate_lora_b.data_ptr(),
self.up_lora_a.data_ptr(),
self.up_lora_b.data_ptr(),
self.down_lora_a.data_ptr(),
self.down_lora_b.data_ptr(),
)
self.cpu_infer.sync()
def save_backward_weights_from_tensors(
self,

View file

@ -15,7 +15,7 @@ import torch
from typing import Optional, Tuple
from abc import ABC, abstractmethod
from ..experts_base import _MoEBase
from ..experts_base import KExpertsCPUBuffer, _MoEBase
class KExpertsSFTBuffer:
@ -98,6 +98,26 @@ class KExpertsSFTBuffer:
cls._shared_buffer = None
class _SFTForwardBufferView:
"""Minimal buffer view consumed by AMXSFTMoEWrapper._make_forward_task."""
__slots__ = ("bsz_tensor", "expert_ids_cpu", "weights_cpu", "input_cpu", "output_cpu")
def __init__(
self,
bsz_tensor: torch.Tensor,
expert_ids_cpu: torch.Tensor,
weights_cpu: torch.Tensor,
input_cpu: torch.Tensor,
output_cpu: torch.Tensor,
):
self.bsz_tensor = bsz_tensor
self.expert_ids_cpu = expert_ids_cpu
self.weights_cpu = weights_cpu
self.input_cpu = input_cpu
self.output_cpu = output_cpu
class BaseSFTMoEWrapper(_MoEBase, ABC):
"""
Base class for SFT MoE CPU operations with concrete buffer management.
@ -357,6 +377,104 @@ class BaseSFTMoEWrapper(_MoEBase, ABC):
return self._return_output(buffer, qlen, output_device)
# ========== Inference-only async forward ==========
def submit_forward_inference(
self,
hidden_states: torch.Tensor,
expert_ids: torch.Tensor,
weights: torch.Tensor,
cuda_stream,
) -> None:
"""
Submit an SFT MoE forward pass for serving.
This path mirrors the normal KT inference wrapper: inputs are copied to
pinned CPU staging buffers, the CPUInfer task is enqueued with the
caller CUDA stream, and sync_forward_inference() returns a persistent
GPU output buffer. It deliberately avoids the training-oriented
torch.cuda.synchronize() in _copy_inputs_to_buffer().
"""
if not hasattr(self.cpu_infer, "submit_with_cuda_stream"):
self.submit_forward(hidden_states, expert_ids, weights, save_for_backward=False)
self._pending_inference_fallback = True
self._pending_inference_fallback_device = hidden_states.device
return
self._validate_forward_inputs(hidden_states, expert_ids, weights)
flat_hidden_states = hidden_states.view(-1, hidden_states.shape[-1])
(
input_tensor_cpu,
expert_ids_cpu,
_deferred_expert_ids_cpu,
weights_cpu,
output_cpu,
bsz_tensor_cpu,
output_gpu,
) = KExpertsCPUBuffer.get_buffer(flat_hidden_states, self.num_experts_per_tok)
current_slot = self.layer_idx % KExpertsCPUBuffer.buffer_depth
bsz_slot_tensor = bsz_tensor_cpu[current_slot]
torch_stream = (
cuda_stream
if isinstance(cuda_stream, torch.cuda.Stream)
else torch.cuda.ExternalStream(cuda_stream, device=flat_hidden_states.device)
)
with torch.cuda.stream(torch_stream):
input_tensor_cpu[current_slot].copy_(flat_hidden_states.to(torch.bfloat16), non_blocking=True)
expert_ids_cpu[current_slot].copy_(expert_ids.to(torch.int64), non_blocking=True)
weights_cpu[current_slot].copy_(weights.to(torch.float32), non_blocking=True)
buffer_view = _SFTForwardBufferView(
bsz_tensor=bsz_slot_tensor,
expert_ids_cpu=expert_ids_cpu[current_slot],
weights_cpu=weights_cpu[current_slot],
input_cpu=input_tensor_cpu[current_slot],
output_cpu=output_cpu[current_slot],
)
self._pending_inference_fallback = False
self._pending_inference_output_cpu = output_cpu[current_slot]
self._pending_inference_output_gpu = output_gpu[current_slot]
self.cpu_infer.submit_with_cuda_stream(
cuda_stream,
self._make_forward_task(buffer_view, save_for_backward=False),
)
def sync_forward_inference(self, cuda_stream) -> torch.Tensor:
"""
Synchronize a serving forward submitted by submit_forward_inference().
Returns a persistent GPU buffer matching the input batch shape. Consumers
on the same CUDA stream will naturally wait for the non-blocking D2H/H2D
staging work ordered through CPUInfer's stream synchronization.
"""
if getattr(self, "_pending_inference_fallback", False):
self._pending_inference_fallback = False
output_device = getattr(self, "_pending_inference_fallback_device", None)
self._pending_inference_fallback_device = None
return self.sync_forward(output_device=output_device)
if not hasattr(self, "_pending_inference_output_cpu"):
raise RuntimeError("No pending inference forward. Call submit_forward_inference() first.")
torch_stream = (
cuda_stream
if isinstance(cuda_stream, torch.cuda.Stream)
else torch.cuda.ExternalStream(cuda_stream, device=self._pending_inference_output_gpu.device)
)
self.cpu_infer.sync_with_cuda_stream(cuda_stream)
with torch.cuda.stream(torch_stream):
self._pending_inference_output_gpu.copy_(self._pending_inference_output_cpu, non_blocking=True)
output = self._pending_inference_output_gpu
del self._pending_inference_output_cpu
del self._pending_inference_output_gpu
return output
# ========== Async backward ==========
def submit_backward_async(

View file

@ -726,6 +726,7 @@ class CompressedSafeTensorLoader(SafeTensorLoader):
"down_scale": down_scales,
}
class GGUFLoader:
"""
GGUF format loader using the official gguf library (gguf.gguf_reader.GGUFReader)

View file

@ -4,6 +4,86 @@ KT-Kernel provides weight conversion tools for CPU-GPU hybrid inference (e.g., i
- **CPU Weights (`convert_cpu_weights.py`)**: Quantize weights to INT4/INT8 with AMX optimization for CPU-resident "cold" experts
- **GPU Weights (`convert_gpu_weights.py`)**: Apply GPTQ/RTN quantization (W4A16/W8A16) for GPU-resident "hot" experts
- **KT Fused Expert LoRA (`convert_kt_to_sglang_adapter.py`)**: Convert KT SFT fused expert LoRA checkpoints into adapter-only SafeTensors directories
---
## KT Fused Expert LoRA Adapter Conversion
KT SFT fused expert LoRA saves MoE expert LoRA tensors in `fused_expert_lora.safetensors` using compact 3D tensors:
```
layers.{L}.experts.gate_lora_a
layers.{L}.experts.gate_lora_b
layers.{L}.experts.up_lora_a
layers.{L}.experts.up_lora_b
layers.{L}.experts.down_lora_a
layers.{L}.experts.down_lora_b
```
Use `convert_kt_to_sglang_adapter.py` to convert raw KT SFT output into one merged SGLang adapter directory:
```bash
python scripts/convert_kt_to_sglang_adapter.py /path/to/kt_adapter /path/to/sglang_adapter \
--base-model-name-or-path /path/to/base_model \
--lora-alpha 16 \
--overwrite
```
Output:
```
sglang_adapter/
├── adapter_config.json
└── adapter_model.safetensors
```
The converter merges the existing non-expert `adapter_model.safetensors` with expanded expert tensors from `fused_expert_lora.safetensors`. Pass this merged directory to SGLang with:
```bash
--enable-lora \
--lora-paths my_lora=/path/to/sglang_adapter
```
The KTransformers SGLang fork will auto-split the merged adapter internally at server startup. Users do not need to pass separate expert and non-expert adapter paths in the normal workflow.
Optional split outputs for debugging:
```bash
python scripts/convert_kt_to_sglang_adapter.py /path/to/kt_adapter /path/to/sglang_adapter \
--base-model-name-or-path /path/to/base_model \
--expert-output-dir /path/to/expert_adapter \
--nonexpert-output-dir /path/to/nonexpert_adapter \
--overwrite
```
Existing PEFT prefixes such as `base_model.model.` are stripped to match SGLang's loader. Scaling is not folded into the LoRA B tensors. Runtime scaling remains `lora_alpha / r`; if the input directory has no `adapter_config.json`, pass `--lora-alpha` explicitly.
This script only converts adapter files. Serving compatibility depends on the KTransformers SGLang runtime branch being used.
### Optional Integration Validation
The unit tests use synthetic tensors and run without model files. To validate a real KT adapter directory, set these environment variables:
```bash
export KT_LORA_ADAPTER_DIR=/path/to/kt_adapter
export KT_LORA_BASE_MODEL=/path/to/base_model
export KT_LORA_ALPHA=16 # required only if the input has no adapter_config.json
```
Then run:
```bash
python -m pytest kt-kernel/test/per_commit/test_convert_kt_to_sglang_adapter_integration.py -q
```
To run a large adapter conversion smoke test, also set:
```bash
export KT_LORA_LARGE_ADAPTER_DIR=/path/to/large_kt_adapter
```
These integration tests check real fused tensor splitting, optional `adapter_model.safetensors` merging, `adapter_config.json` compatibility with `sglang.srt.lora.lora_config.LoRAConfig`, and large-file readability. They intentionally do not start an SGLang server or validate runtime `FusedMoE` LoRA application.
---

View file

@ -0,0 +1,477 @@
#!/usr/bin/env python3
"""Convert KT fused expert LoRA checkpoints into an SGLang adapter directory."""
from __future__ import annotations
import argparse
import json
import os
import re
import shutil
from pathlib import Path
from typing import Dict, Iterable, Mapping
import torch
from safetensors.torch import load_file, save_file
FUSED_EXPERT_LORA_FILE = "fused_expert_lora.safetensors"
ADAPTER_MODEL_FILE = "adapter_model.safetensors"
ADAPTER_CONFIG_FILE = "adapter_config.json"
KT_NAME_MAP = {
"gate_lora_a": ("gate_proj", "lora_A", 1),
"gate_lora_b": ("gate_proj", "lora_B", 2),
"up_lora_a": ("up_proj", "lora_A", 1),
"up_lora_b": ("up_proj", "lora_B", 2),
"down_lora_a": ("down_proj", "lora_A", 1),
"down_lora_b": ("down_proj", "lora_B", 2),
}
TARGET_MODULE_ORDER = [
"q_proj",
"k_proj",
"v_proj",
"o_proj",
"gate_proj",
"up_proj",
"down_proj",
"in_proj_qkv",
"in_proj_z",
"in_proj_b",
"in_proj_a",
"out_proj",
"embed_tokens",
"lm_head",
]
KT_FUSED_KEY_RE = re.compile(r"^layers\.(\d+)\.experts\.([^.]+)$")
def _load_json(path: Path) -> dict:
with path.open("r", encoding="utf-8") as f:
return json.load(f)
def _write_json(path: Path, data: Mapping) -> None:
with path.open("w", encoding="utf-8") as f:
json.dump(data, f, indent=2, sort_keys=True)
f.write("\n")
def _clean_adapter_key(key: str) -> str:
"""Match the existing SGLang converter's PEFT key cleanup."""
key = key.replace("base_model.model.", "")
key = key.replace(".orig_module", "")
return key
def _ordered_target_modules(modules: Iterable[str]) -> list[str]:
seen = set(modules)
ordered = [name for name in TARGET_MODULE_ORDER if name in seen]
ordered.extend(sorted(seen.difference(ordered)))
return ordered
def _infer_target_module_from_key(key: str) -> str | None:
if "lora_embedding_A" in key or "lora_embedding_B" in key:
if "embed_tokens" in key:
return "embed_tokens"
if "lm_head" in key or "unembed_tokens" in key:
return "lm_head"
marker = ".lora_"
if marker not in key:
return None
prefix = key.split(marker, 1)[0]
if "." not in prefix:
return prefix
return prefix.rsplit(".", 1)[-1]
def _merge_tensor(tensors: Dict[str, torch.Tensor], key: str, value: torch.Tensor) -> None:
if key in tensors:
raise ValueError(f"Duplicate output tensor key: {key}")
tensors[key] = value.detach().cpu()
def _load_existing_adapter(input_dir: Path) -> tuple[dict[str, torch.Tensor], set[str]]:
adapter_path = input_dir / ADAPTER_MODEL_FILE
if not adapter_path.exists():
return {}, set()
tensors: dict[str, torch.Tensor] = {}
target_modules: set[str] = set()
for key, value in load_file(str(adapter_path)).items():
cleaned_key = _clean_adapter_key(key)
_merge_tensor(tensors, cleaned_key, value)
target_module = _infer_target_module_from_key(cleaned_key)
if target_module is not None:
target_modules.add(target_module)
return tensors, target_modules
def _convert_fused_expert_lora(
fused_path: Path,
) -> tuple[dict[str, torch.Tensor], int, set[str]]:
if not fused_path.exists():
raise FileNotFoundError(f"Missing {FUSED_EXPERT_LORA_FILE}: {fused_path}")
output: dict[str, torch.Tensor] = {}
ranks: set[int] = set()
expert_counts: set[int] = set()
target_modules: set[str] = set()
for key, tensor in sorted(load_file(str(fused_path)).items()):
match = KT_FUSED_KEY_RE.match(key)
if match is None:
raise ValueError(f"Unexpected key in {FUSED_EXPERT_LORA_FILE}: {key}")
layer_idx, kt_name = match.groups()
if kt_name not in KT_NAME_MAP:
raise ValueError(f"Unsupported KT fused expert LoRA tensor: {key}")
if tensor.dim() != 3:
raise ValueError(f"{key} must be 3D [E, ...], got shape {tuple(tensor.shape)}")
proj_name, lora_name, rank_dim = KT_NAME_MAP[kt_name]
expert_count = int(tensor.shape[0])
rank = int(tensor.shape[rank_dim])
expert_counts.add(expert_count)
ranks.add(rank)
target_modules.add(proj_name)
for expert_idx in range(expert_count):
output_key = (
f"model.layers.{layer_idx}.mlp.experts.{expert_idx}."
f"{proj_name}.{lora_name}.weight"
)
_merge_tensor(output, output_key, tensor[expert_idx].contiguous())
if not output:
raise ValueError(f"No tensors found in {fused_path}")
if len(expert_counts) != 1:
raise ValueError(f"Inconsistent expert counts in {FUSED_EXPERT_LORA_FILE}: {sorted(expert_counts)}")
if len(ranks) != 1:
raise ValueError(f"Inconsistent LoRA ranks in {FUSED_EXPERT_LORA_FILE}: {sorted(ranks)}")
return output, next(iter(ranks)), target_modules
def _build_adapter_config(
input_dir: Path,
rank: int,
target_modules: set[str],
base_model_name_or_path: str,
lora_alpha: float | None,
*,
include_input_target_modules: bool = True,
) -> dict:
config_path = input_dir / ADAPTER_CONFIG_FILE
config = _load_json(config_path) if config_path.exists() else {}
if "lora_alpha" in config:
final_alpha = config["lora_alpha"]
elif lora_alpha is not None:
final_alpha = lora_alpha
else:
raise ValueError(
f"No {ADAPTER_CONFIG_FILE} with lora_alpha found in {input_dir}; "
"pass --lora-alpha to preserve runtime scaling."
)
existing_targets = config.get("target_modules", [])
if include_input_target_modules and isinstance(existing_targets, list):
target_modules.update(str(name).split(".")[-1] for name in existing_targets)
config["peft_type"] = config.get("peft_type", "LORA")
config["r"] = rank
config["lora_alpha"] = final_alpha
config["target_modules"] = _ordered_target_modules(target_modules)
config["bias"] = config.get("bias", "none")
config["task_type"] = config.get("task_type", "CAUSAL_LM")
config["base_model_name_or_path"] = base_model_name_or_path
return config
def _paths_have_ancestor_relationship(left: Path, right: Path) -> bool:
if left == right:
return True
try:
left.relative_to(right)
return True
except ValueError:
pass
try:
right.relative_to(left)
return True
except ValueError:
return False
def _validate_no_ancestor_paths(
paths: Iterable[Path],
*,
label: str,
) -> None:
resolved = list(paths)
for i, left in enumerate(resolved):
for right in resolved[i + 1 :]:
if _paths_have_ancestor_relationship(left, right):
raise ValueError(
f"{label} cannot have ancestor/descendant relationships: "
f"{left} and {right}."
)
def _prepare_output_dir(output_path: Path, input_path: Path, overwrite: bool) -> None:
_validate_output_dir(output_path, input_path, overwrite)
if output_path.exists() and any(output_path.iterdir()):
shutil.rmtree(output_path)
output_path.mkdir(parents=True, exist_ok=True)
def _validate_output_dir(output_path: Path, input_path: Path, overwrite: bool) -> None:
if output_path == input_path:
raise ValueError("Output directory must be different from input directory.")
if _paths_have_ancestor_relationship(output_path, input_path):
raise ValueError(
"Output and input directories cannot be ancestor/descendant of each other: "
f"output={output_path}, input={input_path}."
)
if output_path.exists() and not output_path.is_dir():
raise FileExistsError(f"Output path exists and is not a directory: {output_path}")
if output_path.exists() and any(output_path.iterdir()):
if not overwrite:
raise FileExistsError(f"Output directory is not empty: {output_path}")
def _infer_lora_rank_from_tensor(key: str, tensor: torch.Tensor) -> int | None:
if ".lora_A." in key:
return int(tensor.shape[0])
if ".lora_B." in key:
return int(tensor.shape[1])
return None
def _validate_nonexpert_rank(
existing_tensors: Mapping[str, torch.Tensor],
expert_rank: int,
input_dir: Path,
) -> None:
if not existing_tensors:
return
config_path = input_dir / ADAPTER_CONFIG_FILE
if config_path.exists():
config_rank = _load_json(config_path).get("r")
if config_rank is not None and int(config_rank) != expert_rank:
raise ValueError(
f"Non-expert adapter rank mismatch: adapter_config.json r={config_rank}, "
f"but fused expert LoRA rank={expert_rank}."
)
for key, tensor in existing_tensors.items():
tensor_rank = _infer_lora_rank_from_tensor(key, tensor)
if tensor_rank is None:
continue
if tensor_rank != expert_rank:
raise ValueError(
f"Non-expert adapter tensor rank mismatch for {key}: "
f"tensor rank={tensor_rank}, fused expert LoRA rank={expert_rank}."
)
def _write_adapter(
output_path: Path,
input_path: Path,
tensors: dict[str, torch.Tensor],
config: Mapping,
*,
overwrite: bool,
) -> None:
_prepare_output_dir(output_path, input_path, overwrite)
save_file(tensors, str(output_path / ADAPTER_MODEL_FILE), metadata={"format": "pt"})
_write_json(output_path / ADAPTER_CONFIG_FILE, config)
def convert_kt_to_sglang_adapter(
input_dir: str | os.PathLike,
output_dir: str | os.PathLike,
*,
base_model_name_or_path: str,
lora_alpha: float | None = None,
overwrite: bool = False,
expert_output_dir: str | os.PathLike | None = None,
nonexpert_output_dir: str | os.PathLike | None = None,
) -> dict:
input_path = Path(input_dir).expanduser().resolve()
output_path = Path(output_dir).expanduser().resolve()
expert_output_path = (
Path(expert_output_dir).expanduser().resolve()
if expert_output_dir is not None
else None
)
nonexpert_output_path = (
Path(nonexpert_output_dir).expanduser().resolve()
if nonexpert_output_dir is not None
else None
)
if not input_path.is_dir():
raise FileNotFoundError(f"Input directory not found: {input_path}")
output_paths = [output_path]
output_paths.extend(path for path in (expert_output_path, nonexpert_output_path) if path is not None)
if len(set(output_paths)) != len(output_paths):
raise ValueError("Merged, expert, and non-expert output directories must be distinct.")
_validate_no_ancestor_paths(
output_paths,
label="Merged/expert/non-expert output directories",
)
for path in output_paths:
_validate_output_dir(path, input_path, overwrite)
existing_tensors, existing_targets = _load_existing_adapter(input_path)
fused_tensors, rank, fused_targets = _convert_fused_expert_lora(input_path / FUSED_EXPERT_LORA_FILE)
_validate_nonexpert_rank(existing_tensors, rank, input_path)
if nonexpert_output_path is not None and not existing_tensors:
raise ValueError(
f"Cannot write non-expert adapter: no {ADAPTER_MODEL_FILE} found in {input_path}."
)
tensors: dict[str, torch.Tensor] = {}
for key, value in existing_tensors.items():
_merge_tensor(tensors, key, value)
for key, value in fused_tensors.items():
_merge_tensor(tensors, key, value)
target_modules = set(existing_targets)
target_modules.update(fused_targets)
config = _build_adapter_config(
input_path,
rank,
target_modules,
base_model_name_or_path,
lora_alpha,
)
_write_adapter(output_path, input_path, tensors, config, overwrite=overwrite)
split_outputs: dict[str, dict] = {}
if expert_output_path is not None:
expert_config = _build_adapter_config(
input_path,
rank,
set(fused_targets),
base_model_name_or_path,
lora_alpha,
include_input_target_modules=False,
)
_write_adapter(
expert_output_path,
input_path,
fused_tensors,
expert_config,
overwrite=overwrite,
)
split_outputs["expert"] = {
"output_dir": str(expert_output_path),
"tensor_count": len(fused_tensors),
"target_modules": expert_config["target_modules"],
}
if nonexpert_output_path is not None:
nonexpert_config = _build_adapter_config(
input_path,
rank,
set(existing_targets),
base_model_name_or_path,
lora_alpha,
include_input_target_modules=False,
)
_write_adapter(
nonexpert_output_path,
input_path,
existing_tensors,
nonexpert_config,
overwrite=overwrite,
)
split_outputs["nonexpert"] = {
"output_dir": str(nonexpert_output_path),
"tensor_count": len(existing_tensors),
"target_modules": nonexpert_config["target_modules"],
}
return {
"input_dir": str(input_path),
"output_dir": str(output_path),
"tensor_count": len(tensors),
"rank": rank,
"target_modules": config["target_modules"],
"lora_alpha": config["lora_alpha"],
"split_outputs": split_outputs,
}
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(
description="Convert KT fused expert LoRA weights to an SGLang adapter directory."
)
parser.add_argument("input_dir", help="Directory containing fused_expert_lora.safetensors.")
parser.add_argument("output_dir", help="Destination adapter directory.")
parser.add_argument(
"--base-model-name-or-path",
required=True,
help="Base model path/name to write into adapter_config.json.",
)
parser.add_argument(
"--lora-alpha",
type=float,
default=None,
help="LoRA alpha to use when input adapter_config.json is absent.",
)
parser.add_argument(
"--overwrite",
action="store_true",
help="Remove and recreate output_dir if it already contains files.",
)
parser.add_argument(
"--expert-output-dir",
default=None,
help="Optional destination for a split expert-only adapter directory.",
)
parser.add_argument(
"--nonexpert-output-dir",
default=None,
help="Optional destination for a split non-expert-only adapter directory.",
)
return parser.parse_args()
def main() -> None:
args = parse_args()
summary = convert_kt_to_sglang_adapter(
args.input_dir,
args.output_dir,
base_model_name_or_path=args.base_model_name_or_path,
lora_alpha=args.lora_alpha,
overwrite=args.overwrite,
expert_output_dir=args.expert_output_dir,
nonexpert_output_dir=args.nonexpert_output_dir,
)
print(
"Converted KT fused expert LoRA adapter: "
f"{summary['tensor_count']} tensors, rank={summary['rank']}, "
f"target_modules={summary['target_modules']}"
)
for name, split_summary in summary["split_outputs"].items():
print(
f"Wrote {name} adapter: {split_summary['tensor_count']} tensors, "
f"target_modules={split_summary['target_modules']}, "
f"output_dir={split_summary['output_dir']}"
)
if __name__ == "__main__":
main()

View file

@ -0,0 +1,373 @@
import importlib.util
import os
import sys
from pathlib import Path
import pytest
import torch
from safetensors.torch import load_file, save_file
sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
from ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=5, suite="default")
SCRIPT_PATH = (
Path(__file__).resolve().parents[2]
/ "scripts"
/ "convert_kt_to_sglang_adapter.py"
)
SPEC = importlib.util.spec_from_file_location("convert_kt_to_sglang_adapter", SCRIPT_PATH)
converter = importlib.util.module_from_spec(SPEC)
assert SPEC.loader is not None
SPEC.loader.exec_module(converter)
def _write_full_fused_checkpoint(path: Path, *, rank: int = 3) -> dict[str, torch.Tensor]:
e, h, i = 2, 5, 7
tensors = {
"layers.2.experts.gate_lora_a": torch.arange(e * rank * h, dtype=torch.float32).reshape(e, rank, h),
"layers.2.experts.gate_lora_b": torch.arange(e * i * rank, dtype=torch.float32).reshape(e, i, rank),
"layers.2.experts.up_lora_a": torch.arange(e * rank * h, dtype=torch.float32).reshape(e, rank, h) + 100,
"layers.2.experts.up_lora_b": torch.arange(e * i * rank, dtype=torch.float32).reshape(e, i, rank) + 200,
"layers.2.experts.down_lora_a": torch.arange(e * rank * i, dtype=torch.float32).reshape(e, rank, i) + 300,
"layers.2.experts.down_lora_b": torch.arange(e * h * rank, dtype=torch.float32).reshape(e, h, rank) + 400,
}
save_file(tensors, str(path / converter.FUSED_EXPERT_LORA_FILE))
return tensors
def test_convert_fused_expert_lora_shapes_keys_and_config(tmp_path):
input_dir = tmp_path / "input"
output_dir = tmp_path / "output"
input_dir.mkdir()
fused = _write_full_fused_checkpoint(input_dir)
summary = converter.convert_kt_to_sglang_adapter(
input_dir,
output_dir,
base_model_name_or_path="/models/base",
lora_alpha=16,
)
out = load_file(str(output_dir / converter.ADAPTER_MODEL_FILE))
assert summary["tensor_count"] == 12
assert summary["rank"] == 3
assert summary["target_modules"] == ["gate_proj", "up_proj", "down_proj"]
key = "model.layers.2.mlp.experts.1.gate_proj.lora_A.weight"
assert out[key].shape == (3, 5)
torch.testing.assert_close(out[key], fused["layers.2.experts.gate_lora_a"][1])
key = "model.layers.2.mlp.experts.0.down_proj.lora_B.weight"
assert out[key].shape == (5, 3)
torch.testing.assert_close(out[key], fused["layers.2.experts.down_lora_b"][0])
config = converter._load_json(output_dir / converter.ADAPTER_CONFIG_FILE)
assert config["peft_type"] == "LORA"
assert config["r"] == 3
assert config["lora_alpha"] == 16
assert config["base_model_name_or_path"] == "/models/base"
assert config["target_modules"] == ["gate_proj", "up_proj", "down_proj"]
def test_merges_existing_adapter_and_prefers_existing_lora_alpha(tmp_path):
input_dir = tmp_path / "input"
output_dir = tmp_path / "output"
input_dir.mkdir()
_write_full_fused_checkpoint(input_dir)
existing_tensor = torch.ones(3, 5)
save_file(
{
"base_model.model.model.layers.0.self_attn.q_proj.lora_A.weight": existing_tensor,
},
str(input_dir / converter.ADAPTER_MODEL_FILE),
)
converter._write_json(
input_dir / converter.ADAPTER_CONFIG_FILE,
{
"peft_type": "LORA",
"r": 3,
"lora_alpha": 9,
"target_modules": ["q_proj"],
"bias": "none",
"task_type": "CAUSAL_LM",
"base_model_name_or_path": "old-base",
},
)
summary = converter.convert_kt_to_sglang_adapter(
input_dir,
output_dir,
base_model_name_or_path="/models/base",
lora_alpha=16,
)
out = load_file(str(output_dir / converter.ADAPTER_MODEL_FILE))
cleaned_key = "model.layers.0.self_attn.q_proj.lora_A.weight"
assert cleaned_key in out
torch.testing.assert_close(out[cleaned_key], existing_tensor)
assert summary["lora_alpha"] == 9
config = converter._load_json(output_dir / converter.ADAPTER_CONFIG_FILE)
assert config["lora_alpha"] == 9
assert config["base_model_name_or_path"] == "/models/base"
assert config["target_modules"] == ["q_proj", "gate_proj", "up_proj", "down_proj"]
def test_writes_split_expert_and_nonexpert_adapters(tmp_path):
input_dir = tmp_path / "input"
output_dir = tmp_path / "output"
expert_dir = tmp_path / "expert"
nonexpert_dir = tmp_path / "nonexpert"
input_dir.mkdir()
_write_full_fused_checkpoint(input_dir)
q_proj_tensor = torch.ones(3, 5)
o_proj_tensor = torch.full((5, 3), 2.0)
save_file(
{
"base_model.model.model.layers.0.self_attn.q_proj.lora_A.weight": q_proj_tensor,
"base_model.model.model.layers.0.self_attn.o_proj.lora_B.weight": o_proj_tensor,
},
str(input_dir / converter.ADAPTER_MODEL_FILE),
)
converter._write_json(
input_dir / converter.ADAPTER_CONFIG_FILE,
{
"peft_type": "LORA",
"r": 3,
"lora_alpha": 9,
"target_modules": ["q_proj", "o_proj"],
"bias": "none",
"task_type": "CAUSAL_LM",
"base_model_name_or_path": "old-base",
},
)
summary = converter.convert_kt_to_sglang_adapter(
input_dir,
output_dir,
base_model_name_or_path="/models/base",
expert_output_dir=expert_dir,
nonexpert_output_dir=nonexpert_dir,
)
merged = load_file(str(output_dir / converter.ADAPTER_MODEL_FILE))
expert = load_file(str(expert_dir / converter.ADAPTER_MODEL_FILE))
nonexpert = load_file(str(nonexpert_dir / converter.ADAPTER_MODEL_FILE))
assert summary["tensor_count"] == 14
assert summary["split_outputs"]["expert"]["tensor_count"] == 12
assert summary["split_outputs"]["nonexpert"]["tensor_count"] == 2
assert set(merged) == set(expert) | set(nonexpert)
assert set(expert).isdisjoint(nonexpert)
assert all(".mlp.experts." in key for key in expert)
assert not any(".mlp.experts." in key for key in nonexpert)
cleaned_q_proj_key = "model.layers.0.self_attn.q_proj.lora_A.weight"
cleaned_o_proj_key = "model.layers.0.self_attn.o_proj.lora_B.weight"
torch.testing.assert_close(nonexpert[cleaned_q_proj_key], q_proj_tensor)
torch.testing.assert_close(nonexpert[cleaned_o_proj_key], o_proj_tensor)
expert_config = converter._load_json(expert_dir / converter.ADAPTER_CONFIG_FILE)
nonexpert_config = converter._load_json(nonexpert_dir / converter.ADAPTER_CONFIG_FILE)
assert expert_config["target_modules"] == ["gate_proj", "up_proj", "down_proj"]
assert nonexpert_config["target_modules"] == ["q_proj", "o_proj"]
assert expert_config["base_model_name_or_path"] == "/models/base"
assert nonexpert_config["base_model_name_or_path"] == "/models/base"
def test_requires_lora_alpha_without_input_config(tmp_path):
input_dir = tmp_path / "input"
input_dir.mkdir()
_write_full_fused_checkpoint(input_dir)
with pytest.raises(ValueError, match="pass --lora-alpha"):
converter.convert_kt_to_sglang_adapter(
input_dir,
tmp_path / "output",
base_model_name_or_path="/models/base",
)
def test_rejects_inconsistent_rank(tmp_path):
input_dir = tmp_path / "input"
input_dir.mkdir()
save_file(
{
"layers.0.experts.gate_lora_a": torch.zeros(2, 3, 5),
"layers.0.experts.gate_lora_b": torch.zeros(2, 7, 4),
},
str(input_dir / converter.FUSED_EXPERT_LORA_FILE),
)
with pytest.raises(ValueError, match="Inconsistent LoRA ranks"):
converter.convert_kt_to_sglang_adapter(
input_dir,
tmp_path / "output",
base_model_name_or_path="/models/base",
lora_alpha=8,
)
def test_rejects_unexpected_fused_key(tmp_path):
input_dir = tmp_path / "input"
input_dir.mkdir()
save_file(
{"layers.0.experts.unknown_lora_a": torch.zeros(2, 3, 5)},
str(input_dir / converter.FUSED_EXPERT_LORA_FILE),
)
with pytest.raises(ValueError, match="Unsupported KT fused expert LoRA tensor"):
converter.convert_kt_to_sglang_adapter(
input_dir,
tmp_path / "output",
base_model_name_or_path="/models/base",
lora_alpha=8,
)
def test_rejects_nonempty_output_without_overwrite(tmp_path):
input_dir = tmp_path / "input"
output_dir = tmp_path / "output"
input_dir.mkdir()
output_dir.mkdir()
(output_dir / "existing.txt").write_text("do not remove", encoding="utf-8")
_write_full_fused_checkpoint(input_dir)
with pytest.raises(FileExistsError, match="Output directory is not empty"):
converter.convert_kt_to_sglang_adapter(
input_dir,
output_dir,
base_model_name_or_path="/models/base",
lora_alpha=8,
)
converter.convert_kt_to_sglang_adapter(
input_dir,
output_dir,
base_model_name_or_path="/models/base",
lora_alpha=8,
overwrite=True,
)
assert not (output_dir / "existing.txt").exists()
assert (output_dir / converter.ADAPTER_MODEL_FILE).exists()
def test_rejects_output_same_as_input(tmp_path):
input_dir = tmp_path / "input"
input_dir.mkdir()
_write_full_fused_checkpoint(input_dir)
with pytest.raises(ValueError, match="different from input"):
converter.convert_kt_to_sglang_adapter(
input_dir,
input_dir,
base_model_name_or_path="/models/base",
lora_alpha=8,
overwrite=True,
)
def test_rejects_output_ancestor_of_input(tmp_path):
run_dir = tmp_path / "run"
input_dir = run_dir / "adapter"
output_dir = run_dir
input_dir.mkdir(parents=True)
_write_full_fused_checkpoint(input_dir)
with pytest.raises(ValueError, match="ancestor/descendant"):
converter.convert_kt_to_sglang_adapter(
input_dir,
output_dir,
base_model_name_or_path="/models/base",
lora_alpha=8,
overwrite=True,
)
def test_rejects_output_descendant_of_input(tmp_path):
input_dir = tmp_path / "input"
output_dir = input_dir / "output"
input_dir.mkdir()
_write_full_fused_checkpoint(input_dir)
with pytest.raises(ValueError, match="ancestor/descendant"):
converter.convert_kt_to_sglang_adapter(
input_dir,
output_dir,
base_model_name_or_path="/models/base",
lora_alpha=8,
overwrite=True,
)
def test_rejects_split_output_ancestor_relationship(tmp_path):
input_dir = tmp_path / "input"
output_dir = tmp_path / "output"
expert_dir = tmp_path / "split" / "expert"
nonexpert_dir = tmp_path / "split"
input_dir.mkdir()
_write_full_fused_checkpoint(input_dir)
save_file(
{
"base_model.model.model.layers.0.self_attn.q_proj.lora_A.weight": torch.ones(3, 5),
},
str(input_dir / converter.ADAPTER_MODEL_FILE),
)
converter._write_json(
input_dir / converter.ADAPTER_CONFIG_FILE,
{
"peft_type": "LORA",
"r": 3,
"lora_alpha": 9,
"target_modules": ["q_proj"],
"bias": "none",
"task_type": "CAUSAL_LM",
"base_model_name_or_path": "old-base",
},
)
with pytest.raises(ValueError, match="ancestor/descendant"):
converter.convert_kt_to_sglang_adapter(
input_dir,
output_dir,
base_model_name_or_path="/models/base",
expert_output_dir=expert_dir,
nonexpert_output_dir=nonexpert_dir,
)
def test_rejects_mismatched_nonexpert_rank(tmp_path):
input_dir = tmp_path / "input"
output_dir = tmp_path / "output"
input_dir.mkdir()
_write_full_fused_checkpoint(input_dir, rank=3)
save_file(
{
"base_model.model.model.layers.0.self_attn.q_proj.lora_A.weight": torch.ones(4, 5),
},
str(input_dir / converter.ADAPTER_MODEL_FILE),
)
converter._write_json(
input_dir / converter.ADAPTER_CONFIG_FILE,
{
"peft_type": "LORA",
"r": 4,
"lora_alpha": 9,
"target_modules": ["q_proj"],
"bias": "none",
"task_type": "CAUSAL_LM",
"base_model_name_or_path": "old-base",
},
)
with pytest.raises(ValueError, match="rank mismatch"):
converter.convert_kt_to_sglang_adapter(
input_dir,
output_dir,
base_model_name_or_path="/models/base",
)

View file

@ -0,0 +1,242 @@
import importlib.util
import json
import os
import re
import shutil
import sys
import time
from pathlib import Path
import pytest
import torch
from safetensors.torch import load_file
sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
from ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=60, suite="default")
KT_ADAPTER_ENV = "KT_LORA_ADAPTER_DIR"
KT_BASE_MODEL_ENV = "KT_LORA_BASE_MODEL"
KT_ALPHA_ENV = "KT_LORA_ALPHA"
KT_LARGE_ADAPTER_ENV = "KT_LORA_LARGE_ADAPTER_DIR"
REPO_ROOT = Path(__file__).resolve().parents[3]
SCRIPT_PATH = (
Path(__file__).resolve().parents[2]
/ "scripts"
/ "convert_kt_to_sglang_adapter.py"
)
SPEC = importlib.util.spec_from_file_location("convert_kt_to_sglang_adapter", SCRIPT_PATH)
converter = importlib.util.module_from_spec(SPEC)
assert SPEC.loader is not None
SPEC.loader.exec_module(converter)
SGLANG_EXPERT_KEY_RE = re.compile(
r"^model\.layers\.(\d+)\.mlp\.experts\.(\d+)\."
r"(gate_proj|up_proj|down_proj)\.lora_[AB]\.weight$"
)
def _required_adapter_dir(env_name: str) -> Path:
value = os.environ.get(env_name)
if not value:
pytest.skip(f"Set {env_name} to run real adapter integration tests.")
path = Path(value).expanduser().resolve()
if not path.is_dir():
pytest.fail(f"{env_name} is not a directory: {path}")
if not (path / converter.FUSED_EXPERT_LORA_FILE).is_file():
pytest.fail(f"{env_name} must contain {converter.FUSED_EXPERT_LORA_FILE}: {path}")
return path
def _base_model_name_or_path() -> str:
value = os.environ.get(KT_BASE_MODEL_ENV)
if not value:
pytest.fail(f"Set {KT_BASE_MODEL_ENV} before running real adapter integration tests.")
return value
def _optional_lora_alpha() -> float | None:
value = os.environ.get(KT_ALPHA_ENV)
if value in (None, ""):
return None
try:
return float(value)
except ValueError:
pytest.fail(f"{KT_ALPHA_ENV} must be numeric, got: {value!r}")
def _lora_alpha_for_input(input_dir: Path) -> float | None:
alpha = _optional_lora_alpha()
if (input_dir / converter.ADAPTER_CONFIG_FILE).exists():
return alpha
if alpha is None:
pytest.fail(
f"{input_dir} has no {converter.ADAPTER_CONFIG_FILE}; set {KT_ALPHA_ENV}."
)
return alpha
def _convert_real_adapter(input_dir: Path, tmp_path: Path, output_name: str = "output") -> tuple[Path, dict]:
output_dir = tmp_path / output_name
summary = converter.convert_kt_to_sglang_adapter(
input_dir,
output_dir,
base_model_name_or_path=_base_model_name_or_path(),
lora_alpha=_lora_alpha_for_input(input_dir),
)
return output_dir, summary
def _load_json(path: Path) -> dict:
with path.open("r", encoding="utf-8") as f:
return json.load(f)
def _link_or_copy(source: Path, dest: Path) -> None:
try:
os.symlink(source, dest)
except OSError:
shutil.copy2(source, dest)
def _assert_config_shape(output_dir: Path, summary: dict) -> dict:
config = _load_json(output_dir / converter.ADAPTER_CONFIG_FILE)
assert config["peft_type"] == "LORA"
assert config["r"] == summary["rank"]
assert config["lora_alpha"] == summary["lora_alpha"]
assert config["base_model_name_or_path"] == _base_model_name_or_path()
assert {"gate_proj", "up_proj", "down_proj"}.issubset(config["target_modules"])
return config
def _assert_fused_tensors_preserved(input_dir: Path, output_dir: Path) -> int:
fused_tensors = load_file(str(input_dir / converter.FUSED_EXPERT_LORA_FILE))
output_tensors = load_file(str(output_dir / converter.ADAPTER_MODEL_FILE))
checked = 0
for input_key, input_tensor in sorted(fused_tensors.items()):
match = converter.KT_FUSED_KEY_RE.match(input_key)
assert match is not None, input_key
layer_idx, kt_name = match.groups()
assert kt_name in converter.KT_NAME_MAP, input_key
assert input_tensor.dim() == 3, input_key
proj_name, lora_name, _rank_dim = converter.KT_NAME_MAP[kt_name]
for expert_idx in range(input_tensor.shape[0]):
output_key = (
f"model.layers.{layer_idx}.mlp.experts.{expert_idx}."
f"{proj_name}.{lora_name}.weight"
)
assert SGLANG_EXPERT_KEY_RE.match(output_key), output_key
assert output_key in output_tensors
assert output_tensors[output_key].shape == input_tensor[expert_idx].shape
assert output_tensors[output_key].dtype == input_tensor.dtype
assert torch.equal(output_tensors[output_key], input_tensor[expert_idx].cpu())
checked += 1
assert checked > 0
return checked
@pytest.mark.requires_model
def test_real_adapter_conversion_preserves_fused_tensors_and_config(tmp_path):
input_dir = _required_adapter_dir(KT_ADAPTER_ENV)
output_dir, summary = _convert_real_adapter(input_dir, tmp_path)
checked = _assert_fused_tensors_preserved(input_dir, output_dir)
config = _assert_config_shape(output_dir, summary)
output_tensors = load_file(str(output_dir / converter.ADAPTER_MODEL_FILE))
assert summary["tensor_count"] == len(output_tensors)
assert checked <= summary["tensor_count"]
assert config["target_modules"] == summary["target_modules"]
@pytest.mark.requires_model
def test_real_adapter_directory_merges_existing_adapter_model(tmp_path):
input_dir = _required_adapter_dir(KT_ADAPTER_ENV)
existing_adapter_path = input_dir / converter.ADAPTER_MODEL_FILE
if not existing_adapter_path.exists():
pytest.skip(f"{input_dir} has no {converter.ADAPTER_MODEL_FILE} to merge.")
output_dir, _summary = _convert_real_adapter(input_dir, tmp_path)
input_tensors = load_file(str(existing_adapter_path))
output_tensors = load_file(str(output_dir / converter.ADAPTER_MODEL_FILE))
for input_key, input_tensor in input_tensors.items():
output_key = converter._clean_adapter_key(input_key)
assert output_key in output_tensors
assert output_tensors[output_key].shape == input_tensor.shape
assert output_tensors[output_key].dtype == input_tensor.dtype
assert torch.equal(output_tensors[output_key], input_tensor.cpu())
@pytest.mark.requires_model
def test_real_fused_conversion_without_input_config_uses_env_alpha(tmp_path):
input_dir = _required_adapter_dir(KT_ADAPTER_ENV)
alpha = _optional_lora_alpha()
if alpha is None:
pytest.skip(f"Set {KT_ALPHA_ENV} to validate conversion without input config.")
no_config_input = tmp_path / "input_without_config"
no_config_input.mkdir()
_link_or_copy(
input_dir / converter.FUSED_EXPERT_LORA_FILE,
no_config_input / converter.FUSED_EXPERT_LORA_FILE,
)
existing_adapter_path = input_dir / converter.ADAPTER_MODEL_FILE
if existing_adapter_path.exists():
_link_or_copy(existing_adapter_path, no_config_input / converter.ADAPTER_MODEL_FILE)
output_dir, summary = _convert_real_adapter(no_config_input, tmp_path, "output_without_config")
assert summary["lora_alpha"] == alpha
config = _assert_config_shape(output_dir, summary)
assert config["lora_alpha"] == alpha
_assert_fused_tensors_preserved(no_config_input, output_dir)
@pytest.mark.requires_model
def test_sglang_lora_config_loader_accepts_converted_adapter(tmp_path):
input_dir = _required_adapter_dir(KT_ADAPTER_ENV)
output_dir, summary = _convert_real_adapter(input_dir, tmp_path)
sglang_python = REPO_ROOT / "third_party" / "sglang" / "python"
sys.path.insert(0, str(sglang_python))
try:
from sglang.srt.lora.lora_config import LoRAConfig
except Exception as exc:
pytest.fail(f"Unable to import SGLang LoRAConfig: {exc}")
lora_config = LoRAConfig(str(output_dir))
output_tensors = load_file(str(output_dir / converter.ADAPTER_MODEL_FILE))
assert lora_config.r == summary["rank"]
assert lora_config.lora_alpha == summary["lora_alpha"]
assert lora_config.target_modules == summary["target_modules"]
assert len(output_tensors) == summary["tensor_count"]
@pytest.mark.requires_model
def test_large_adapter_conversion_smoke(tmp_path, record_property):
input_dir = _required_adapter_dir(KT_LARGE_ADAPTER_ENV)
start_time = time.perf_counter()
output_dir, summary = _convert_real_adapter(input_dir, tmp_path, "large_output")
duration_seconds = time.perf_counter() - start_time
output_path = output_dir / converter.ADAPTER_MODEL_FILE
output_tensors = load_file(str(output_path))
config = _assert_config_shape(output_dir, summary)
record_property("conversion_seconds", round(duration_seconds, 3))
record_property("output_bytes", output_path.stat().st_size)
record_property("tensor_count", summary["tensor_count"])
assert len(output_tensors) == summary["tensor_count"]
assert config["target_modules"] == summary["target_modules"]

2
third_party/sglang vendored

@ -1 +1 @@
Subproject commit ebaff7729b9e41c29d94f8d19a53473d321dc566
Subproject commit 7dc02b12fdde70f911a69f68e230a85f9fce1775