Summary
servekit prepare followed by servekit launch -- ... --load-format sharded_state cannot serve an MLA
model. On Kimi-K3 at TP=4 / PP=8 / EP=4 across 8 nodes, prepare succeeds and writes a complete 1.5 TB
checkpoint (32 ranks, all gates pass), and the serve job then dies during model init on every pipeline
stage with:
KeyError: 'language_model.model.layers.67.self_attn.attn_mha.kv_b_proj.weight'
ShardedStateLoader.load_model -> param_data = state_dict[key].data
Unlike #14 nothing is silently zeroed -- this is a hard failure, which is the good case. But it is the same
family: the checkpoint is written from a model whose module graph is not the one that reads it back.
This is an upstream SGLang interaction rather than a servekit bug per se -- prepare/launch only drive
SGLang's own ShardedStateLoader -- but it makes sharded_state presharding unusable for MLA models, and
servekit is where it surfaces.
Root cause
DeepSeek-V2-style MLA attention, which Kimi-K3 derives from, constructs attn_mha with an empty
kv_b_proj and fills it in lazily on the first forward:
# __init__
self.attn_mha = RadixAttention(...)
self.attn_mha.kv_b_proj = None
# first forward
if self.attn_mha.kv_b_proj is None:
self.attn_mha.kv_b_proj = self.kv_b_projSo whether attn_mha.kv_b_proj.weight appears in model.state_dict() depends on whether a forward has run:
| side |
has forwarded? |
attn_mha.kv_b_proj |
key in state_dict |
save (prepare serves until ready) |
yes |
aliased to self.kv_b_proj |
present -> written to the checkpoint |
| load (fresh model, pre-warmup) |
no |
None |
absent |
ShardedStateLoader.load_model iterates the checkpoint's keys and indexes the fresh model's state dict,
so a key the checkpoint has and the model does not is an immediate KeyError.
Summary
servekit preparefollowed byservekit launch -- ... --load-format sharded_statecannot serve an MLAmodel. On Kimi-K3 at TP=4 / PP=8 / EP=4 across 8 nodes,
preparesucceeds and writes a complete 1.5 TBcheckpoint (32 ranks, all gates pass), and the serve job then dies during model init on every pipeline
stage with:
Unlike #14 nothing is silently zeroed -- this is a hard failure, which is the good case. But it is the same
family: the checkpoint is written from a model whose module graph is not the one that reads it back.
This is an upstream SGLang interaction rather than a servekit bug per se --
prepare/launchonly driveSGLang's own
ShardedStateLoader-- but it makessharded_statepresharding unusable for MLA models, andservekit is where it surfaces.
Root cause
DeepSeek-V2-style MLA attention, which Kimi-K3 derives from, constructs
attn_mhawith an emptykv_b_projand fills it in lazily on the first forward:So whether
attn_mha.kv_b_proj.weightappears inmodel.state_dict()depends on whether a forward has run:attn_mha.kv_b_projprepareserves until ready)self.kv_b_projNoneShardedStateLoader.load_modeliterates the checkpoint's keys and indexes the fresh model's state dict,so a key the checkpoint has and the model does not is an immediate
KeyError.