No Significant Improvement Observed in Model Training Speed

Nobody has claimed this yet.

I am trying to speedup inference and training of a `mistralai/Mistral-Small-3.1-24B-Instruct-2503` model.

Simply replacing `AutoModelForCausalLM` with `AutoLigerKernelForCausalLM` does not lead to any speedup in my sampling speed or memory usage. I am also using DeepSpeed for distributed training.

```
model = AutoLigerKernelForCausalLM.from_pretrained(
"mistralai/Mistral-Small-3.1-24B-Instruct-2503",
torch_dtype=TorchDtype.float32,
attn_implementation="sdpa",
)
```

I have also tried this with the same result:

```
model = AutoModelForCausalLM.from_pretrained(
"mistralai/Mistral-Small-3.1-24B-Instruct-2503",
torch_dtype=TorchDtype.float32,
attn_implementation="sdpa",
)
apply_liger_kernel_to_mistral(
rope=True,
cross_entropy=False,
fused_linear_cross_entropy=True,
rms_norm=True,
swiglu=True,
model=model,
)
```

Am I missing anything? Thanks for any help.
Should I expect to see the speedup and memory optimization in the autoregressive generative sampling or in the backward pass? or in both?

########################

Python version: 3.12.9
PyTorch version: 2.6.0+cu124
CUDA version: 12.4
Triton version: 3.2.0
Transformers version: 4.51.1
DeepSpeed version: 0.15.4

Start by reproducing the Mistral-Small-3.1-24B setup using AutoLigerKernelForCausalLM and apply_liger_kernel_to_mistral with the listed PyTorch, CUDA, Triton, Transformers, and DeepSpeed versions. Compare autoregressive sampling, training forward, backward, and memory usage against AutoModelForCausalLM, then establish which measurements should improve and document the cause of any discrepancy.

Written by the indexing model from the issue text.