feat(m3): support M3 + DSpark (#1190)

.github/workflows/k8s-dispatch.yml

@@ -29,6 +29,7 @@ on:
           - test/ci/eval/kimi-k2.5-nvfp4-eagle3-evalscope-ocr-bench.yaml
           - test/ci/eval/kimi-k2.5-mxfp4-eagle3-evalscope-aime25-amd.yaml
           - test/ci/eval/kimi-k3-mxfp4-tp8ep8-evalscope-aime26-amd.yaml
+          - test/ci/eval/minimax-m3-nvfp4-dspark-evalscope-aime26.yaml
           - test/ci/eval/minimax-m3-nvfp4-evalscope-aime25.yaml
           - test/ci/eval/qwen3.5-122b-a10b-nvfp4-epd-1e1p2d-evalscope-ocr-bench.yaml
           - test/ci/eval/qwen3.5-122b-a10b-nvfp4-evalscope-ocr-bench.yaml

.github/workflows/slurm-dispatch.yml

@@ -40,6 +40,7 @@ on:
           - test/ci/eval/kimi-k3-mxfp4-tp8-two-node-evalscope-aime26-gb300-slurm.yaml
           - test/ci/eval/kimi-k3-nvfp4-dspark-tp8-two-node-evalscope-aime26-gb300-slurm.yaml
           - test/ci/eval/kimi-k3-nvfp4-tp8-two-node-evalscope-aime26-gb300-slurm.yaml
+          - test/ci/eval/minimax-m3-nvfp4-dspark-evalscope-aime26.yaml
           - test/ci/eval/minimax-m3-nvfp4-evalscope-aime25.yaml
           - test/ci/eval/qwen3.5-122b-a10b-nvfp4-epd-1e1p2d-evalscope-ocr-bench.yaml
           - test/ci/eval/qwen3.5-122b-a10b-nvfp4-evalscope-ocr-bench.yaml

docs/configuration/server.md

@@ -174,7 +174,7 @@ the values accepted by the bundled `tokenspeed-smg` package.
 | Parameter | Purpose |
 | --- | --- |
 | `--speculative-config` | JSON speculative decoding configuration. |
-| `--speculative-algorithm` | Speculative algorithm, such as `EAGLE3`, `MTP`, or `DFLASH`. |
+| `--speculative-algorithm` | Speculative algorithm, such as `EAGLE3`, `MTP`, `DFLASH`, or `DSPARK`. |
 | `--speculative-draft-model-path` | Draft model path or repo ID. |
 | `--speculative-draft-model-quantization` | Draft model quantization. Defaults to `unquant`. |
 | `--speculative-num-steps` | Number of draft model steps. Defaults to `3`. |
@@ -185,6 +185,21 @@ the values accepted by the bundled `tokenspeed-smg` package.
 Prefer `--speculative-config` for recipe-style launches because it keeps method,
 draft model, and token count together.
 
+`DFLASH` and `DSPARK` are block drafters: one draft forward proposes a whole
+block instead of one token per step, so their two token counts are coupled.
+`--speculative-num-draft-tokens` is the verify width -- one anchor row plus one
+row per drafted token -- and `--speculative-num-steps` must be one less. The
+draft checkpoint's `block_size` fixes both, and a mismatch is rejected at
+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 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
+mask applied by the draft's own layers, never a cache-retention policy of its
+own.
+
 ## Observability
 
 | Parameter | Purpose |

docs/recipes/models.md

@@ -71,6 +71,8 @@ MiniMax M3 uses 128-token MSA blocks. TokenSpeed configures its dense and sparse
 attention layers automatically; select the dense backend with
 `--attention-backend` and run with `--disable-kvstore`.
 
+### EAGLE3 draft
+
 ```bash
 tokenspeed serve nvidia/MiniMax-M3-NVFP4 \
     --tensor-parallel-size 4 \
@@ -95,6 +97,60 @@ tokenspeed serve nvidia/MiniMax-M3-NVFP4 \
     --port 8000
 ```
 
+### DSpark draft
+
+`nvidia/MiniMax-M3-DSpark` is a six-layer Qwen3-shaped GQA block drafter with a
+vanilla Markov head. Keep the target launch shape and swap the speculative
+options:
+
+```bash
+tokenspeed serve nvidia/MiniMax-M3-NVFP4 \
+    --tensor-parallel-size 4 \
+    --max-model-len 262144 \
+    --max-num-seqs 16 \
+    --max-prefill-tokens 8192 \
+    --chunked-prefill-size 8192 \
+    --gpu-memory-utilization 0.95 \
+    --disable-cuda-graph-padding \
+    --attention-backend trtllm \
+    --kv-cache-dtype fp8 \
+    --moe-backend flashinfer_trtllm \
+    --speculative-algorithm DSPARK \
+    --speculative-draft-model-path nvidia/MiniMax-M3-DSpark \
+    --speculative-num-steps 7 \
+    --speculative-eagle-topk 1 \
+    --speculative-num-draft-tokens 8 \
+    --disable-kvstore \
+    --block-size 128 \
+    --trust-remote-code \
+    --host 0.0.0.0 \
+    --port 8000
+```
+
+Notes:
+
+- Launch this checkpoint with `--speculative-num-draft-tokens 8` and
+  `--speculative-num-steps 7`: the verify window is one anchor row plus seven
+  draft queries, and the step count is always the verify width minus one. The
+  width is checked against the checkpoint's `block_size` at startup, so a
+  mismatched launch fails fast instead of drafting a wrong-width block.
+- `--block-size 128` is the target's MSA page size. The draft writes its KV at
+  the target's cache locations and shares the target's page table, so it
+  inherits that page size; do not set a separate draft block size.
+- The draft's 1024-token sliding window is an attention mask its own layers
+  apply. It is deliberately not a cache-retention policy, because the draft's
+  pages are the target's pages.
+- Target features are captured from the residual stream after each layer in
+  `dflash_config.target_layer_ids` (`[1, 12, 23, 35, 46, 57]` of M3's 60
+  layers) and concatenated in ascending layer order to feed the draft's `fc`.
+- The draft checkpoint stores fp32 master weights. It is loaded in the target's
+  dtype rather than the standalone fp32-to-fp16 default, because the two
+  exchange hidden states and share the target's embedding and LM head.
+- Measured on 4x GB200 with the launch above: gsm8k `mean_acc` 0.9727 versus
+  0.9773 without speculative decoding (paired disagreement 3 vs 9, McNemar
+  p ~ 0.15 -- within run-to-run noise), at a mean accepted length of 4.99 of 8
+  and about 1.7x decode throughput at 16 concurrent requests.
+
 ## Kimi K2.5 / K2.6
 
 Kimi-style MoE launches usually need remote code, long context, reasoning and

python/tokenspeed/runtime/engine/event_loop.py

@@ -1027,15 +1027,20 @@ class EventLoop:
     ) -> ModelConfig:
         server_args = self.server_args
         quantization = server_args.quantization
+        dtype = server_args.dtype
         if is_draft_worker:
             quantization = server_args.speculative_draft_model_quantization
+            if dtype == "auto":
+                # A draft is fed the target's hidden states and borrows its
+                # embedding and LM head, so the two dtypes have to agree.
+                dtype = self.model_config.dtype
         return ModelConfig(
             model_path,
             trust_remote_code=server_args.trust_remote_code,
             revision=server_args.revision,
             context_length=server_args.max_model_len,
             model_override_args=server_args.hf_overrides,
-            dtype=server_args.dtype,
+            dtype=dtype,
             quantization=quantization,
             server_args=server_args,
             is_draft_worker=is_draft_worker,

python/tokenspeed/runtime/execution/cache_loc_kernel.py

@@ -33,7 +33,7 @@ from tokenspeed.runtime.utils import get_colorful_logger
 logger = get_colorful_logger(__name__)
 
 
-@triton.jit
+@triton.jit(do_not_specialize=["max_pages"])
 def compute_out_cache_loc_kernel(
     # Input pointers
     input_lengths_ptr,  # [batch_size] or None for uniform mode
@@ -45,7 +45,7 @@ def compute_out_cache_loc_kernel(
     # Scalars
     uniform_input_length,  # used when input_lengths_ptr is None
     page_size: tl.constexpr,
-    max_pages: tl.constexpr,
+    max_pages,  # runtime: constexpr here recompiles per page-table width
     window_pages: tl.constexpr,  # 0 = full history; >0 = sliding ring width
     BLOCK_SIZE: tl.constexpr,
 ):
@@ -79,6 +79,11 @@ def compute_out_cache_loc_kernel(
     IS the construction that makes addressing a slid-out column
     impossible. Columns the scheduler punched to the null page read as
     page 0 and route to slot 0 either way.
+
+    ``max_pages``, by contrast, is a runtime scalar: the page table's width
+    grows with context, so specializing on it recompiles the kernel once per
+    distinct width. It only feeds the overflow test, the clamp and the row
+    stride, none of which need a compile-time constant.
     """
     # Program ID represents which request we're processing
     req_idx = tl.program_id(0)
@@ -167,7 +172,7 @@ def compute_out_cache_loc(
     )
 
 
-@triton.jit
+@triton.jit(do_not_specialize=["max_pages"])
 def fused_decode_input_prep_kernel(
     # Inputs
     req_pool_indices_ptr,  # [batch_size]
@@ -180,7 +185,7 @@ def fused_decode_input_prep_kernel(
     # Scalars
     uniform_input_length,
     page_size: tl.constexpr,
-    max_pages: tl.constexpr,
+    max_pages,  # runtime: constexpr here recompiles per page-table width
     BLOCK_SIZE: tl.constexpr,
 ):
     """One launch fuses the decode-uniform path's four small kernels.
@@ -268,7 +273,7 @@ def fused_decode_input_prep(
     )
 
 
-@triton.jit
+@triton.jit(do_not_specialize=["max_pages"])
 def dflash_prepare_decode_kernel(
     output_tokens_ptr,
     accept_lengths_ptr,
@@ -282,7 +287,7 @@ def dflash_prepare_decode_kernel(
     verify_width: tl.constexpr,
     draft_query_width: tl.constexpr,
     page_size: tl.constexpr,
-    max_pages: tl.constexpr,
+    max_pages,  # runtime: constexpr here recompiles per page-table width
     max_draft_prefix,
     block_ids_stride: tl.constexpr,
     BLOCK_SIZE: tl.constexpr,

python/tokenspeed/runtime/layers/attention/configs/mha.py

@@ -82,6 +82,12 @@ class MHAConfig(BaseAttnConfig):
             # Target-stack labels don't fit the draft depth; drop so the pool falls back to full attn
             layer_types = ()
         sliding_window_tokens = getattr(hf_config, "sliding_window", None)
+        if draft_block_decode:
+            # A block drafter writes its KV at the target's cache locations, so
+            # it cannot carry a retention policy of its own; its window is an
+            # attention mask the draft applies in its own layers.
+            layer_types = ()
+            sliding_window_tokens = None
         return cls(
             device=server_args.device,
             context_len=model_config.context_len + server_args.spec_context_pad,

python/tokenspeed/runtime/models/base/causal_lm.py

@@ -143,6 +143,39 @@ class BaseCausalLM(nn.Module):
 
             self.model.layers_to_capture = [val + 1 for val in layer_ids]
 
+    def set_dflash_layers_to_capture(
+        self,
+        layer_ids: list[int],
+        incremental_callback=None,
+        slot_bufs: list | None = None,
+    ) -> None:
+        """Capture the target hidden states a DFLASH/DSpark draft consumes.
+
+        Checkpoints name layer *outputs*, but a layer captures the residual
+        entering it -- hence the ``+ 1`` shift, same as EAGLE3.
+        """
+
+        num_layers = len(self.model.layers)
+        if len(set(layer_ids)) != len(layer_ids):
+            raise ValueError("DFLASH target_layer_ids must be unique.")
+        invalid = [val for val in layer_ids if val < 0 or val + 1 >= num_layers]
+        if invalid:
+            raise ValueError(
+                "DFLASH target_layer_ids must map to capturable target layer "
+                f"outputs. Got invalid ids {invalid}; valid range is "
+                f"[0, {num_layers - 2}] for {num_layers} target layers."
+            )
+
+        self.capture_aux_hidden_states = True
+        capture_layers = sorted(val + 1 for val in layer_ids)
+        self.model.layers_to_capture = capture_layers
+        # The draft concatenates captures in ascending layer order.
+        self.model._dflash_capture_idx_map = {
+            layer_idx: i for i, layer_idx in enumerate(capture_layers)
+        }
+        self.model._dflash_incremental_callback = incremental_callback
+        self.model._dflash_slot_bufs = slot_bufs
+
     @torch.no_grad()
     def forward(
         self,

python/tokenspeed/runtime/models/base/transformer_model.py

@@ -71,6 +71,12 @@ class BaseTransformerModel(nn.Module):
         self.layers = self.resolve_layers(config, quant_config, prefix)
         self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
         self.layers_to_capture: list[int] = []
+        # DFLASH incremental projection, populated by
+        # set_dflash_layers_to_capture().
+        self._dflash_capture_idx_map: dict[int, int] = {}
+        self._dflash_incremental_callback = None
+        self._dflash_slot_bufs: list[torch.Tensor] | None = None
+        self._dflash_incr_active = False
 
         self._compile_decoder_stack()
 
@@ -154,6 +160,25 @@ class BaseTransformerModel(nn.Module):
             prefix=add_prefix("layers", prefix),
         )
 
+    def _notify_dflash_capture(
+        self, layer_idx: int, aux_hidden_states: list[torch.Tensor]
+    ) -> None:
+        """Hand the capture layer ``layer_idx`` just appended to the drafter.
+
+        An idle forward skips the attention block, so nothing was appended.
+        """
+        callback = self._dflash_incremental_callback
+        slot_bufs = self._dflash_slot_bufs
+        capture_idx = self._dflash_capture_idx_map.get(layer_idx)
+        if callback is None or slot_bufs is None or capture_idx is None:
+            return
+        if len(aux_hidden_states) != capture_idx + 1:
+            return
+        captured = aux_hidden_states[capture_idx]
+        num_tokens = captured.shape[0]
+        slot_bufs[capture_idx][:num_tokens].copy_(captured)
+        callback(capture_idx, num_tokens)
+
     def forward(
         self,
         input_ids: torch.Tensor,
@@ -193,6 +218,8 @@ class BaseTransformerModel(nn.Module):
 
         for i, layer in enumerate(self.layers):
 
+            capture = i in self.layers_to_capture
+
             with get_global_expert_distribution_recorder().with_current_layer(i):
 
                 hidden_states, residual = layer(
@@ -201,11 +228,12 @@ class BaseTransformerModel(nn.Module):
                     ctx,
                     out_cache_loc,
                     residual,
-                    aux_hidden_states=(
-                        aux_hidden_states if i in self.layers_to_capture else None
-                    ),
+                    aux_hidden_states=aux_hidden_states if capture else None,
                 )
 
+            if capture and self._dflash_incr_active:
+                self._notify_dflash_capture(i, aux_hidden_states)
+
         if not ctx.forward_mode.is_idle():
 
             if residual is None:

python/tokenspeed/runtime/utils/hf_transformers_utils.py

@@ -231,6 +231,31 @@ def _restore_raw_glm_dsa_fields(config: PretrainedConfig, raw_config: dict) -> N
             setattr(config, key, raw_config[key])
 
 
+def _restore_raw_dflash_fields(config: PretrainedConfig, raw_config: dict) -> None:
+    """Re-assert a DFLASH/DSpark draft's sliding window.
+
+    ``Qwen3DSparkModel`` parses as ``Qwen3Config``, which nulls
+    ``sliding_window`` unless ``use_sliding_window`` is set -- a flag these
+    checkpoints never write, since they carry the window in ``dflash_config``.
+    """
+    dflash_config = raw_config.get("dflash_config")
+    if not isinstance(dflash_config, dict):
+        return
+    if getattr(config, "sliding_window", None) is not None:
+        return
+
+    sliding_window = raw_config.get("sliding_window")
+    if sliding_window is None and dflash_config.get("use_swa"):
+        sliding_window = dflash_config.get("swa_window_size")
+    if sliding_window is None:
+        return
+
+    config.sliding_window = int(sliding_window)
+    if hasattr(config, "use_sliding_window"):
+        # transformers gates the window on this flag; keep the two consistent.
+        config.use_sliding_window = True
+
+
 def get_config(
     model: str,
     trust_remote_code: bool,
@@ -327,6 +352,7 @@ def get_config(
 
     _materialize_architectures(config, raw_config)
     _restore_raw_glm_dsa_fields(config, raw_config)
+    _restore_raw_dflash_fields(config, raw_config)
 
     # extract 'text_config'
     text_config = get_hf_text_config(config)

test/ci/eval/minimax-m3-nvfp4-dspark-evalscope-aime26.yaml

@@ -0,0 +1,58 @@
+api_version: ci.tokenspeed.io/v1
+name: eval-minimax-m3-nvfp4-dspark-aime26
+type: eval
+workflow_stage: model-test
+triggers:
+  - manual
+  - nightly
+runner:
+  labels:
+    - b200-4gpu
+env:
+  CI: "true"
+install:
+  - bash test/ci_system/install_deps.sh
+server:
+  command: >-
+    ts serve
+    --model nvidia/MiniMax-M3-NVFP4
+    --attn-tp-size 4
+    --moe-tp-size 4
+    --max-model-len 262144
+    --gpu-memory-utilization 0.95
+    --attention-backend trtllm
+    --kv-cache-dtype fp8
+    --moe-backend flashinfer_trtllm
+    --speculative-algorithm DSPARK
+    --speculative-draft-model-path nvidia/MiniMax-M3-DSpark
+    --speculative-num-steps 7
+    --speculative-eagle-topk 1
+    --speculative-num-draft-tokens 8
+    --disable-kvstore
+    --block-size 128
+    --trust-remote-code
+    --enable-cache-report
+    --host 127.0.0.1
+    --port 8000
+  ready:
+    url: http://127.0.0.1:8000/readiness
+    timeout: 1800
+    interval: 10
+eval:
+  install:
+    - python3 -m uv venv --seed --clear /tmp/evalscope-perf && python3 -m uv pip install --python /tmp/evalscope-perf/bin/python 'evalscope[perf]'
+  command: >-
+    HF_HOME=${RUNNER_TEMP:-/tmp}/hf-eval-cache
+    /tmp/evalscope-perf/bin/evalscope eval
+    --model nvidia/MiniMax-M3-NVFP4
+    --api-url http://127.0.0.1:8000/v1
+    --api-key EMPTY_TOKEN
+    --datasets aime26
+    --dataset-hub huggingface
+    --dataset-args '{"aime26":{"dataset_id":"math-ai/aime26"}}'
+    --eval-batch-size 16
+    --stream
+    --generation-config '{"do_sample":true,"temperature":1.0,"top_p":0.95,"max_tokens":163840,"timeout":1800}'
+report:
+  github_step_summary: true
+score_threshold: 0.86

test/ci_system/test_eval_configs.py

@@ -34,7 +34,7 @@ DATASETS = {
         "dataset_args": {"dataset_id": "math-ai/aime25"},
     },
     "aime26": {
-        "count": 7,
+        "count": 8,
         "dataset_args": {"dataset_id": "math-ai/aime26"},
     },
     "gpqa_diamond": {

test/runtime/test_dspark_config.py

@@ -1,8 +1,9 @@
-"""Semantic coverage for the DSpark drafter's config and CLI contract.
+"""Semantic coverage for the DSpark drafter's config and wiring contract.
 
 These are the cheap guards that catch a mis-specified launch before eight GPUs
 are committed to loading a 1T-parameter target: block geometry, Markov-head
-resolution, algorithm dispatch, and the draft-worker architecture rewrite.
+resolution, algorithm dispatch, the draft-worker architecture rewrite, and the
+target-side capture the draft is fed from.
 """
 
 from __future__ import annotations
@@ -29,6 +30,8 @@ from tokenspeed.runtime.layers.attention.configs.base import (
 from tokenspeed.runtime.layers.attention.configs.mla import (
     resolve_mla_kv_cache_dtype,
 )
+from tokenspeed.runtime.models.base.causal_lm import BaseCausalLM
+from tokenspeed.runtime.models.base.transformer_model import BaseTransformerModel
 from tokenspeed.runtime.models.dspark import _get_markov_params
 from tokenspeed.runtime.utils.hf_transformers_utils import get_config
 
@@ -251,6 +254,217 @@ def test_non_spec_archs_still_get_the_nextn_rewrite(tmp_path) -> None:
     assert config.architectures[0] == "Qwen3ForCausalLMNextN"
 
 
+# --------------------------------------------------------------------------
+# Sliding-window geometry of a windowed DSpark draft
+# --------------------------------------------------------------------------
+
+
+def _write_swa_draft_config(tmp_path, **overrides) -> str:
+    """A DSpark draft whose layers are all sliding, as MiniMax-M3's ships."""
+    return _write_config(
+        tmp_path,
+        architectures=["Qwen3DSparkModel"],
+        layer_types=["sliding_attention", "sliding_attention"],
+        dflash_config={
+            "mask_token_id": 127,
+            "target_layer_ids": [1, 12],
+            "use_swa": True,
+            "swa_window_size": 1024,
+            "markov_rank": 16,
+        },
+        **overrides,
+    )
+
+
+@pytest.mark.parametrize("declare_top_level", (True, False))
+def test_dspark_draft_keeps_the_window_transformers_would_null(
+    tmp_path, declare_top_level: bool
+) -> None:
+    """``Qwen3Config`` drops ``sliding_window`` unless ``use_sliding_window``.
+
+    DSpark checkpoints never write that flag, so without the restore every
+    consumer -- draft layer construction, the draft attention config, the cache
+    recipe -- sees None on a sliding draft. The window is read from the raw
+    config when present, else from ``dflash_config.swa_window_size``.
+    """
+    extra = {"sliding_window": 1024} if declare_top_level else {}
+    path = _write_swa_draft_config(tmp_path, **extra)
+
+    config = get_config(path, trust_remote_code=False, is_draft_worker=True)
+
+    assert config.sliding_window == 1024
+
+
+# --------------------------------------------------------------------------
+# Draft dtype
+# --------------------------------------------------------------------------
+
+
+def _draft_dtype(server_dtype: str, target_dtype: torch.dtype) -> object:
+    from tokenspeed.runtime.engine import event_loop as event_loop_module
+
+    loop = event_loop_module.EventLoop.__new__(event_loop_module.EventLoop)
+    loop.server_args = SimpleNamespace(
+        dtype=server_dtype,
+        quantization=None,
+        speculative_draft_model_quantization=None,
+        trust_remote_code=True,
+        revision=None,
+        max_model_len=262144,
+        hf_overrides="{}",
+    )
+    loop.model_config = SimpleNamespace(dtype=target_dtype)
+
+    with mock.patch.object(event_loop_module, "ModelConfig") as model_config:
+        loop._load_model_config("draft", is_draft_worker=True)
+
+    return model_config.call_args.kwargs["dtype"]
+
+
+def test_draft_inherits_the_targets_dtype_instead_of_its_own() -> None:
+    """A DSpark draft stored as fp32 master weights must not land on fp16.
+
+    It is fed the target's hidden states and borrows its LM head, so "auto"
+    has to resolve against the target; the standalone fp32 -> fp16 rule would
+    leave the first GEMM mixing bf16 and fp16 with no kernel.
+    """
+    assert _draft_dtype("auto", torch.bfloat16) is torch.bfloat16
+
+
+def test_an_explicit_dtype_still_wins_for_the_draft() -> None:
+    assert _draft_dtype("bfloat16", torch.bfloat16) == "bfloat16"
+
+
+# --------------------------------------------------------------------------
+# Draft cache grouping
+# --------------------------------------------------------------------------
+
+
+def _draft_attn_config(algorithm: str, layer_types: tuple[str, ...]):
+    from tokenspeed.runtime.layers.attention.configs.mha import MHAConfig
+
+    server_args = SimpleNamespace(
+        speculative_algorithm=algorithm,
+        speculative_num_steps=7,
+        speculative_num_draft_tokens=8,
+        kv_cache_dtype="fp8_e4m3",
+        kv_cache_quant_method="none",
+        device="cuda",
+        attention_backend="trtllm",
+        drafter_attention_backend="trtllm",
+        spec_context_pad=0,
+        prefix_granularity=128,
+        max_num_seqs=16,
+        data_parallel_size=None,
+        max_cudagraph_capture_size=80,
+        chunked_prefill_size=8192,
+        disaggregation_mode="null",
+        attn_tp_size=4,
+        mapping=SimpleNamespace(attn=SimpleNamespace(tp_size=4, dp_size=1)),
+    )
+    model_config = SimpleNamespace(
+        hf_config=SimpleNamespace(layer_types=layer_types, sliding_window=1024),
+        num_attention_layers=len(layer_types),
+        context_len=4096,
+        num_attention_heads=32,
+        num_key_value_heads=8,
+        head_dim=128,
+        dtype=torch.bfloat16,
+    )
+    return MHAConfig.generate(server_args, model_config, is_draft=True)
+
+
+def test_block_draft_shares_the_targets_retention() -> None:
+    """A DSpark draft's window is a mask, not a cache-retention policy.
+
+    Its KV rows are written at the target's cache locations, so they live and
+    die with the target's pages; a sliding cache group of its own would both
+    evict rows the target still owns and collide with the target's planes.
+    """
+    config = _draft_attn_config("DSPARK", ("sliding_attention",) * 6)
+
+    assert config.layer_types == ()
+    assert config.sliding_window_tokens is None
+
+
+def test_a_non_block_draft_keeps_its_own_labels() -> None:
+    config = _draft_attn_config("EAGLE3", ("sliding_attention",) * 6)
+
+    assert config.layer_types == ("sliding_attention",) * 6
+    assert config.sliding_window_tokens == 1024
+
+
+# --------------------------------------------------------------------------
+# Target-side capture on the base stack
+# --------------------------------------------------------------------------
+
+
+class _CaptureModel:
+    """Stand-in for the transformer stack: the capture state, no modules."""
+
+    def __init__(self, num_layers: int) -> None:
+        self.layers = [object()] * num_layers
+        self.layers_to_capture = []
+        self._dflash_capture_idx_map = {}
+        self._dflash_incremental_callback = None
+        self._dflash_slot_bufs = None
+        self._dflash_incr_active = False
+
+    notify = BaseTransformerModel._notify_dflash_capture
+
+
+class _CaptureCausalLM:
+    """Stand-in exposing only the setter under test."""
+
+    def __init__(self, num_layers: int) -> None:
+        self.model = _CaptureModel(num_layers)
+        self.capture_aux_hidden_states = False
+
+    set_dflash_layers_to_capture = BaseCausalLM.set_dflash_layers_to_capture
+
+
+def test_taps_shift_by_one_and_sort_for_positional_concat() -> None:
+    """MiniMax-M3's DSpark taps, shuffled: they name completed-layer outputs,
+    and the draft concatenates the captures in ascending layer order."""
+    causal_lm = _CaptureCausalLM(num_layers=60)
+
+    causal_lm.set_dflash_layers_to_capture([57, 1, 35, 12, 46, 23])
+
+    assert causal_lm.model.layers_to_capture == [2, 13, 24, 36, 47, 58]
+    assert causal_lm.model._dflash_capture_idx_map == {
+        2: 0,
+        13: 1,
+        24: 2,
+        36: 3,
+        47: 4,
+        58: 5,
+    }
+    assert causal_lm.capture_aux_hidden_states is True
+
+
+def test_each_capture_reaches_the_drafter_in_concat_order() -> None:
+    slot_bufs = [torch.zeros(4, 3) for _ in range(2)]
+    seen: list[tuple[int, int]] = []
+    causal_lm = _CaptureCausalLM(num_layers=8)
+    causal_lm.set_dflash_layers_to_capture(
+        [1, 5],
+        incremental_callback=lambda idx, num_tokens: seen.append((idx, num_tokens)),
+        slot_bufs=slot_bufs,
+    )
+    model = causal_lm.model
+    model._dflash_incr_active = True
+
+    aux_hidden_states: list[torch.Tensor] = []
+    aux_hidden_states.append(torch.ones(2, 3))
+    model.notify(2, aux_hidden_states)
+    aux_hidden_states.append(torch.full((2, 3), 2.0))
+    model.notify(6, aux_hidden_states)
+
+    assert seen == [(0, 2), (1, 2)]
+    assert torch.equal(slot_bufs[0][:2], torch.ones(2, 3))
+    assert torch.equal(slot_bufs[1][:2], torch.full((2, 3), 2.0))
+
+
 # --------------------------------------------------------------------------
 # CLI contract
 # --------------------------------------------------------------------------

test/runtime/test_sliding_cache_loc.py

@@ -1,14 +1,17 @@
-"""Sliding-window write-location kernel.
+"""Write-location kernels: the sliding-window ring and the page-table width.
 
-The SWA draft's page table is a ring: absolute position ``p`` lives in
-column ``(p // P) % window_pages``. These tests pin the three properties the
-multi-step draft depends on: ring mapping (including wrap-around), null-hole
-routing to the safe dummy slot, and agreement with the full-history kernel
-inside the window (before any wrap).
+The SWA draft's page table is a ring: absolute position ``p`` lives in column
+``(p // P) % window_pages``. The first half pins ring mapping (including
+wrap-around), null-hole routing to the dummy slot, and agreement with the
+full-history kernel before any wrap.
+
+The last test pins that table's width as a runtime argument: ``tl.constexpr``
+there recompiles the kernel once per distinct width.
 """
 
 from __future__ import annotations
 
+import inspect
 import os
 import sys
 
@@ -26,8 +29,11 @@ register_cuda_ci(est_time=10, suite="runtime-1gpu")
 from test.runtime.conftest import requires_cuda
 
 from tokenspeed.runtime.execution.cache_loc_kernel import (
+    compute_out_cache_loc_kernel,
     compute_out_cache_loc_sliding,
     compute_out_cache_loc_uniform,
+    dflash_prepare_decode_kernel,
+    fused_decode_input_prep_kernel,
 )
 from tokenspeed.runtime.execution.draft_page_staging import CacheView
 
@@ -169,3 +175,23 @@ def test_capture_replay_address_stability() -> None:
     graph.replay()
     torch.cuda.synchronize()
     assert int(out[0]) == 77 * _P
+
+
+# --------------------------------------------------------------------------
+# Page-table width stays a runtime argument
+# --------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize(
+    "kernel",
+    [
+        compute_out_cache_loc_kernel,
+        fused_decode_input_prep_kernel,
+        dflash_prepare_decode_kernel,
+    ],
+    ids=["out_cache_loc", "fused_decode_input_prep", "dflash_prepare_decode"],
+)
+def test_max_pages_is_not_constexpr(kernel) -> None:
+    """Re-annotating it ``tl.constexpr`` is the regression to catch."""
+    parameter = inspect.signature(kernel.fn).parameters["max_pages"]
+    assert parameter.annotation is inspect.Parameter.empty