sgl-project/sglang · critical · RuntimeError

Nemotron 3.5 DFLASH draft requires its checkpoint embedding.

Error message

Nemotron 3.5 DFLASH draft requires its checkpoint embedding.

What it means

Thrown by the DFLASH v2 worker when the draft model is flagged as a Nemotron 3.5 draft (is_nemotron_35_draft) but get_input_embeddings() returns None. Nemotron 3.5 DFLASH drafts must use their own checkpoint's embedding matrix instead of the target model's, so a missing embedding makes the worker unable to embed draft inputs. It usually indicates a broken/partially-loaded draft checkpoint or a model class that doesn't implement the embedding accessor correctly.

Source

Thrown at python/sglang/srt/speculative/dflash_worker_v2.py:185

        torch.gather(gathered_ids.view(self.tp_size, n), 0, best_rank, out=selected)
        self.out[:n].copy_(selected.view(-1))


def _commit_accept(candidates, accept_len, bonus_tokens):
    """The committed block: drafted tokens shifted left, the bonus at the accept
    boundary. Returns it with the commit lengths."""
    out_tokens = torch.empty_like(candidates, dtype=torch.int64)
    out_tokens[:, :-1].copy_(candidates[:, 1:])
    out_tokens[:, -1].fill_(0)
    out_tokens.scatter_(1, accept_len.to(torch.int64)[:, None], bonus_tokens[:, None])
    return out_tokens, accept_len.to(torch.int32) + 1


def _resolve_dflash_embedding_module(draft_model, target_model):
    if getattr(draft_model, "is_nemotron_35_draft", False):
        embed_module = draft_model.get_input_embeddings()
        if embed_module is None:
            raise RuntimeError(
                "Nemotron 3.5 DFLASH draft requires its checkpoint embedding."
            )
        return embed_module
    return target_model.get_input_embeddings()


def _is_all_greedy(sampling_info) -> bool:
    return sampling_info is None or sampling_info.is_all_greedy


def _selector_lattice(draft_model, pred_hidden, anchor_token_ids):
    # Flattened to [N, H] and viewed back because the radix top-k kernel is 2D.
    bs, num_pred = pred_hidden.shape[0], pred_hidden.shape[1]
    candidate_ids, unary_logits = draft_model.compute_candidates(
        pred_hidden.reshape(-1, pred_hidden.shape[-1])
    )
    candidate_ids = candidate_ids.view(bs, num_pred, -1)
    return candidate_ids, draft_model.candidate_selector.build_lattice(

View on GitHub (pinned to 0132848349)

Solutions

1. Verify the draft checkpoint actually contains the embedding weights (inspect the safetensors index for model.embed_tokens.weight) and that files aren't truncated
2. Check the draft model's get_input_embeddings() implementation returns the real nn.Embedding module (not None due to tying/config flags)
3. Re-download or re-convert the draft checkpoint with the embedding included
4. If embeddings are intentionally tied, fix the model class so get_input_embeddings resolves the tied weight

draft = worker.draft_worker.model
if getattr(draft, "is_nemotron_35_draft", False):
    assert draft.get_input_embeddings() is not None, (
        "Nemotron 3.5 draft checkpoint is missing its embedding weights"
    )

try:
    worker = DFlashWorker(...)
except RuntimeError as e:
    if "requires its checkpoint embedding" in str(e):
        log.error("Draft checkpoint incomplete; re-download the draft weights")
        raise

When it happens

Trigger: Loading a Nemotron 3.5 DFLASH draft checkpoint whose weights lack the embedding (e.g. sharded file missing, or tie_word_embeddings handling stripped it); a custom draft model class overriding get_input_embeddings to return None; quantized/converted checkpoints that drop embed_tokens.

Common situations: Swapping in a draft checkpoint converted or quantized with a tool that omitted embedding weights; mixing draft/target checkpoints of different revisions; custom Nemotron draft model implementations that never set the embedding module.

Related errors

• DFLASH mask_token_id is outside the target vocab size. mask_
• DFLASH speculative decoding only supports CUDA and NPU devic
• next_token_logits row count mismatch. Expected {bs * draft_t
• candidates and next_token_logits must be on the same device,
• uniform_samples shape mismatch. Expected {(bs, draft_token_n

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