Skip to content

vllm.model_executor.models.qwen3_dspark

Qwen3 DSpark draft model for semi-autoregressive drafting.

DSpark drafts a whole block in one parallel pass (DFlash-style: context-KV precompute + a non-causal query-block forward) and then injects intra-block dependency with a lightweight sequential Markov head.

The parallel backbone is a standard Qwen3 decoder stack reused from the DFlash Qwen3 draft (see qwen3_dflash.py). DSpark adds: * markov_head: low-rank V x r / r x V transition bias added to the base logits, sampled left-to-right by the speculator (the sequential stage).

DSparkMarkovHead is shared with the DSV4-style DSpark model.

Classes:

DSparkMarkovHead

Bases: Module

Sequential transition-bias head (low-rank V x r, r x V).

markov_w1[token] embeds the previously sampled token (target vocab, vocab_size); markov_w2 projects it to a draft-vocab bias (draft_vocab_size) added to the base draft logits. The two sizes coincide for full-vocab drafts.

Methods:

  • bias –

    Vocab-size transition bias from a Markov embedding ([B, r] -> [B, V]).

  • embed –

    r-dim Markov embedding of token_ids ([B] -> [B, r]).

Source code in vllm/model_executor/models/qwen3_dspark.py 
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
class DSparkMarkovHead(nn.Module):
    """Sequential transition-bias head (low-rank V x r, r x V).

    ``markov_w1[token]`` embeds the previously sampled token (target vocab,
    ``vocab_size``); ``markov_w2`` projects it to a draft-vocab bias
    (``draft_vocab_size``) added to the base draft logits. The two sizes
    coincide for full-vocab drafts.
    """

    def __init__(
        self,
        vocab_size: int,
        draft_vocab_size: int,
        markov_rank: int,
        prefix: str,
    ) -> None:
        super().__init__()
        # TODO(ben): profile for which (if any) it makes sense to replicate or TP-shard
        self.markov_w1 = VocabParallelEmbedding(
            vocab_size, markov_rank, prefix=maybe_prefix(prefix, "markov_w1")
        )
        self.markov_w2 = ParallelLMHead(
            draft_vocab_size, markov_rank, prefix=maybe_prefix(prefix, "markov_w2")
        )

    def embed(self, token_ids: torch.Tensor) -> torch.Tensor:
        """r-dim Markov embedding of ``token_ids`` ([B] -> [B, r])."""
        return self.markov_w1(token_ids)

    def bias(self, markov_embed: torch.Tensor, logits_processor) -> torch.Tensor:
        """Vocab-size transition bias from a Markov embedding ([B, r] -> [B, V])."""
        return logits_processor(self.markov_w2, markov_embed)

bias(markov_embed, logits_processor)

Vocab-size transition bias from a Markov embedding ([B, r] -> [B, V]).

Source code in vllm/model_executor/models/qwen3_dspark.py 
65
66
67
def bias(self, markov_embed: torch.Tensor, logits_processor) -> torch.Tensor:
    """Vocab-size transition bias from a Markov embedding ([B, r] -> [B, V])."""
    return logits_processor(self.markov_w2, markov_embed)

embed(token_ids)

r-dim Markov embedding of token_ids ([B] -> [B, r]).

Source code in vllm/model_executor/models/qwen3_dspark.py 
61
62
63
def embed(self, token_ids: torch.Tensor) -> torch.Tensor:
    """r-dim Markov embedding of ``token_ids`` ([B] -> [B, r])."""
    return self.markov_w1(token_ids)

Qwen3DSparkModel

Bases: DFlashQwen3Model

DFlash Qwen3 backbone + DSpark Markov head.

Source code in vllm/model_executor/models/qwen3_dspark.py 
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
class Qwen3DSparkModel(DFlashQwen3Model):
    """DFlash Qwen3 backbone + DSpark Markov head."""

    def __init__(
        self,
        *,
        vllm_config: VllmConfig,
        start_layer_id: int = 0,
        prefix: str = "",
    ) -> None:
        super().__init__(
            vllm_config=vllm_config, start_layer_id=start_layer_id, prefix=prefix
        )
        config = self.config
        draft_vocab_size = (
            getattr(config, "draft_vocab_size", None) or config.vocab_size
        )
        self.markov_head = DSparkMarkovHead(
            config.vocab_size,
            draft_vocab_size,
            config.markov_rank,
            prefix=maybe_prefix(prefix, "markov_head"),
        )