[PD] Transfer DFlash draft KV separately from DCP target indices

docs/cookbook/autoregressive/Moonshotai/Kimi-K3.mdx

@@ -70,7 +70,7 @@ Pick your hardware, then the deployment shape and operating point. Node count fo
 
 `Long-Context` appears only under the `Prefill` PD mode; for long-context unified serving on B200, start from High-Throughput and raise `--context-length`.
 
-**Spec Decode** — layers onto the strategy without changing it, on every platform except B200. DSPARK proposes 7 draft tokens per step (tune in the Playground) and requires `pp_size == 1`, so on B200 it also drops the pipeline and re-lays the same 16 GPUs flat: PP2 × TP8 → TP16, PP2 × DCPEP8 → DCPEP16. DFLASH has no published draft checkpoint. The win is largest on short interactive traffic and fades as the prompt grows.
+**Spec Decode** — layers onto the strategy without changing it, on every platform except B200. DSPARK proposes 7 draft tokens per step (tune in the Playground) and requires `pp_size == 1`, so on B200 it also drops the pipeline and re-lays the same 16 GPUs flat: PP2 × TP8 → TP16, PP2 × DCPEP8 → DCPEP16. DFLASH uses the [Kimi-K3-DFlash](https://huggingface.co/modal-labs/Kimi-K3-DFlash) draft at block size 8; it also requires `pp_size == 1` and no DP attention, but does not re-lay pipelined recipes, so it is unavailable on the pipelined B200 Unified recipes. It has only been validated on Blackwell, so the Hopper and AMD recipes don't offer it yet. The win is largest on short interactive traffic and fades as the prompt grows.
 
 <Note>
 `--mamba-full-memory-ratio` is the one sizing flag, computed live: set your average request length in the [Mamba ratio calculator](#mamba-ratio-calculator); everything else follows the panels, and the result is pinned into the command. (The Ascend NPU recipes don't use the calculator and set no ratio at all — the Ascend path sizes both pools itself: the A3 Series pins `--max-mamba-cache-size` instead, and the A5 recipe serves with the radix cache off.)

docs/src/snippets/configs/moonshotai/kimi-k3.jsx

@@ -334,16 +334,28 @@ export const config = {
         {
           id: "dflash",
           label: "DFLASH",
-          // Listed so the axis is complete, but not selectable: no K3 DFLASH draft
-          // checkpoint has been published, so there is nothing to point
-          // --speculative-draft-model-path at. DFLASH is also CUDA-only, rejects DP
-          // attention, and requires pp_size == 1.
-          disabled: true,
-          disableReason:
-            "No K3 DFLASH draft checkpoint published yet — DSPARK is the available speculative path.",
-          flags: [
+          // DFLASH requires pp_size == 1 and rejects DP attention off NPU, and it
+          // does not re-lay a pipelined cell the way DSPARK does, so those
+          // recipes are unavailable. The NPU recipes ship DSPARK only. K3 DFLASH
+          // has only been validated on Blackwell, so Hopper and AMD stay off.
+          disabled: (s) =>
+            config.isNpuHw(s) ||
+            ["h100", "h200", "mi350x", "mi355x"].includes(s.hw) ||
+            config.isPipelined(s) ||
+            !!config.flagOf(config.cellFor(s), "--enable-dp-attention"),
+          disableReason: (s) =>
+            config.isNpuHw(s)
+              ? "Only DSPARK is supported on this recipe."
+              : ["h100", "h200", "mi350x", "mi355x"].includes(s.hw)
+                ? "K3 DFLASH has not been validated on this hardware yet; use DSPARK."
+                : "DFLASH requires pp_size == 1 and no DP attention. Pick a recipe that runs a single pipeline stage without DP attention, or use DSPARK.",
+          flags: (s) => [
             "--speculative-algorithm DFLASH",
-            "--speculative-draft-model-path <dflash-draft>",
+            "--speculative-draft-model-path modal-labs/Kimi-K3-DFlash",
+            // 8 is the recommended default block size for this draft.
+            "--speculative-dflash-block-size 8",
+            // Same ReplaySSM rule as DSPARK: the PD prefill role rejects the flag.
+            ...(s.pdMode !== "prefill" ? ["--enable-linear-replayssm-spec"] : []),
           ],
         },
       ],
@@ -785,7 +797,7 @@ export const config = {
         //   DSPARK  --speculative-dspark-block-size N   (gamma, == proposed)
         //   DFLASH  --speculative-dflash-block-size N+1 (verify window)
         //   EAGLE   --speculative-num-steps N           (chain; topk>1 is a tree)
-        // Only DSPARK is selectable today, so only its form is emitted.
+        // Only DSPARK's form is emitted; the DFLASH option pins its block size.
         id: "proposedDraftTokens", title: "Proposed Draft Tokens",
         // The NPU recipes pin the shipped block size (7).
         showWhen: (b) => b.spec === "dspark" && !config.isNpuHw(b),

python/sglang/srt/disaggregation/base/conn.py

@@ -16,6 +16,7 @@
 
 class StateType(str, enum.Enum):
     MAMBA = "mamba"
+    DFLASH_KV = "dflash_kv"
     QSA_PENDING = "qsa_pending"
     QSA_COMPRESSED = "qsa_compressed"
     SWA = "swa"

python/sglang/srt/disaggregation/decode.py

@@ -51,6 +51,10 @@
     HiCacheRestoreGatedKVReceiver,
     HiCacheRestoreResult,
 )
+from sglang.srt.disaggregation.dflash_kv import (
+    lists_draft_as_kv_entries,
+    resolve_dflash_draft_transfer,
+)
 from sglang.srt.disaggregation.utils import (
     DisaggregationMode,
     KVClassType,
@@ -398,6 +402,11 @@ def __init__(
         self.pp_size = get_parallel().pp_size
         self.num_reserved_decode_tokens = num_reserved_decode_tokens
         self.transfer_backend = transfer_backend
+        self.dflash_draft_transfer = resolve_dflash_draft_transfer(
+            allocator=token_to_kv_pool_allocator,
+            draft_worker=scheduler.draft_worker,
+            transfer_backend=transfer_backend,
+        )
         # Queue for requests pending pre-allocation
         self.queue: List[DecodeRequest] = []
         self.retracted_queue: List[Req] = []
@@ -596,7 +605,15 @@ def _init_kv_manager(self) -> CommonKVManager:
             kv_item_lens += device_kv_item_lens[c4_layer_num:]
             kv_data_mem_kinds += ["VRAM"] * len(device_kv_data_ptrs[c4_layer_num:])
         num_draft_entries = 0
-        if self.draft_token_to_kv_pool is not None:
+        draft_kv_pool = (
+            self.draft_token_to_kv_pool
+            if lists_draft_as_kv_entries(
+                mode=DisaggregationMode.DECODE,
+                dflash_draft_transfer=self.dflash_draft_transfer,
+            )
+            else None
+        )
+        if draft_kv_pool is not None:
             # Draft KV shares target virtual ids. Unified target KV is transferred
             # with physical ids, so it needs a separate draft index vector.
             draft_kv_data_ptrs, draft_kv_data_lens, draft_kv_item_lens = (
@@ -614,7 +631,7 @@ def _init_kv_manager(self) -> CommonKVManager:
         kv_args.num_draft_entries = num_draft_entries
         kv_args.kv_layer_ids = build_kv_layer_ids(
             token_to_kv_pool=self.token_to_kv_pool,
-            draft_token_to_kv_pool=self.draft_token_to_kv_pool,
+            draft_token_to_kv_pool=draft_kv_pool,
             num_draft_entries=num_draft_entries,
             num_hidden_layers=self.scheduler.model_config.num_hidden_layers,
         )
@@ -632,6 +649,7 @@ def _init_kv_manager(self) -> CommonKVManager:
             self.draft_token_to_kv_pool,
             total_kv_layers=self.scheduler.model_config.num_hidden_layers,
             req_to_token_pool=getattr(self, "req_to_token_pool", None),
+            dflash_draft_transfer=self.dflash_draft_transfer,
         )
 
         kv_args.ib_device = get_disagg().disaggregation_ib_device
@@ -1504,6 +1522,16 @@ def pop_preallocated(
 
             seq_len = origin_input_len
 
+            def _draft_payload():
+                return self.dflash_draft_transfer.page_indices(
+                    req_to_token=self.req_to_token_pool.req_to_token[
+                        decode_req.req.kv.req_pool_idx
+                    ],
+                    prefix_len=total_prefix_len,
+                    seq_len=seq_len,
+                    wire_page_size=self.token_to_kv_pool.page_size,
+                )
+
             def _mamba_payload():
                 return [
                     self.req_to_token_pool.translate_mamba_indices(
@@ -1574,6 +1602,7 @@ def _request_state_payload():
                     clear_request_state(int(decode_req.req.kv.req_pool_idx))
             payloads = {
                 StateType.MAMBA: _mamba_payload,
+                StateType.DFLASH_KV: _draft_payload,
                 StateType.QSA_PENDING: _qsa_pending_payload,
                 StateType.QSA_COMPRESSED: _full_kv_pages_payload,
                 StateType.SWA: _swa_payload,

python/sglang/srt/disaggregation/dflash_kv.py

@@ -0,0 +1,105 @@
+"""Full DFlash draft KV transfer in the draft pool's logical token-id domain."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, List, Optional, Tuple
+
+import msgspec
+import numpy as np
+import torch
+
+from sglang.srt.disaggregation.utils import DisaggregationMode, TransferBackend
+from sglang.srt.mem_cache.common import kv_to_page_indices
+from sglang.srt.runtime_context import get_parallel
+from sglang.srt.speculative.dflash_utils import (
+    get_dflash_attention_sliding_window_size,
+    get_dflash_layer_types,
+)
+
+if TYPE_CHECKING:
+    from sglang.srt.mem_cache.allocator.base import BaseTokenToKVPoolAllocator
+    from sglang.srt.mem_cache.memory_pool import MHATokenToKVPool
+    from sglang.srt.speculative.dflash_worker_v2 import DFlashWorkerV2
+
+
+class DFlashDraftTransfer(msgspec.Struct, frozen=True):
+    """Full DFlash draft KV sent as its own PD state component.
+
+    Under DCP the target KV travels as rank-local rows, while the full draft
+    pool is replicated and indexed by the target allocator's logical token ids.
+    """
+
+    pool: MHATokenToKVPool
+    # Draft tokens decode attends to; None transfers the whole prompt.
+    window: Optional[int]
+
+    def buffer_infos(self, wire_page_size: int) -> Tuple[List, List, List]:
+        """Draft buffers in target wire-page units; draft allocation pages are wider."""
+        if wire_page_size <= 0 or self.pool.page_size % wire_page_size:
+            raise ValueError("Draft allocation pages must contain whole transfer pages")
+        ptrs, lengths, item_lengths = self.pool.get_contiguous_buf_infos()
+        ratio = self.pool.page_size // wire_page_size
+        if any(length % ratio for length in item_lengths):
+            raise ValueError("Draft buffer page bytes must divide into transfer pages")
+        return ptrs, lengths, [length // ratio for length in item_lengths]
+
+    def page_indices(
+        self,
+        *,
+        req_to_token: torch.Tensor,
+        prefix_len: int,
+        seq_len: int,
+        wire_page_size: int,
+    ) -> np.ndarray:
+        """Wire pages of one request's draft KV past the decode prefix."""
+        # The decode prefix already holds its draft KV, like the SWA payload.
+        start = prefix_len
+        if self.window is not None:
+            start = max(prefix_len, seq_len - self.window)
+        start = start // wire_page_size * wire_page_size
+        return kv_to_page_indices(req_to_token[start:seq_len], wire_page_size)
+
+
+def resolve_dflash_draft_transfer(
+    *,
+    allocator: BaseTokenToKVPoolAllocator,
+    draft_worker: Optional[DFlashWorkerV2],
+    transfer_backend: TransferBackend,
+) -> Optional[DFlashDraftTransfer]:
+    # Without DCP the target's KV indices are logical too, so the draft rides
+    # the target KV entries (num_draft_entries) instead.
+    if (
+        allocator.full_draft_kv_pool is None
+        or transfer_backend != TransferBackend.NIXL
+        or get_parallel().attn_dcp_size == 1
+    ):
+        return None
+    return DFlashDraftTransfer(
+        pool=allocator.full_draft_kv_pool,
+        window=draft_transfer_window(
+            draft_worker.draft_model_runner.model_config.hf_config
+        ),
+    )
+
+
+def lists_draft_as_kv_entries(
+    *,
+    mode: DisaggregationMode,
+    dflash_draft_transfer: Optional[DFlashDraftTransfer],
+) -> bool:
+    """Whether this side registers the draft pool among its KV entries.
+
+    Neither side knows the peer's DCP size when it registers, so decode keeps
+    the entries for a DCP=1 prefill's relayout path and also registers
+    DFLASH_KV for a same-DCP prefill, which sends only the latter.
+    """
+    return mode == DisaggregationMode.DECODE or dflash_draft_transfer is None
+
+
+def draft_transfer_window(draft_hf_config) -> Optional[int]:
+    """Draft KV tokens decode needs, or None unless every draft layer slides."""
+    layer_types = get_dflash_layer_types(draft_hf_config)
+    if not layer_types or any(t != "sliding_attention" for t in layer_types):
+        return None
+    # The helper returns window_left; the HF window also counts the current token.
+    return get_dflash_attention_sliding_window_size(draft_hf_config) + 1

python/sglang/srt/disaggregation/nixl/conn.py

@@ -1579,6 +1579,9 @@ def register_buffer_to_engine(self):
         if not self.aux_descs:
             raise Exception("NIXL memory registration failed for aux tensors")
 
+        # Decode lists DFlash draft buffers both as KV entries and as DFLASH_KV
+        # state; register each buffer once.
+        kv_regions = set(zip(self.kv_args.kv_data_ptrs, self.kv_args.kv_data_lens))
         state_addrs = []
         for comp_ptrs, comp_lens in zip(
             self.kv_args.state_data_ptrs or [],
@@ -1587,6 +1590,8 @@ def register_buffer_to_engine(self):
             for state_data_ptr, state_data_len in zip(comp_ptrs, comp_lens):
                 if state_data_ptr == 0 or state_data_len == 0:
                     continue
+                if (state_data_ptr, state_data_len) in kv_regions:
+                    continue
                 state_addrs.append(
                     (state_data_ptr, state_data_len, self.kv_args.gpu_id, "")
                 )
@@ -1606,6 +1611,16 @@ def _add_remote_peer(self, decode_kv_args: KVArgsRegisterInfo):
         decode_kv_args.requires_dcp_relayout = self.requires_dcp_relayout(
             decode_kv_args.dst_dcp_size, decode_kv_args.dst_dcp_rank
         )
+        # Draft MHA heads shard by attention TP even when the target is MLA.
+        if (
+            StateType.DFLASH_KV in self.kv_args.state_types
+            and decode_kv_args.decode_tp_size != self.attn_tp_size
+        ):
+            raise RuntimeError(
+                "PD transfer of full DFlash draft KV requires equal prefill and "
+                f"decode attention TP sizes, got prefill={self.attn_tp_size}, "
+                f"decode={decode_kv_args.decode_tp_size}."
+            )
         if decode_kv_args.requires_dcp_relayout:
             self._init_dcp_pack_buffers_once(decode_kv_args.dst_dcp_size)
         self.decode_kv_args_table[agent_name] = decode_kv_args
@@ -2660,6 +2675,7 @@ def maybe_send_extra(
                     force_flat=True,
                 )
             elif st in (
+                StateType.DFLASH_KV,
                 StateType.SWA,
                 StateType.QSA_PENDING,
                 StateType.QSA_COMPRESSED,
@@ -2691,7 +2707,12 @@ def maybe_send_extra(
                     dst_gpu_id=dst_gpu_id,
                     notif=comp_notif,
                     state_type=st,
-                    force_flat=st in (StateType.QSA_PENDING, StateType.QSA_COMPRESSED),
+                    force_flat=st
+                    in (
+                        StateType.DFLASH_KV,
+                        StateType.QSA_PENDING,
+                        StateType.QSA_COMPRESSED,
+                    ),
                     src_layer_ids=src_lids,
                     dst_layer_ids=dst_lids,
                     dst_item_lens=dst_lens,

python/sglang/srt/disaggregation/prefill.py

@@ -42,6 +42,10 @@
     compute_grid_segments,
     staging_grid_tokens,
 )
+from sglang.srt.disaggregation.dflash_kv import (
+    lists_draft_as_kv_entries,
+    resolve_dflash_draft_transfer,
+)
 from sglang.srt.disaggregation.utils import (
     FAKE_BOOTSTRAP_HOST,
     DisaggregationMode,
@@ -183,9 +187,14 @@ def __init__(
         self.scheduler = scheduler
         self.scheduler_stage_metrics = scheduler_stage_metrics
         self.max_total_num_tokens = (
-            self.scheduler.tp_worker.model_runner.effective_max_total_num_tokens
+            self.scheduler.tp_worker.model_runner.effective_logical_max_total_num_tokens
         )
         self.transfer_backend = transfer_backend
+        self.dflash_draft_transfer = resolve_dflash_draft_transfer(
+            allocator=self.scheduler.token_to_kv_pool_allocator,
+            draft_worker=self.scheduler.draft_worker,
+            transfer_backend=transfer_backend,
+        )
         if envs.SGLANG_DISAGG_STAGING_BUFFER.get():
             if self.is_mla_backend:
                 raise RuntimeError(
@@ -273,6 +282,14 @@ def _init_kv_manager(self) -> CommonKVManager:
             else None
         )
         num_draft_entries = 0
+        dflash_draft_transfer = (
+            self.dflash_draft_transfer if transfer_draft_cache else None
+        )
+        if not lists_draft_as_kv_entries(
+            mode=DisaggregationMode.PREFILL,
+            dflash_draft_transfer=dflash_draft_transfer,
+        ):
+            draft_kv_pool = None
         if draft_kv_pool is not None:
             # Draft KV shares target virtual ids. Unified target KV is transferred
             # with physical ids, so it needs a separate draft index vector.
@@ -314,6 +331,7 @@ def _init_kv_manager(self) -> CommonKVManager:
             self.draft_token_to_kv_pool if transfer_draft_cache else None,
             self.scheduler.model_config.num_hidden_layers,
             req_to_token_pool=req_to_token_pool,
+            dflash_draft_transfer=dflash_draft_transfer,
         )
 
         kv_manager_class = get_kv_class(self.transfer_backend, KVClassType.MANAGER)
@@ -1352,6 +1370,16 @@ def _send_kv_chunk(
             seq_len = min(req.extend_range.end, transfer_input_len)
             c128_seq_len = transfer_input_len
 
+            def _draft_payload():
+                return self.disagg_prefill_bootstrap_queue.dflash_draft_transfer.page_indices(
+                    req_to_token=self.req_to_token_pool.req_to_token[
+                        req.kv.req_pool_idx
+                    ],
+                    prefix_len=req.disagg_decode_prefix_len,
+                    seq_len=seq_len,
+                    wire_page_size=self.token_to_kv_pool.page_size,
+                )
+
             def _mamba_payload():
                 return [
                     self.req_to_token_pool.translate_mamba_indices(
@@ -1419,6 +1447,7 @@ def _request_state_payload():
             )
             payloads = {
                 StateType.MAMBA: _mamba_payload,
+                StateType.DFLASH_KV: _draft_payload,
                 StateType.QSA_PENDING: _qsa_pending_payload,
                 StateType.QSA_COMPRESSED: _full_kv_pages_payload,
                 StateType.SWA: _swa_payload,