[WIP] feat: add DFlash2 draft serving support

docs/configuration/server.md

@@ -194,6 +194,16 @@ startup rather than silently drafting a wrong-width block. A checkpoint with
 `block_size: 8` therefore wants `--speculative-num-draft-tokens 8
 --speculative-num-steps 7`.
 
+A checkpoint whose architecture is `DFlash2DraftModel` uses the same `DFLASH`
+launch method. TokenSpeed selects its grouped-convolution and candidate-selector
+runtime from the checkpoint architecture; no separate algorithm flag is needed.
+Greedy requests use the selector's best path. Stochastic requests cache the
+realized top-K selector distribution per request and use lossless `p/q`
+rejection plus `relu(p-q)` residual sampling. Select `--sampling-backend
+flashinfer` or `flashinfer_full` for that path; the Triton sampling backends
+currently fail fast for DFlash2 stochastic verification instead of silently
+falling back to target-only acceptance.
+
 A block drafter writes its KV at the target's cache locations, so it shares the
 target's page table: `--block-size` is a target-side choice and the draft
 follows it. Any sliding window the draft checkpoint declares is an attention

docs/recipes/models.md

@@ -210,6 +210,14 @@ attention. It does not yet expose an equivalent of SGLang's
 `--speculative-dflash-draft-window-size`; add such a flag before relying on
 bounded draft attention for long-context deployments.
 
+Official DFlash2 checkpoints that declare `DFlash2DraftModel` use the same
+`--speculative-algorithm DFLASH` launch. Their grouped dynamic convolutions and
+candidate selector are enabled automatically from the draft architecture. Add
+`--sampling-backend flashinfer` (or `flashinfer_full`) for nonzero-temperature
+sampling so the selector's sparse q distribution participates in lossless
+rejection verification. `greedy` remains the lower-overhead temperature-zero
+path.
+
 ## Kimi K3
 
 Kimi-K3 combines a MoonViT vision encoder with a hybrid KDA

python/tokenspeed/runtime/configs/model_config.py

@@ -284,6 +284,21 @@ def _resolve_attention_family(
     return None
 
 
+def _is_dflash2_mla(
+    hf_config: PretrainedConfig,
+    hf_text_config: PretrainedConfig,
+) -> bool:
+    architectures = _model_architectures(hf_config, hf_text_config)
+    dflash_config = getattr(hf_text_config, "dflash_config", None) or getattr(
+        hf_config, "dflash_config", None
+    )
+    return (
+        "DFlash2DraftModel" in architectures
+        and isinstance(dflash_config, dict)
+        and dflash_config.get("attention_mode") == "mla"
+    )
+
+
 def _apply_attention_family_defaults(
     server_args: ServerArgs,
     spec: _AttentionFamilySpec,
@@ -541,6 +556,8 @@ def __init__(
         if attention_family is not None:
             _apply_attention_family_defaults(server_args, attention_family)
             attention_family.configure(self)
+        elif _is_dflash2_mla(self.hf_config, self.hf_text_config):
+            configure_mla_attention(self)
         elif "MiniCPM3ForCausalLM" in self.hf_config.architectures:
             self.head_dim = 128
             self.attention_arch = AttentionArch.MLA

python/tokenspeed/runtime/execution/drafter/__init__.py

@@ -46,6 +46,7 @@ def get_drafter_impl(spec_algo: str, model: torch.nn.Module) -> type[BaseDrafter
     # Imports are local: drafter modules pull in kernel ops and model code,
     # and this package init must stay importable from lightweight contexts.
     from tokenspeed.runtime.execution.drafter.dflash import DFlash
+    from tokenspeed.runtime.execution.drafter.dflash2 import DFlash2
     from tokenspeed.runtime.execution.drafter.dspark import DSpark
     from tokenspeed.runtime.execution.drafter.eagle import Eagle
     from tokenspeed.runtime.models.inkling_nextn import (
@@ -59,6 +60,12 @@ def get_drafter_impl(spec_algo: str, model: torch.nn.Module) -> type[BaseDrafter
         "DSPARK": DSpark,
     }
 
+    if spec_algo == "DFLASH":
+        from tokenspeed.runtime.models.dflash2 import DFlash2DraftModel
+
+        if isinstance(model, DFlash2DraftModel):
+            return DFlash2
+
     # "MTP" covers two algorithms:
     # (1) Eagle-like MTP (e.g. DeepSeek) stays on Eagle in eagle.py;
     # (2) Vanilla MTP (e.g. Inkling) with multi-layer weights stays on Mtp in mtp.py.

python/tokenspeed/runtime/execution/drafter/dflash2.py

@@ -0,0 +1,244 @@
+# Copyright (c) 2026 LightSeek Foundation
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+"""DFlash2 proposal path on top of TokenSpeed's native DFlash block runtime."""
+
+from __future__ import annotations
+
+import torch
+
+from tokenspeed.runtime.execution.drafter.dflash import DFlash
+from tokenspeed.runtime.execution.forward_batch_info import ForwardMode
+from tokenspeed.runtime.layers.logits_processor import LogitsMetadata, LogitsProcessor
+from tokenspeed.runtime.sampling.draft_distribution import SparseDraftDistribution
+from tokenspeed.runtime.utils.nvtx import nvtx_range
+
+
+def _walk_best_path(
+    candidate_ids: torch.Tensor,
+    scores: torch.Tensor,
+    anchor_token_ids: torch.Tensor,
+    out: torch.Tensor,
+) -> torch.Tensor:
+    """Walk a fixed DFlash2 candidate lattice without host-side tensor reads."""
+    batch_size, num_steps, top_k = candidate_ids.shape
+    out[:, 0].copy_(anchor_token_ids)
+    previous = torch.zeros(batch_size, dtype=torch.int64, device=candidate_ids.device)
+    for step in range(num_steps):
+        transitions = torch.gather(
+            scores[:, step],
+            1,
+            previous[:, None, None].expand(-1, 1, top_k),
+        ).squeeze(1)
+        previous = torch.argmax(transitions, dim=-1)
+        token = torch.gather(candidate_ids[:, step], 1, previous[:, None]).squeeze(1)
+        out[:, step + 1].copy_(token)
+    return out
+
+
+def _walk_sampled_path(
+    candidate_ids: torch.Tensor,
+    scores: torch.Tensor,
+    anchor_token_ids: torch.Tensor,
+    temperatures: torch.Tensor,
+    coins: torch.Tensor,
+    out: torch.Tensor,
+) -> torch.Tensor:
+    """Sample the selector path and return each realized sparse q row."""
+    batch_size, num_steps, top_k = candidate_ids.shape
+    out[:, 0].copy_(anchor_token_ids)
+    previous = torch.zeros(batch_size, dtype=torch.int64, device=candidate_ids.device)
+    realized_probs = torch.empty(
+        (batch_size, num_steps, top_k),
+        dtype=torch.float32,
+        device=candidate_ids.device,
+    )
+    stochastic = temperatures > 0
+    safe_temperatures = torch.where(
+        stochastic, temperatures, torch.ones_like(temperatures)
+    )
+    for step in range(num_steps):
+        transitions = (
+            torch.gather(
+                scores[:, step],
+                1,
+                previous[:, None, None].expand(-1, 1, top_k),
+            )
+            .squeeze(1)
+            .float()
+        )
+        probabilities = torch.softmax(transitions / safe_temperatures[:, None], dim=-1)
+        realized_probs[:, step].copy_(probabilities)
+        sampled = (probabilities.cumsum(dim=-1) < coins[:, step, None]).sum(dim=-1)
+        sampled = sampled.clamp(max=top_k - 1)
+        greedy = torch.argmax(transitions, dim=-1)
+        previous = torch.where(stochastic, sampled, greedy)
+        token = torch.gather(candidate_ids[:, step], 1, previous[:, None]).squeeze(1)
+        out[:, step + 1].copy_(token)
+    return realized_probs
+
+
+class DFlash2(DFlash):
+    """DFlash block runtime with the DFlash2 top-k transition selector."""
+
+    def __init__(self, *args, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+        self.candidate_selector = getattr(self.model, "candidate_selector", None)
+        if self.candidate_selector is None:
+            raise ValueError(
+                "DFlash2 requires a draft model with candidate_selector weights."
+            )
+        if self.draft_query_width != self.spec_num_tokens:
+            raise ValueError("DFlash2 requires the anchor-plus-mask DFlash layout.")
+
+        config = self.model.config
+        nested = getattr(config, "dflash_config", {}) or {}
+        self.selector_top_k = int(nested.get("selector_top_k"))
+        self.output_multiplier = float(nested.get("output_multiplier", 1.0))
+        self.final_logit_softcapping = float(
+            nested.get("final_logit_softcapping") or 0.0
+        )
+        self.candidate_logits_processor: LogitsProcessor | None = None
+        self.sampling_backend = None
+        self.dflash2_verify_mode = "unwired"
+        self._draft_candidate_ids_pool: torch.Tensor | None = None
+        self._draft_candidate_probs_pool: torch.Tensor | None = None
+
+    def wire_target(self, target_model) -> None:
+        super().wire_target(target_model)
+        self.candidate_logits_processor = LogitsProcessor(
+            self.model.config,
+            logit_scale=self.output_multiplier,
+            tp_rank=self.logits_processor.tp_rank,
+            tp_size=self.logits_processor.tp_size,
+            tp_group=self.logits_processor.tp_group,
+        )
+        self.candidate_logits_processor.final_logit_softcapping = (
+            self.final_logit_softcapping if self.final_logit_softcapping > 0 else None
+        )
+
+    def wire_sampling_backend(self, sampling_backend) -> None:
+        mode = str(getattr(sampling_backend, "dflash2_verify_mode", "unsupported"))
+        if mode not in ("greedy", "rejection"):
+            raise ValueError(
+                f"{type(sampling_backend).__name__} cannot verify DFlash2 q; use "
+                "--sampling-backend greedy, flashinfer, or flashinfer_full"
+            )
+        self.sampling_backend = sampling_backend
+        self.dflash2_verify_mode = mode
+        if mode == "rejection":
+            pool_rows = int(self.runtime_states.valid_cache_lengths.shape[0])
+            shape = (pool_rows, self.spec_num_tokens - 1, self.selector_top_k)
+            self._draft_candidate_ids_pool = torch.zeros(
+                shape, dtype=torch.int64, device=self.device
+            )
+            self._draft_candidate_probs_pool = torch.zeros(
+                shape, dtype=torch.float32, device=self.device
+            )
+
+    def get_draft_distribution(self, base_ctx) -> SparseDraftDistribution | None:
+        if (
+            self.dflash2_verify_mode != "rejection"
+            or base_ctx.bs == base_ctx.num_extends
+        ):
+            return None
+        if (
+            self._draft_candidate_ids_pool is None
+            or self._draft_candidate_probs_pool is None
+        ):
+            raise RuntimeError("DFlash2 rejection buffers were not initialized")
+        req_pool_indices = self.input_buffers.req_pool_indices_buf[
+            base_ctx.num_extends : base_ctx.bs
+        ]
+        return SparseDraftDistribution(
+            candidate_ids=self._draft_candidate_ids_pool.index_select(
+                0, req_pool_indices
+            ),
+            probabilities=self._draft_candidate_probs_pool.index_select(
+                0, req_pool_indices
+            ),
+        )
+
+    def _compute_candidates(
+        self, hidden_states: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if self.candidate_logits_processor is None:
+            raise RuntimeError("DFlash2 must be wired to the target before drafting.")
+        metadata = LogitsMetadata(forward_mode=ForwardMode.DECODE)
+        logits = self.candidate_logits_processor._get_logits(
+            hidden_states, self.lm_head, metadata
+        )
+        unary_logits, candidate_ids = torch.topk(
+            logits, self.selector_top_k, dim=-1, sorted=False
+        )
+        return candidate_ids, unary_logits
+
+    @nvtx_range("dflash2_sample_block", color="purple")
+    def _sample_block(
+        self,
+        draft_hidden: torch.Tensor,
+        block_ids: torch.Tensor,
+        next_tokens: torch.Tensor,
+    ) -> torch.Tensor:
+        hidden_states = draft_hidden[:, 1:, :]
+        batch_size, num_steps, _ = hidden_states.shape
+        candidate_ids, unary_logits = self._compute_candidates(
+            hidden_states.reshape(-1, self.hidden_size)
+        )
+        candidate_ids = candidate_ids.view(batch_size, num_steps, self.selector_top_k)
+        unary_logits = unary_logits.view_as(candidate_ids)
+        anchor_token_ids = (
+            block_ids[:, 0]
+            .to(torch.int64)
+            .clamp(0, int(self.model.config.vocab_size) - 1)
+        )
+        scores = self.candidate_selector(
+            candidate_ids,
+            unary_logits,
+            hidden_states,
+            anchor_token_ids,
+        )
+        if self.dflash2_verify_mode == "rejection":
+            if self.sampling_backend is None:
+                raise RuntimeError("DFlash2 sampling backend is not wired")
+            temperatures, coins = self.sampling_backend.dflash2_proposal_state(
+                self.input_buffers.req_pool_indices_buf[:batch_size],
+                batch_size,
+                num_steps,
+            )
+            realized_probs = _walk_sampled_path(
+                candidate_ids,
+                scores,
+                anchor_token_ids,
+                temperatures,
+                coins,
+                next_tokens,
+            )
+            req_pool_indices = self.input_buffers.req_pool_indices_buf[:batch_size]
+            self._draft_candidate_ids_pool.index_copy_(
+                0, req_pool_indices, candidate_ids
+            )
+            self._draft_candidate_probs_pool.index_copy_(
+                0, req_pool_indices, realized_probs
+            )
+        else:
+            _walk_best_path(candidate_ids, scores, anchor_token_ids, next_tokens)
+        next_tokens.clamp_(min=0, max=int(self.model.config.vocab_size) - 1)
+        return next_tokens

python/tokenspeed/runtime/execution/model_executor.py

@@ -438,6 +438,9 @@ def __init__(
                 vocab_size=config.vocab_size,
             )
             self.drafter.wire_target(self.model_runner.model)
+            wire_sampling_backend = getattr(self.drafter, "wire_sampling_backend", None)
+            if wire_sampling_backend is not None:
+                wire_sampling_backend(self.sampling_backend)
             MultimodalRuntime.wire_drafter(
                 self.input_buffers, self.model_runner.model_config
             )
@@ -722,6 +725,7 @@ def _run_sampling(
         sampling_info: SamplingBatchInfo,
         ctx: ForwardContext,
         candidates: torch.Tensor | None = None,
+        draft_distribution=None,
     ):
         if self.drafter is None:
             return self.sampling_backend.sample(logits_output, sampling_info)
@@ -734,7 +738,7 @@ def _run_sampling(
 
         if num_extends == 0:
             output_tokens, accept_lengths = self.sampling_backend.verify(
-                logits_output, sampling_info, candidates
+                logits_output, sampling_info, candidates, draft_distribution
             )
             accept_lengths = self._apply_force_single_token_verify(
                 accept_lengths, 0, num_decodes, ctx.decode_input_ids
@@ -748,7 +752,10 @@ def _run_sampling(
         )
         decode_out = LogitsProcessorOutput(next_token_logits=logits[num_extends:])
         decode_tokens, decode_accept = self.sampling_backend.verify(
-            decode_out, sampling_info[num_extends:], candidates
+            decode_out,
+            sampling_info[num_extends:],
+            candidates,
+            draft_distribution,
         )
         decode_accept = self._apply_force_single_token_verify(
             decode_accept, num_extends, num_decodes, ctx.decode_input_ids
@@ -851,12 +858,22 @@ def _forward_step(
             if self.config.spec_algo is not None
             else None
         )
+        draft_distribution = (
+            self.drafter.get_draft_distribution(ctx)
+            if self.drafter is not None
+            and hasattr(self.drafter, "get_draft_distribution")
+            else None
+        )
 
         if self.capturable_grammar is not None:
             self.capturable_grammar.wait_bitmask()
 
         output_tokens, accept_lengths = self._run_sampling(
-            logits_output, sampling_info, ctx, candidates
+            logits_output,
+            sampling_info,
+            ctx,
+            candidates,
+            draft_distribution,
         )
 
         # Backstop: flag any request whose sampled id falls outside [0, vocab)