kvcache-ai / kvcache-ai/ktransformers

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

Open
#2,151 1 comment 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
19.5k
Forks
1.6k
Avg merge
19h 32m
Merged PRs (30d)
27

Description

### 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.

Contributor guide

Open the contributing guide

Research direction

Start in kt-kernel/python/utils/moe_kernel.py at GeneralMoEWrapper.load_weights(), then compare its cpu_save branch with SafeTensorLoader.load_experts() in kt-kernel/python/utils/loader.py and the existing load_weights_from_tensors() pattern. Done means the cpu_save path uses the loader's returned keys and produces the expected contiguous tensor layout without the reported KeyError.

Written by the indexing model from the issue text.

Assessment

Tech stack
numpy, python
Domain
machine-learning
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Quiet
Clarity
Clearly specified
Newbie friendliness
25/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.