[None][feat] Add Qwen3-based DSpark drafter (DeepSpec dense checkpoints)

tensorrt_llm/_torch/models/dspark/draft.py

@@ -38,6 +38,27 @@
 from .heads import confident_prefix_length
 
 
+def resolve_noise_token_id(mask_token_id: Optional[int], config, ckpt_attr: str) -> int:
+    """Resolve the DSpark noise/mask token id.
+
+    ``mask_token_id`` is either the speculative config's value (set explicitly
+    or resolved from the checkpoint during ``DSparkDecodingConfig`` validation)
+    or ``None``, in which case it falls back to the drafter checkpoint's
+    ``ckpt_attr`` attribute. Both missing is a checkpoint/config error: an
+    embedding id of ``vocab_size`` is out of range for ``embed_tokens`` and
+    would only surface later as an opaque device-side assert.
+    """
+    if mask_token_id is None:
+        mask_token_id = getattr(config, ckpt_attr, None)
+    if mask_token_id is None:
+        raise ValueError(
+            f"DSpark drafter checkpoint config has neither a resolved "
+            f"mask_token_id nor a `{ckpt_attr}` attribute; the DeepSpec "
+            f"checkpoint's config.json is expected to carry `{ckpt_attr}`."
+        )
+    return int(mask_token_id)
+
+
 def build_draft_input_ids(
     bonus_token_ids: torch.Tensor, *, block_size: int, noise_token_id: int
 ) -> torch.Tensor:

tensorrt_llm/_torch/models/dspark/heads.py

@@ -217,14 +217,22 @@ class DSparkConfidenceHead(nn.Module):
     Markov head's previous-token embedding. Output is a single logit per position.
     """
 
-    def __init__(self, *, hidden_size: int, markov_rank: int = 0, with_markov: bool = False):
+    def __init__(
+        self,
+        *,
+        hidden_size: int,
+        markov_rank: int = 0,
+        with_markov: bool = False,
+        bias: bool = False,
+    ):
         super().__init__()
         self.with_markov = bool(with_markov)
         input_dim = int(hidden_size) + (int(markov_rank) if with_markov else 0)
-        # The checkpoint stores ``proj`` as a bias-free bf16 weight, but the
+        # The V4-Pro checkpoint stores ``proj`` as a bias-free bf16 weight; the
+        # DeepSpec Qwen3 drafter checkpoints carry a bias. Either way the
         # confidence score is computed in fp32 (mirrors the DeepSpec reference
         # ``Linear(input_dim, 1, dtype=torch.float32)`` with the fp32 matmul).
-        self.proj = nn.Linear(input_dim, 1, bias=False, dtype=torch.float32)
+        self.proj = nn.Linear(input_dim, 1, bias=bias, dtype=torch.float32)
 
     def forward(
         self, hidden_states: torch.Tensor, prev_embeddings: Optional[torch.Tensor] = None

tensorrt_llm/_torch/models/modeling_speculative.py

@@ -1912,11 +1912,24 @@ def get_draft_model(model_config, draft_config, lm_head, model):
         return DFlashForCausalLM(draft_config)
     elif spec_dec_mode.is_dspark():
         # Lazy import to avoid a cycle (modeling_dspark -> modeling_deepseekv4 ->
-        # modeling_speculative). The DSpark draft reuses the target's aux streams.
-        # The draft stage count (n_mtp_layers) is not in the HF config, so derive
-        # it from the checkpoint's mtp.* namespace.
-        from .modeling_dspark import (DSparkForCausalLM, count_dspark_stages,
+        # modeling_speculative).
+        from .modeling_dspark import (DSparkForCausalLM, Qwen3DSparkForCausalLM,
+                                      count_dspark_stages,
                                       validate_dspark_eplb_layer_base)
+
+        # Dense drafters (e.g. dspark_qwen3_8b_block7) are separate checkpoints. The
+        # DeepSeek-V4 drafter lives in the target checkpoint's mtp.* namespace.
+        draft_arches = getattr(draft_config.pretrained_config, "architectures",
+                               None) or []
+        # The drafter's own ModelConfig carries spec_config=None,
+        # so the validated speculative-config values are passed in explicitly here.
+        if any("Qwen3DSpark" in arch for arch in draft_arches):
+            return Qwen3DSparkForCausalLM(
+                draft_config,
+                block_size=model_config.spec_config.block_size,
+                mask_token_id=model_config.spec_config.mask_token_id,
+                ctx_window_size=model_config.spec_config.ctx_window_size,
+                serving_max_seq_len=model_config.max_seq_len)
         num_stages = count_dspark_stages(
             model_config.spec_config.speculative_model)
         validate_dspark_eplb_layer_base(model_config, draft_config)
@@ -1925,6 +1938,7 @@ def get_draft_model(model_config, draft_config, lm_head, model):
             getattr(model, "aux_stream_dict", None),
             num_stages=num_stages,
             block_size=model_config.spec_config.block_size,
+            mask_token_id=model_config.spec_config.mask_token_id,
         )
     elif spec_dec_mode.is_draft_target_one_model():
         # Keep the draft LM head vocab-sharded so greedy draft sampling uses the

tensorrt_llm/llmapi/llm_args.py

@@ -2837,6 +2837,15 @@ class DSparkDecodingConfig(DecodingBaseConfig):
         "read from the draft model config (dspark_markov_head_type), "
         "defaulting to \"vanilla\".")
 
+    ctx_window_size: Optional[PositiveInt] = Field(
+        default=None,
+        description=
+        "Ring-window length, in tokens, of the worker-owned per-layer context "
+        "K/V buffer used by the dense (Qwen3-style) DSpark drafter. Clamped to "
+        "[block_size + 2, max_position_embeddings]. Only applies to the dense "
+        "Qwen3 DSpark drafter (ignored by the DeepSeek-V4 drafter, which uses "
+        "the checkpoint's own sliding_window instead). Defaults to 2048.")
+
     # NOTE: confidence-based dynamic drafting (the draft model's confidence head
     # that truncates the proposed block) is NOT enabled in this PR. The user-facing
     # ``enable_confidence_head`` / ``confidence_threshold`` knobs are intentionally
@@ -5726,11 +5735,23 @@ def validate_speculative_config(self):
                     with open(draft_config_path) as f:
                         draft_cfg = json.load(f)
                     dspark_cfg = draft_cfg.get("dspark_config", {})
+                    # DeepSpec-released dense drafter checkpoints (e.g.
+                    # Qwen3DSparkModel) use unprefixed top-level keys in their
+                    # own config.json; gate that fallback on the checkpoint's
+                    # architectures (mirroring the dispatch in
+                    # get_draft_model) so a V4-style checkpoint that happens
+                    # to carry an unrelated top-level key with the same name
+                    # (e.g. `block_size`) doesn't get silently misread.
+                    draft_arches = draft_cfg.get("architectures") or []
+                    is_dense_drafter = any("Qwen3DSpark" in arch
+                                           for arch in draft_arches)
 
                     def _dspark_get(key, top_level_key):
                         value = dspark_cfg.get(key)
                         if value is None:
                             value = draft_cfg.get(top_level_key)
+                        if value is None and is_dense_drafter:
+                            value = draft_cfg.get(key)
                         return value
 
                     # The checkpoint's ``dspark_target_layer_ids`` is

tensorrt_llm/usage/llm_args_golden_manifest.json

@@ -1731,6 +1731,13 @@
       "kind": "value",
       "path": "speculative_config.block_size"
     },
+    {
+      "allowed_values": [],
+      "annotation": "Optional[Annotated[int, Gt(gt=0)]]",
+      "converter": "",
+      "kind": "value",
+      "path": "speculative_config.ctx_window_size"
+    },
     {
       "allowed_values": [
         "AUTO",

tests/integration/defs/accuracy/references/gsm8k.yaml

@@ -174,6 +174,8 @@ Qwen3/Qwen3-8B:
     accuracy: 87.1114
   - spec_dec_algo: DFlash
     accuracy: 87.1114
+  - spec_dec_algo: DSpark
+    accuracy: 87.1114
   - quant_algo: FP8
     kv_cache_quant_algo: FP8
     accuracy: 87.1114

tests/integration/defs/accuracy/test_llm_api_pytorch.py

@@ -4637,6 +4637,30 @@ def test_dflash(self):
             task = GSM8K(self.MODEL_NAME)
             task.evaluate(llm)
 
+    @skip_pre_hopper
+    def test_dspark(self):
+        pytorch_config = dict(
+            max_batch_size=8,
+            disable_overlap_scheduler=False,
+            cuda_graph_config=CudaGraphConfig(max_batch_size=8,
+                                              enable_padding=True),
+        )
+        kv_cache_config = KvCacheConfig(enable_block_reuse=False,
+                                        free_gpu_memory_fraction=0.6)
+
+        dspark_model_dir = f"{llm_models_root()}/dspark_qwen3_8b_block7"
+        target_model_dir = f"{llm_models_root()}/Qwen3/Qwen3-8B"
+
+        spec_config = DSparkDecodingConfig(max_draft_len=7,
+                                           speculative_model=dspark_model_dir)
+
+        with LLM(model=target_model_dir,
+                 **pytorch_config,
+                 kv_cache_config=kv_cache_config,
+                 speculative_config=spec_config) as llm:
+            task = GSM8K(self.MODEL_NAME)
+            task.evaluate(llm)
+
     @skip_pre_blackwell
     @pytest.mark.parametrize("tp_size,pp_size,ep_size,attention_dp",
                              [(1, 1, 1, False)],

tests/integration/test_lists/test-db/l0_h100.yml

@@ -157,6 +157,7 @@ l0_h100:
   - accuracy/test_llm_api_pytorch.py::TestQwen3_30B_A3B::test_dummy_load_format
   - accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_eagle3[eagle3_one_model=True-enable_chunked_prefill=False-enable_max_concurrency=False-enable_draft_len_schedule=False]
   - accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_dflash
+  - accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_dspark
   - accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_dflash
   - accuracy/test_llm_api_pytorch.py::TestQwen3_5_4B::test_bf16
   - accuracy/test_llm_api_pytorch.py::TestQwen3_5_4B::test_fp8

tests/unittest/_torch/speculative/hw_agnostic/test_dspark_attention.py

@@ -86,7 +86,9 @@ def fake_decoder_layer_init(
         fake_decoder_layer_init,
     )
     model_config = types.SimpleNamespace(
-        pretrained_config=types.SimpleNamespace(vocab_size=128, hc_mult=2),
+        pretrained_config=types.SimpleNamespace(
+            vocab_size=128, hc_mult=2, dspark_noise_token_id=127
+        ),
         spec_config=None,
     )