sLSTM `_recurrent_kernel_` layout differs between `cuda` and `vanilla` backends, breaking checkpoint loading across backends

[gabrielhae]

Summary

The exposed shape of sLSTMCell._recurrent_kernel_ differs depending on the backend chosen in sLSTMLayerConfig, even for the identical config. A model trained with one backend fails to load_state_dict when the cell is instantiated with the other, because the same parameter is exposed in a different (transposed) layout.

Environment

• xlstm version: 2.0.5
• torch: 2.13.0+cu126
• python: 3.14, Linux (CUDA available)

Repro

import torch
from xlstm.blocks.slstm.cell import sLSTMCell
from xlstm.blocks.slstm.layer import sLSTMLayerConfig

cfg = sLSTMLayerConfig(embedding_dim=16, num_heads=4,
                       conv1d_kernel_size=4, batch_size=8)

for backend in ("vanilla", "cuda"):
    cell = sLSTMCell(sLSTMLayerConfig(embedding_dim=16, num_heads=4,
                      conv1d_kernel_size=4, batch_size=8, backend=backend))
    print(backend, tuple(cell._recurrent_kernel_.shape))

Output:

vanilla (4, 16, 4)
cuda    (4, 4, 16)

Expected vs actual

• Expected: for the same config, the exposed parameter shape should be consistent across backends, so a checkpoint saved from a cuda-trained
  model can be loaded into a vanilla-instantiated cell (and vice versa).
• Actual: _recurrent_kernel_ is exposed as (num_heads, num_gates*head_dim, head_dim) for vanilla but (num_heads, head_dim, num_gates*head_dim) for cuda — a transpose of dims 1 and 2.

Impact

Cross-backend loading is not possible out of the box. A checkpoint trained on a GPU worker (backend="cuda") fails to load on a CPU-only machine where the cell is built with backend="vanilla":

RuntimeError: Error(s) in loading state_dict for ...:
  size mismatch for ...slstm_cell._recurrent_kernel_:
    copying a param with shape torch.Size([4, 4, 16]) from checkpoint,
    the shape in current model is torch.Size([4, 16, 4])

Notes

I traced the difference to sLSTMCell_cuda._recurrent_kernel_int2ext (and the corresponding ext2int) reshaping the ParameterProxy-exposed tensor differently from the base sLSTMCell. It would be helpful if the exposed external layout were identical across backends (or documented as
backend-specific), so loaders can be backend-agnostic.