GeneralMoEWrapper.load_weights() cpu_save path reads nonexistent dict keys, guaranteed KeyError

Author: AmirF194Created Aug 9, 2026Updated Aug 18, 2026

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:

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):

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.

Source: kvcache-ai/ktransformers