feat(dspark): expose forced-rejection diagnostics

docs/configuration/server.md

@@ -222,6 +222,12 @@ is intentionally verbose and should only be enabled while debugging. For
 decode-only batches it also logs the anchor, draft candidates, target verify
 tokens, and their position-wise matches.
 
+Set `TOKENSPEED_FORCE_SINGLE_TOKEN_VERIFY=1` to keep the full speculative
+proposal and target-verify work while committing exactly one target token per
+decode round. This is a diagnostic for measuring speculative-round cost
+independently of acceptance, not a serving policy. It is read once at process
+startup and is disabled by default.
+
 ### Per-Request Stats
 
 `--enable-log-request-stats` enriches the scheduler's per-request finish line for

python/tokenspeed/runtime/execution/input_buffer.py

@@ -32,13 +32,15 @@
 from tokenspeed.runtime.execution.forward_batch_info import compute_position_triton
 from tokenspeed.runtime.multimodal.inputs import Modality, substitute_mm_pad_
 from tokenspeed.runtime.utils import get_colorful_logger
+from tokenspeed.runtime.utils.env import envs
 from tokenspeed.runtime.utils.nvtx import nvtx_range
 
 if TYPE_CHECKING:
     from tokenspeed.runtime.execution.runtime_states import RuntimeStates
 
 
 logger = get_colorful_logger(__name__)
+FORCE_SINGLE_TOKEN_VERIFY = envs.TOKENSPEED_FORCE_SINGLE_TOKEN_VERIFY.get()
 
 
 class InputBuffers:
@@ -99,7 +101,11 @@ def __init__(
             self.out_cache_loc_buf = torch.full(
                 (max_num_tokens,), dummy_kv_slot, dtype=torch.int32
             )
-            self.force_single_token_verify_buf = torch.zeros(max_bs, dtype=torch.bool)
+            # A process-wide diagnostic can start every row forced; the
+            # ordinary path leaves this false and remote recovery sets rows.
+            self.force_single_token_verify_buf = torch.full(
+                (max_bs,), FORCE_SINGLE_TOKEN_VERIFY, dtype=torch.bool
+            )
             self.extend_prefix_lens_buf = torch.zeros(max_bs, dtype=torch.int32)
             self.extend_seq_lens_buf = torch.zeros(max_bs, dtype=torch.int32)
 
@@ -296,7 +302,7 @@ def write_decode_input_ids(
                 )
             self.force_single_token_verify_buf[
                 row_offset : row_offset + expected_count
-            ] = force_single_token
+            ] = (force_single_token | FORCE_SINGLE_TOKEN_VERIFY)
             runtime_states.remote_spec_candidate_ready[decode_req_pool_indices] = False
 
         # Decode-only fast path: one fused Triton kernel writes out_cache_loc,

python/tokenspeed/runtime/execution/model_executor.py

@@ -49,7 +49,10 @@
     ForwardMode,
 )
 from tokenspeed.runtime.execution.forward_thread import ForwardThread
-from tokenspeed.runtime.execution.input_buffer import InputBuffers
+from tokenspeed.runtime.execution.input_buffer import (
+    FORCE_SINGLE_TOKEN_VERIFY,
+    InputBuffers,
+)
 from tokenspeed.runtime.execution.model_runner import ModelRunner
 from tokenspeed.runtime.execution.multimodal_runtime import MultimodalRuntime
 from tokenspeed.runtime.execution.nan_guard import NanGuard
@@ -818,7 +821,9 @@ def _apply_force_single_token_verify(
         row_count: int,
         decode_input_ids: list[int] | None,
     ) -> torch.Tensor:
-        if decode_input_ids is None or row_count <= 0:
+        if row_count <= 0 or (
+            decode_input_ids is None and not FORCE_SINGLE_TOKEN_VERIFY
+        ):
             return accept_lengths
         force_mask = self.input_buffers.force_single_token_verify_buf[
             row_offset : row_offset + row_count

python/tokenspeed/runtime/utils/env.py

@@ -241,6 +241,7 @@ class Envs:
     TOKENSPEED_PROFILE_WITH_STACK = EnvBool(True)
     TOKENSPEED_TEST_REQUEST_TIME_STATS = EnvBool(False)
     TOKENSPEED_LOG_SPEC_ACCEPT_LENGTHS = EnvBool(False)
+    TOKENSPEED_FORCE_SINGLE_TOKEN_VERIFY = EnvBool(False)
     TOKENSPEED_PROFILER_DIR = EnvStr("/tmp")
     TOKENSPEED_CI_SMALL_KV_SIZE = EnvInt(-1)
     TOKENSPEED_NVTX = EnvBool(False)

test/runtime/test_force_single_token_verify.py

@@ -0,0 +1,116 @@
+"""Forced-rejection control for acceptance-independent speculative timing."""
+
+from __future__ import annotations
+
+import os
+import sys
+from types import SimpleNamespace
+from unittest import mock
+
+import torch
+
+_TEST_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+sys.path.insert(0, _TEST_DIR)
+sys.path.insert(0, os.path.dirname(_TEST_DIR))
+
+from ci_system.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=5, suite="runtime-1gpu")
+
+from test.runtime.conftest import requires_cuda
+
+import tokenspeed.runtime.execution.input_buffer as input_buffer_module
+import tokenspeed.runtime.execution.model_executor as model_executor_module
+from tokenspeed.runtime.execution.input_buffer import InputBuffers
+from tokenspeed.runtime.execution.model_executor import ModelExecutor
+
+
+def _executor(force_mask: torch.Tensor) -> ModelExecutor:
+    executor = ModelExecutor.__new__(ModelExecutor)
+    executor.input_buffers = SimpleNamespace(force_single_token_verify_buf=force_mask)
+    return executor
+
+
+def test_local_decode_keeps_acceptance_unchanged_by_default() -> None:
+    accept_lengths = torch.tensor([4, 2], dtype=torch.int32)
+    executor = _executor(torch.tensor([True, True]))
+
+    with mock.patch.object(model_executor_module, "FORCE_SINGLE_TOKEN_VERIFY", False):
+        actual = executor._apply_force_single_token_verify(
+            accept_lengths,
+            row_offset=0,
+            row_count=2,
+            decode_input_ids=None,
+        )
+
+    assert actual is accept_lengths
+
+
+def test_remote_recovery_mask_forces_only_marked_rows() -> None:
+    executor = _executor(torch.tensor([False, True, False]))
+
+    with mock.patch.object(model_executor_module, "FORCE_SINGLE_TOKEN_VERIFY", False):
+        actual = executor._apply_force_single_token_verify(
+            torch.tensor([4, 5], dtype=torch.int32),
+            row_offset=1,
+            row_count=2,
+            decode_input_ids=[7, -1],
+        )
+
+    assert actual.tolist() == [1, 5]
+
+
+def test_global_control_forces_local_decode_rows() -> None:
+    executor = _executor(torch.tensor([True, True]))
+
+    with mock.patch.object(model_executor_module, "FORCE_SINGLE_TOKEN_VERIFY", True):
+        actual = executor._apply_force_single_token_verify(
+            torch.tensor([8, 3], dtype=torch.int32),
+            row_offset=0,
+            row_count=2,
+            decode_input_ids=None,
+        )
+
+    assert actual.tolist() == [1, 1]
+
+
+def test_input_buffer_initializes_global_force_mask() -> None:
+    with mock.patch.object(input_buffer_module, "FORCE_SINGLE_TOKEN_VERIFY", True):
+        buffers = InputBuffers(
+            max_bs=3,
+            max_num_tokens=8,
+            page_size=4,
+            dummy_kv_slot=0,
+            state_write_padding_pool_index=0,
+            device="cpu",
+        )
+
+    assert buffers.force_single_token_verify_buf.tolist() == [True, True, True]
+
+
+@requires_cuda
+def test_global_force_is_cuda_graph_capturable() -> None:
+    executor = _executor(torch.ones(2, dtype=torch.bool, device="cuda"))
+    accept_lengths = torch.tensor([8, 3], dtype=torch.int32, device="cuda")
+
+    with mock.patch.object(model_executor_module, "FORCE_SINGLE_TOKEN_VERIFY", True):
+        # Warm up the elementwise selection before capture.
+        executor._apply_force_single_token_verify(
+            accept_lengths,
+            row_offset=0,
+            row_count=2,
+            decode_input_ids=None,
+        )
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            captured = executor._apply_force_single_token_verify(
+                accept_lengths,
+                row_offset=0,
+                row_count=2,
+                decode_input_ids=None,
+            )
+
+    accept_lengths.fill_(6)
+    graph.replay()
+    torch.cuda.synchronize()
+    assert captured.tolist() == [1, 1]