sgl-project/sglang · error · NotImplementedError

Eagle3 MLA draft post_load_weights only supports float dtype

Error message

Eagle3 MLA draft post_load_weights only supports float dtypes, got {w.dtype}

What it means

Eagle3MLA.post_load_weights re-implements the DeepseekV2 bf16 fast path for splitting kv_b_proj into w_kc/w_vc, and it requires the weight to be a float dtype (bf16/fp16/fp32) because it calls .unflatten/.split on the raw tensor. Quantized or integer-packed kv_b_proj weights (GGUF/awq/marlin tensors) hit this NotImplementedError.

Source

Thrown at python/sglang/srt/models/kimi_k25_eagle3.py:510

                logger.warning("Eagle3 MLA: skipping unexpected weight %s", name)
                continue
            param = params_dict[mapped_name]
            weight_loader = getattr(param, "weight_loader", default_weight_loader)
            weight_loader(param, loaded_weight)

        self.post_load_weights()

    def post_load_weights(self) -> None:
        """Split kv_b_proj into w_kc / w_vc tensors used by MLA absorb_core.

        DeepseekV2 normally does this in DeepseekV2WeightLoaderMixin.post_load_weights;
        we re-implement the bf16 fast-path directly here to keep the eagle3 draft
        path independent of the full DeepseekV2 weight loader.
        """
        self_attn = self.model.midlayer.self_attn
        w = self_attn.kv_b_proj.weight
        if w.dtype not in (torch.bfloat16, torch.float16, torch.float32):
            raise NotImplementedError(
                f"Eagle3 MLA draft post_load_weights only supports float dtypes, got {w.dtype}"
            )
        w_kc, w_vc = w.unflatten(
            0, (-1, self_attn.qk_nope_head_dim + self_attn.v_head_dim)
        ).split([self_attn.qk_nope_head_dim, self_attn.v_head_dim], dim=1)
        self_attn.w_kc = w_kc.transpose(1, 2).contiguous().transpose(1, 2)
        self_attn.w_vc = w_vc.contiguous().transpose(1, 2)


EntryClass = [Eagle3DeepseekV2ForCausalLM]

View on GitHub (pinned to 0132848349)

Solutions

  1. Use an unquantized (bf16/fp16) EAGLE3 draft checkpoint for kv_b_proj.
  2. Pre-cast kv_b_proj to torch.bfloat16 before load if you authored the checkpoint.
  3. Extend post_load_weights to dequantize integer dtypes before the split if quantized drafts must be supported.

Example fix

// before
# kv_b_proj.weight loaded as GGUF/quantized dtype -> NotImplementedError

// after
# convert to bf16 when preparing the draft checkpoint
w = w.to(torch.bfloat16)
Defensive patterns

Strategy: validation

Validate before calling

import torch, safetensors
w = safetensors.torch.load_file(draft_path + "/kv_b_proj.safetensors")["kv_b_proj.weight"]
assert w.dtype in (torch.bfloat16, torch.float16, torch.float32), f"bad dtype {w.dtype}"

Prevention

When it happens

Trigger: Loading an EAGLE3 MLA draft checkpoint whose midlayer self_attn.kv_b_proj.weight is not bfloat16/float16/float32 — e.g. a GGUF-quantized or otherwise integer-dtype kv_b_proj — during load_weights -> post_load_weights.

Common situations: Loading a quantized (GGUF/awq/GPTQ) Kimi K2.5 EAGLE3 draft; mixing a quantized target-model loader with the eagle3 draft path; exotic torch versions/serialization that round-trip weights as uint8/int8.

Related errors


AI-assisted analysis of sgl-project/sglang@0132848349 (2026-08-28). Data as JSON: /api/errors/9d105f45e93c01d9. Report an issue: GitHub.