[Feature] Add MESH expert residency with io_uring direct I/O for MoE inference

.gitignore

@@ -30,4 +30,5 @@ build*
 CMakeFiles/
 kvc2/
 sched/
-*.png
+*.png
+.trae/

TROUBLESHOOTING_LOG.md

@@ -0,0 +1,62 @@
+# 踩坑日志
+
+## 2026-06-21: 35B TP>1 输出坍缩为 `!`
+
+### 现象
+- 35B AMXINT4 模型在 TP>1（TP2/TP4）时，输出全部坍缩为 `!`
+- TP1 正常
+- layer 3 的 GPU w2_weight 包含 NaN/Inf
+
+### 根因
+启动脚本中 `--model` 和 `--kt-weight-path` 都指向 AMXINT4 格式 checkpoint（`Qwen3.5-35B-A3B-AMXINT4-NUMA2-MESH`）。
+
+AMXINT4 格式的 checkpoint key 为 `blk.N.ffn_down_exps.E.numa.T.weight` + `.scale`，按 NUMA/TP 预切分存储。sglang 标准 weight_loader（`_weight_loader_physical` → `_weight_loader_impl` → `_load_w2`）只处理标准格式 key（`experts.N.down_proj.`），不识别 AMXINT4 格式。
+
+因此 GPU 的 `w2_weight`（由 `UnquantizedFusedMoEMethod.create_weights` 创建为 `torch.empty`）从未被 checkpoint 填充，保持未初始化状态。TP1 时恰好不含 NaN，TP>1 时 layer 3 的 w2_weight 包含 NaN/Inf，导致输出 NaN 坍缩为 `!`。
+
+### 修复
+- `--model` 改为标准 bf16 格式 checkpoint（`Qwen3.5-35B-A3B-Unfused`）
+- `--kt-weight-path` 保持 AMXINT4 格式（供 CPU AMX 内核使用）
+
+```
+--model /mnt/data2/tmp/qujing_mesh/Qwen3.5-35B-A3B-Unfused/
+--kt-weight-path /mnt/data2/models/Qwen3.5-35B-A3B-AMXINT4-NUMA2-MESH/
+```
+
+### 验证
+- DIAG 输出确认 `_weight_loader_physical` 和 `_load_w2` 被正确调用
+- w2_weight 被正确加载（nan=False, inf=False）
+- TP 切分正确（shard_dim=1, shard_size=256, tp_rank, narrow 后 shape=[2048, 256]）
+- CPU expert 被正确跳过（mask=False → SKIP）
+- TP2 服务成功启动，英文/中文/长文本输出全部正常
+
+### 关键代码路径
+- GPU 权重创建：`UnquantizedFusedMoEMethod.create_weights()` → `w2_weight = torch.empty(...)`（未初始化）
+- GPU 权重加载：sglang 标准 weight_loader 只处理 `experts.N.down_proj.` 格式 key
+- CPU 权重加载：kt-kernel 的 `SafeTensorLoader`（`kt-kernel/python/utils/loader.py`）解析 AMXINT4 格式 key
+- KTEPWrapperMethod：`create_weights` 调用 `gpu_method.create_weights`，`process_weights_after_loading` 调用 `wrapper.load_weights`（CPU 权重）
+
+### 教训
+1. **`--model` 必须指向标准 bf16 checkpoint**，sglang 标准 weight_loader 才能识别并加载 GPU 权重
+2. **`--kt-weight-path` 指向 AMXINT4 checkpoint**，供 kt-kernel 的 AMX 内核使用
+3. **MESH 模式下同样需要修复 `--model` 路径**，因为 GPU 专家（GE>0）仍需要从标准 checkpoint 加载 w2_weight
+4. **`torch.empty` 不会自动清零**，未初始化的权重在 TP>1 时可能包含 NaN/Inf
+5. **本地代码必须与远程同步**：远程 .venv 是 pip install 的旧版，本地子模块是最新 commit，需要 rsync 同步
+
+### 排查过程中排除的假设
+- ❌ 不是输入 NaN/Inf 传播
+- ❌ 不是 w13_weight NaN
+- ❌ 不是 topk_weights NaN
+- ❌ 不是 CPU 专家问题
+- ❌ 不是 `--kt-enable-dynamic-expert-update` 导致
+- ❌ 不是 `_prepare_weight_bf16` 导致（`load()` 未被调用）
+- ❌ 不是 `process_weights_after_loading` 导致（bf16 no-op）
+- ✅ **是 `--model` 指向 AMXINT4 格式 checkpoint，GPU w2_weight 从未被填充**
+
+### 涉及文件
+- `third_party/sglang/python/sglang/srt/layers/moe/fused_moe_triton/layer.py` — weight_loader 实现
+- `third_party/sglang/python/sglang/srt/layers/moe/kt_ep_wrapper.py` — KTEPWrapperMethod
+- `third_party/sglang/python/sglang/srt/layers/quantization/unquant.py` — UnquantizedFusedMoEMethod
+- `third_party/sglang/python/sglang/srt/models/qwen3_5.py` — 模型定义
+- `kt-kernel/python/utils/loader.py` — SafeTensorLoader（CPU 权重加载）
+- `kt-kernel/python/utils/amx.py` — AMXMoEWrapper

align_safetensors.py

@@ -0,0 +1,61 @@
+#!/usr/bin/env python3
+"""Align safetensors header to 512-byte boundary for MESH O_DIRECT."""
+import struct
+import os
+import sys
+import shutil
+
+
+def align_safetensors(src, dst):
+    with open(src, 'rb') as f:
+        header_len = struct.unpack('<Q', f.read(8))[0]
+        header_bytes = f.read(header_len)
+        # Read rest in chunks to avoid loading 1.5G into memory
+        data = f.read()
+
+    # Pad header with spaces so data_start = 8 + padded_len is 512-aligned
+    target = ((header_len + 8 + 511) // 512) * 512
+    padded_len = target - 8
+    padding = padded_len - header_len
+    padded_header = header_bytes + b' ' * padding
+
+    with open(dst, 'wb') as f:
+        f.write(struct.pack('<Q', padded_len))
+        f.write(padded_header)
+        f.write(data)
+
+
+def main():
+    src_dir = sys.argv[1]
+    dst_dir = sys.argv[2]
+    os.makedirs(dst_dir, exist_ok=True)
+
+    # Copy non-safetensors files (config.json, etc.)
+    for f in os.listdir(src_dir):
+        if not f.endswith('.safetensors') and not f.endswith('.tmp'):
+            src_f = os.path.join(src_dir, f)
+            dst_f = os.path.join(dst_dir, f)
+            if os.path.isfile(src_f):
+                shutil.copy2(src_f, dst_f)
+
+    # Align safetensors files
+    files = sorted([f for f in os.listdir(src_dir)
+                    if f.endswith('.safetensors') and not f.endswith('.tmp')])
+
+    for i, f in enumerate(files):
+        src = os.path.join(src_dir, f)
+        dst = os.path.join(dst_dir, f)
+        align_safetensors(src, dst)
+        # Verify alignment
+        with open(dst, 'rb') as fp:
+            hl = struct.unpack('<Q', fp.read(8))[0]
+            aligned = (hl + 8) % 512 == 0
+        size_gib = os.path.getsize(dst) / 1024 ** 3
+        print(f'[{i+1}/{len(files)}] {f} -> {size_gib:.2f} GiB, aligned={aligned}')
+        sys.stdout.flush()
+
+    print('All done!')
+
+
+if __name__ == '__main__':
+    main()

bench_decode_stream.py

@@ -0,0 +1,74 @@
+#!/usr/bin/env python3
+"""Decode-only benchmark using stream mode.
+Usage: python3 bench_decode_stream.py <port> [max_tokens]
+Returns: decode_tok_s (excluding prefill/first-token latency)
+"""
+import requests, time, json, sys
+
+port = sys.argv[1] if len(sys.argv) > 1 else "50052"
+max_tokens = int(sys.argv[2]) if len(sys.argv) > 2 else 512
+url = f"http://localhost:{port}/v1/completions"
+
+# Warmup
+try:
+    r = requests.post(url, json={
+        "model": "default", "prompt": "Hello",
+        "max_tokens": 16, "temperature": 0
+    }, timeout=60)
+    if r.status_code != 200:
+        print(f"WARMUP_FAIL: {r.status_code}")
+        sys.exit(1)
+except Exception as e:
+    print(f"WARMUP_FAIL: {e}")
+    sys.exit(1)
+
+time.sleep(2)
+
+# Decode benchmark with stream
+data = {
+    "model": "default",
+    "prompt": "Write a detailed essay about artificial intelligence and its impact on society, including historical context, current developments, and future implications:",
+    "max_tokens": max_tokens,
+    "temperature": 0,
+    "stream": True,
+}
+
+start = time.time()
+first_token_time = None
+token_count = 0
+
+try:
+    r = requests.post(url, json=data, stream=True, timeout=300)
+    for line in r.iter_lines():
+        if line:
+            line = line.decode()
+            if line.startswith("data: "):
+                line = line[6:]
+                if line == "[DONE]":
+                    break
+                try:
+                    chunk = json.loads(line)
+                    text = chunk["choices"][0].get("text", "")
+                    if text:
+                        token_count += 1
+                        if first_token_time is None:
+                            first_token_time = time.time()
+                except:
+                    pass
+except Exception as e:
+    print(f"DECODE_FAIL: {e}")
+    sys.exit(1)
+
+end = time.time()
+
+if first_token_time and token_count > 1:
+    decode_elapsed = end - first_token_time
+    decode_tps = (token_count - 1) / decode_elapsed  # exclude first token (prefill)
+    total_elapsed = end - start
+    print(f"decode_tok_s: {decode_tps:.2f}")
+    print(f"decode_tokens: {token_count}")
+    print(f"total_elapsed: {total_elapsed:.2f}s")
+    print(f"decode_elapsed: {decode_elapsed:.2f}s")
+else:
+    print(f"DECODE_FAIL: no tokens generated")
+    sys.exit(1)