Make sLSTM checkpoints portable across backends (#127)

[jaideeppyne]

What

Fixes #127. A checkpoint saved with one sLSTM backend cannot be loaded into a cell using another backend:

from xlstm.blocks.slstm.cell import sLSTMCell_cuda, sLSTMCell_vanilla, sLSTMCellConfig
cfg = sLSTMCellConfig(hidden_size=16, num_heads=4)
cuda = sLSTMCell_cuda(cfg, skip_backend_init=True)
sLSTMCell_vanilla(cfg, skip_backend_init=True).load_state_dict(cuda.state_dict())
# RuntimeError: size mismatch for _recurrent_kernel_: ... torch.Size([4, 4, 16]) vs torch.Size([4, 16, 4])

Why

sLSTMCell_vanilla and sLSTMCell_cuda keep the recurrent kernel and bias in different internal layouts:

tensor	vanilla internal	cuda internal	canonical / external
recurrent kernel	(H, G·head_dim, head_dim)	(H, head_dim, G·head_dim)	(H, head_dim, G, head_dim)
bias	(H·G·head_dim,) gate-major	(H·G·head_dim,) head-major	(H, G, head_dim)

state_dict serializes the internal _recurrent_kernel_/_bias_ parameters directly. Loading across backends therefore fails: the recurrent kernel raises a size mismatch, and — more insidiously — the bias has the same flat shape in both backends, so it loads without error but with a scrambled gate/head ordering (silent weight corruption).

The cell already defines a single canonical (external) layout via its _..._int2ext/_..._ext2int conversions, exposes it through the recurrent_kernel/bias proxies, and sLSTMCell_cuda.__init__ documents cross-backend conversion (skip_backend_init) as an intended workflow — so the checkpoint simply needs to be stored in that canonical layout.

How

Two hooks registered in sLSTMCellBase.__init__:

• _register_state_dict_hook — on save, convert _recurrent_kernel_/_bias_ to the canonical external layout (int2ext).
• _register_load_state_dict_pre_hook — on load, convert an external-layout checkpoint back into the current backend's internal layout (ext2int).

External tensors are identified by rank (recurrent kernel 4D, bias 3D). Legacy checkpoints saved in the old internal layout (recurrent kernel 3D, bias 1D) are left untouched, so they still load into a cell of their original backend — the change is backward compatible.

Verification (CPU, via skip_backend_init=True)

• cuda → vanilla and vanilla → cuda now load with no error, and the canonical weights match exactly (allclose) — including the bias, which was previously corrupted silently.
• Same-backend save/load round-trips remain bit-for-bit identical.
• A legacy internal-layout checkpoint still loads into a cell of its original backend.

Tests

tests/test_slstm_backend_state_dict.py covers all four cases above. The two cross-backend tests fail on main (RuntimeError / corrupted bias) and pass with the fix; same-backend and legacy tests pass both ways. They run on CPU.

---

Disclosure: this change was prepared with AI assistance and reviewed/verified by me before submission.

[github-actions]

All contributors have signed the CLA ✍️ ✅
_{Posted by the CLA Assistant Lite bot.}

[jaideeppyne]

I have read the CLA Document and I hereby sign the CLA

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 93efad0246

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[chatgpt-codex-connector]

Materialize canonical tensors before returning the state dict

When a vanilla-backed cell is exported through a serializer that requires contiguous tensors, such as safetensors.torch.save_file(model.state_dict(), ...), this hook now causes saving to fail. Both vanilla int2ext methods return non-contiguous views produced by permute, and the hook places those views directly in the state dict; typical multi-head configurations therefore trigger safetensors' non-contiguous-tensor rejection. Materialize the converted values with contiguous() before storing them.

Useful? React with 👍 / 👎.

[jaideeppyne]

Good catch, thanks. The vanilla int2ext conversions return non-contiguous permute views, so the exported state_dict held non-contiguous tensors and safetensors.torch.save_file would reject them.

Fixed in a7e54a9 by calling .contiguous() on the converted recurrent kernel and bias in the state-dict hook. Verified that safetensors.torch.save_file(cell.state_dict(), ...) now succeeds for both backends and that a cell of the other backend still loads the resulting checkpoint. Added a test asserting the exported tensors are contiguous.