kvcache-ai / kvcache-ai/ktransformers
GeneralMoEWrapper.load_weights() cpu_save path reads nonexistent dict keys, guaranteed KeyError
- 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
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