[GH-ISSUE #2151] GeneralMoEWrapper.load_weights() cpu_save path reads nonexistent dict keys, guaranteed KeyError #1294

New issue

Open

opened 2026-10-29 14:50:55 +00:00 by bsakaguchi · 1 comment

bsakaguchi commented

2026-10-29 14:50:55 +00:00

Owner

Originally created by @AmirF194 on GitHub (Aug 9, 2026).
Original GitHub issue: https://github.com/kvcache-ai/ktransformers/issues/2151

Reminder

I have read the above rules and searched the existing issues.

System Info

Not hardware-dependent: the crash happens in pure Python before any native kt_kernel_ext call, so it reproduces on any platform with the package importable (kt-kernel Python layer, any CPU variant).

Reproduction

GeneralMoEWrapper.load_weights() (kt-kernel/python/utils/moe_kernel.py) has two branches that call self.safetensor_loader.load_experts(base_key). The first, in the if self.load_merged_weight: block, reads the returned dict correctly:

self.gate_weights = w["gate"]
self.up_weights = w["up"]
self.down_weights = w["down"]

The second, in the if self.cpu_save: block about 80 lines later, reads different keys that load_experts() never returns:

self.gate_proj = torch.cat(w["gate_weight"], dim=0).contiguous()
self.up_proj = torch.cat(w["up_weight"], dim=0).contiguous()
self.down_proj = torch.cat(w["down_weight"], dim=0).contiguous()

SafeTensorLoader.load_experts() (kt-kernel/python/utils/loader.py) returns a dict with keys up, gate, down, up_scale, gate_scale, down_scale (no _weight suffix), so load_weights() raises KeyError: 'gate_weight' whenever it is called with cpu_save=True and merged safetensor weights are present. Minimal repro against current main (constructing GeneralMoEWrapper with a mocked loader returning that exact key set, since the crash happens before any native call is reached):

Traceback (most recent call last):
  File "kt_kernel/utils/moe_kernel.py", line 289, in load_weights
    self.gate_proj = torch.cat(w["gate_weight"], dim=0).contiguous()
                               ~^^^^^^^^^^^^^^^
KeyError: 'gate_weight'

Both keys trace back to the same introducing commit (9bc00e5, #1587) and have never been touched since; the correctly-keyed branch 80 lines earlier in the same function was never updated to match.

An identical, byte-for-byte copy of the same wrong keys exists in kt-kernel/python/utils/amx.py (AMXMoEWrapper, lines ~514-516), also from 9bc00e5. I'm leaving that one alone since amx.py is currently touched by open PR #2111.

Others

I have a fix ready (correct key names plus converting the loader's per-expert numpy arrays into the contiguous [expert_num, intermediate, hidden] tensor the native kernel expects, matching load_weights_from_tensors()'s existing torch.stack pattern) and will open a PR referencing this issue.

Originally created by @AmirF194 on GitHub (Aug 9, 2026). Original GitHub issue: https://github.com/kvcache-ai/ktransformers/issues/2151 ### Reminder - [x] I have read the above rules and searched the existing issues. ### System Info Not hardware-dependent: the crash happens in pure Python before any native kt_kernel_ext call, so it reproduces on any platform with the package importable (kt-kernel Python layer, any CPU variant). ### Reproduction `GeneralMoEWrapper.load_weights()` (`kt-kernel/python/utils/moe_kernel.py`) has two branches that call `self.safetensor_loader.load_experts(base_key)`. The first, in the `if self.load_merged_weight:` block, reads the returned dict correctly: ```python self.gate_weights = w["gate"] self.up_weights = w["up"] self.down_weights = w["down"] ``` The second, in the `if self.cpu_save:` block about 80 lines later, reads different keys that `load_experts()` never returns: ```python self.gate_proj = torch.cat(w["gate_weight"], dim=0).contiguous() self.up_proj = torch.cat(w["up_weight"], dim=0).contiguous() self.down_proj = torch.cat(w["down_weight"], dim=0).contiguous() ``` `SafeTensorLoader.load_experts()` (`kt-kernel/python/utils/loader.py`) returns a dict with keys `up`, `gate`, `down`, `up_scale`, `gate_scale`, `down_scale` (no `_weight` suffix), so `load_weights()` raises `KeyError: 'gate_weight'` whenever it is called with `cpu_save=True` and merged safetensor weights are present. Minimal repro against current `main` (constructing `GeneralMoEWrapper` with a mocked loader returning that exact key set, since the crash happens before any native call is reached): ```text Traceback (most recent call last): File "kt_kernel/utils/moe_kernel.py", line 289, in load_weights self.gate_proj = torch.cat(w["gate_weight"], dim=0).contiguous() ~^^^^^^^^^^^^^^^ KeyError: 'gate_weight' ``` Both keys trace back to the same introducing commit (`9bc00e5`, `#1587`) and have never been touched since; the correctly-keyed branch 80 lines earlier in the same function was never updated to match. An identical, byte-for-byte copy of the same wrong keys exists in `kt-kernel/python/utils/amx.py` (`AMXMoEWrapper`, lines ~514-516), also from `9bc00e5`. I'm leaving that one alone since `amx.py` is currently touched by open PR #2111. ### Others I have a fix ready (correct key names plus converting the loader's per-expert numpy arrays into the contiguous `[expert_num, intermediate, hidden]` tensor the native kernel expects, matching `load_weights_from_tensors()`'s existing `torch.stack` pattern) and will open a PR referencing this issue.

bsakaguchi commented

2026-10-29 14:50:56 +00:00

Author

Owner

@Anai-Guo commented on GitHub (Aug 18, 2026):

Thanks for flagging the identical copy in amx.py here. I've opened #2161 to fix the AMXMoEWrapper.load_weights() cpu_save branch the same way #2152 handles GeneralMoEWrapper — switching the key to blk.{layer_idx} (the model.layers. prefix and its fallback both ValueError in SafeTensorLoader.load_experts()) and stacking the per-expert numpy arrays from w["gate"|"up"|"down"][0] instead of the nonexistent *_weight keys. It only touches those lines, so it's orthogonal to #2111.

@Anai-Guo commented on GitHub (Aug 18, 2026): Thanks for flagging the identical copy in `amx.py` here. I've opened #2161 to fix the `AMXMoEWrapper.load_weights()` `cpu_save` branch the same way #2152 handles `GeneralMoEWrapper` — switching the key to `blk.{layer_idx}` (the `model.layers.` prefix and its fallback both `ValueError` in `SafeTensorLoader.load_experts()`) and stacking the per-expert numpy arrays from `w["gate"|"up"|"down"][0]` instead of the nonexistent `*_weight` keys. It only touches those lines, so it's orthogonal to #2111.