diff --git a/docs/configuration/server.md b/docs/configuration/server.md index 64f2082491..6e439ffa8d 100644 --- a/docs/configuration/server.md +++ b/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 diff --git a/python/tokenspeed/runtime/execution/input_buffer.py b/python/tokenspeed/runtime/execution/input_buffer.py index 5a324def39..107bed7520 100644 --- a/python/tokenspeed/runtime/execution/input_buffer.py +++ b/python/tokenspeed/runtime/execution/input_buffer.py @@ -32,6 +32,7 @@ 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: @@ -39,6 +40,7 @@ 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, diff --git a/python/tokenspeed/runtime/execution/model_executor.py b/python/tokenspeed/runtime/execution/model_executor.py index 07ff0a363c..74e91d3c15 100644 --- a/python/tokenspeed/runtime/execution/model_executor.py +++ b/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 diff --git a/python/tokenspeed/runtime/utils/env.py b/python/tokenspeed/runtime/utils/env.py index 144ccda584..91b53a2d3f 100755 --- a/python/tokenspeed/runtime/utils/env.py +++ b/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) diff --git a/test/runtime/test_force_single_token_verify.py b/test/runtime/test_force_single_token_verify.py new file mode 100644 index 0000000000..568f39f07d --- /dev/null +++ b/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]