[Bug] `sharded_state` cannot save and load an MLA model

[youssef62]

Checklist

• I searched related issues but found no solution.
• The bug persists in the latest version.
• Issues without environment info and a minimal reproducible demo are hard to resolve and may receive no feedback.
• If this is not a bug report but a general question, please start a discussion at https://github.com/sgl-project/sglang/discussions. Otherwise, it will be closed.
• Please use English. Otherwise, it will be closed.

Describe the bug

This issue is the same as the closed #25332 for which pr #25373 exists but did not get much attention.

Saving a sharded_state checkpoint of a DeepSeek-V2-style (MLA) model after any forward has run writes
the kv_b_proj weight under self_attn.attn_mha.kv_b_proj.weight and never under
self_attn.kv_b_proj.weight. A freshly constructed model has no such key, so loading that checkpoint back
with load_format="sharded_state" dies during model init:

File "python/sglang/srt/model_loader/loader.py", line 1796, in load_model
    param_data = state_dict[key].data
                 ~~~~~~~~~~^^^^^
KeyError: 'model.layers.22.self_attn.attn_mha.kv_b_proj.weight'

Reproduction

Reproduced on a single GPU with deepseek-ai/DeepSeek-V2-Lite, TP=1, PP=1.

Repro

"""
usage:
    python3 mla_sharded_state_repro.py save /path/to/DeepSeek-V2-Lite /path/to/out [--skip-forward]
    python3 mla_sharded_state_repro.py load /path/to/out

"""

import os
import shutil
import sys

from sglang import Engine


def main():
    if sys.argv[1] == "save":
        model_path, out_dir = sys.argv[2], sys.argv[3]
        os.makedirs(out_dir, exist_ok=True)

        llm = Engine(
            model_path=model_path,
            trust_remote_code=True,
            attention_backend="triton",
            disable_cuda_graph=True,
        )
        if "--skip-forward" not in sys.argv:
            llm.generate("hi", {"max_new_tokens": 1})
        llm.save_sharded_model(path=out_dir, pattern=None, max_size=5 * 1024**3)
        llm.shutdown()

        for name in os.listdir(model_path):
            if os.path.splitext(name)[1] not in (".bin", ".pt", ".safetensors"):
                shutil.copy(f"{model_path}/{name}", f"{out_dir}/{name}")

    else:  # load
        llm = Engine(
            model_path=sys.argv[2],
            load_format="sharded_state",
            trust_remote_code=True,
            attention_backend="triton",
        )
        print(llm.generate("hi", {"max_new_tokens": 1}))
        llm.shutdown()
        print("=== loaded fine ===")


if __name__ == "__main__":
    main()

Run:

python3 mla_sharded_state_repro.py save /path/to/DeepSeek-V2-Lite /path/to/out
python3 mla_sharded_state_repro.py load /path/to/out

# control: saving before any forward round-trips fine
python3 mla_sharded_state_repro.py save /path/to/DeepSeek-V2-Lite /path/to/ctl --skip-forward
python3 mla_sharded_state_repro.py load /path/to/ctl

Result

sglang 0.0.0.dev1+gc4271c3fe (lmsysorg/sglang:nightly-dev-cu13), GH200 aarch64, CUDA 13, 1 GPU,
deepseek-ai/DeepSeek-V2-Lite (27 layers), TP=1 PP=1:

save after a forward:

KeyError: 'model.layers.22.self_attn.attn_mha.kv_b_proj.weight'

--skip-forward:

=== loaded fine ===

Environment

Python: 3.12.3 (main, Jun 19 2026, 12:46:00) [GCC 13.3.0]
CUDA available: True
GPU 0: NVIDIA GH200 120GB
GPU 0 Compute Capability: 9.0
CUDA_HOME: /usr/local/cuda
NVCC: Cuda compilation tools, release 13.0, V13.0.88
CUDA Driver Version: 590.48.01
PyTorch: 2.13.0+cu130
sglang: 0.0.0.dev1+gc4271c3fe
sglang-kernel: 0.4.6.post1
flashinfer_python: 0.6.17
flashinfer_cubin: 0.6.17
flashinfer_jit_cache: 0.6.17+cu130
triton: 3.7.1
transformers: 5.12.1
torchao: 0.17.0
numpy: 2.3.5
aiohttp: 3.14.3
fastapi: 0.141.1
huggingface_hub: 1.27.0
interegular: 0.3.3
modelscope: 1.39.1
orjson: 3.11.9
outlines: 0.1.11
packaging: 26.3
psutil: 7.2.2
pydantic: 2.13.4
python-multipart: 0.0.32
pyzmq: 27.1.0
uvicorn: 0.52.3
uvloop: 0.22.1
vllm: Module Not Found
xgrammar: 0.2.1
openai: 2.6.1
tiktoken: 0.13.0
anthropic: 0.122.0
litellm: Module Not Found
torchcodec: Module Not Found
NVIDIA Topology:
�[4mGPU0 GPU1 GPU2 GPU3 CPU Affinity NUMA Affinity GPU NUMA ID�[0m
GPU0 X NV6 NV6 NV6 0-7 0 4
GPU1 NV6 X NV6 NV6 1 12
GPU2 NV6 NV6 X NV6 2 20
GPU3 NV6 NV6 NV6 X 3 28

Legend:

X = Self
SYS = Connection traversing PCIe as well as the SMP interconnect between NUMA nodes (e.g., QPI/UPI)
NODE = Connection traversing PCIe as well as the interconnect between PCIe Host Bridges within a NUMA node
PHB = Connection traversing PCIe as well as a PCIe Host Bridge (typically the CPU)
PXB = Connection traversing multiple PCIe bridges (without traversing the PCIe Host Bridge)
PIX = Connection traversing at most a single PCIe bridge
NV# = Connection traversing a bonded set of # NVLinks

ulimit soft: 1000000