From b1171c804b96474eafd3e2f0354b8ec67bfa8ab7 Mon Sep 17 00:00:00 2001 From: mokashliu Date: Mon, 17 Aug 2026 09:10:46 +0800 Subject: [PATCH] Add graph-aware adaptive K for DFlash Select DFlash verification lengths from CUDA graph coverage, profiled target and drafter costs, accepted-prefix history, and calibrated runtime overhead. Skip drafting at K=0 and capture a dedicated ordinary-decode graph. Co-authored-by: OpenAI Codex Signed-off-by: mokashliu --- .../adaptive_verification.md | 32 +- tests/test_config.py | 22 + tests/test_sampling_params.py | 20 + .../spec_decode/test_adaptive_verification.py | 298 ++++++++++++ tests/v1/spec_decode/test_dynamic_sd_cug.py | 52 +- tests/v1/worker/test_gpu_batch_ordering.py | 84 +++- vllm/config/speculative.py | 20 +- vllm/model_executor/models/qwen3_5.py | 9 +- vllm/model_executor/models/qwen3_next.py | 21 +- vllm/sampling_params.py | 1 + vllm/v1/attention/selector.py | 13 +- vllm/v1/worker/gpu/cudagraph_utils.py | 28 +- vllm/v1/worker/gpu/model_runner.py | 242 ++++++++-- .../gpu/spec_decode/adaptive_verification.py | 4 +- .../gpu/spec_decode/dflash/adaptive_k.py | 446 ++++++++++++++++++ .../gpu/spec_decode/dflash/speculator.py | 1 + 16 files changed, 1224 insertions(+), 69 deletions(-) create mode 100644 vllm/v1/worker/gpu/spec_decode/dflash/adaptive_k.py diff --git a/docs/features/speculative_decoding/adaptive_verification.md b/docs/features/speculative_decoding/adaptive_verification.md index 58c9d57cee71..b420b6728157 100644 --- a/docs/features/speculative_decoding/adaptive_verification.md +++ b/docs/features/speculative_decoding/adaptive_verification.md @@ -4,7 +4,11 @@ Speculative decoding buys fewer decode steps with more compute. At batch size 1 That matters because per-position acceptance decays fast. While the GPU is memory-bound that slot is effectively free and worth the gamble; once it saturates the gamble has a real throughput cost. The crossover moves with load and with workload-dependent acceptance rates, so no static `num_speculative_tokens` is right across concurrencies. -Adaptive verification decides per step how much of the draft to verify instead. Every (request, position) draft slot is scored by its *survival probability*, the running product of that request's per-position confidences, and the highest-scoring slots are admitted until a global budget is spent. Slots compete across requests: position 5 of a confident request can outrank position 1 of a doubtful one, so one request keeps its full block while another could be trimmed after a token or two. +Adaptive verification decides per step how much of the draft to verify instead. +DSpark scores every (request, position) draft slot by its *survival +probability*, the running product of that request's per-position confidences, +and admits the highest-scoring slots until a global budget is spent. DFlash +uses its observed accepted-prefix survival to choose one K for the next batch. The budget itself comes from a cost model profiled at startup. vLLM measures what a step costs at each shape, then picks the token count that maximizes expected accepted tokens per second. @@ -12,7 +16,13 @@ The practical effect is that one configuration holds up across the whole load ra ## Support -Adaptive verification needs per-position acceptance estimates, so today it is only supported for DSpark with a **confidence head**. +Adaptive verification supports: + +- DSpark checkpoints with a **confidence head**. It trims the current + verification batch per request. +- DFlash checkpoints. It chooses a uniform K from batch size, accepted-prefix + history, and the profiled draft/verify costs. K=0 skips the DFlash forward and + runs ordinary target decoding for the next step. ## Usage @@ -32,11 +42,23 @@ vllm serve deepseek-ai/DeepSeek-V4-Flash-DSpark \ Set `enable_adaptive_verification: false` to verify the full block for every request. +For DFlash, use the same flag with `"method": "dflash"`. The configured +`num_speculative_tokens` remains the maximum K. DFlash drafts all mask positions +in one parallel forward, so lowering a nonzero K reduces target verification +work but not draft-model work; K=0 is the only choice that skips the drafter. +The runtime considers compact graph buckets (K=0, 1, 3, 7, ... and the +configured maximum) rather than capturing a graph for every integer K. + ## Requirements and limitations -- The attention backend must tolerate device-decided query lengths, since the CPU lengths only bound them from above. Backends that plan off the CPU lengths are excluded by the attention selector, and rejected at startup for models that hard-wire their backend. -- Full cudagraphs are required: step costs are profiled from captured graphs, so `--enforce-eager` is rejected at startup. -- Not supported with LoRA (the per-token LoRA mapping is built from CPU-side boundaries), pipeline parallelism (cost curves and confidences exist only on the last rank), or output logprobs (to be fixed). +- DSpark requires an attention backend that tolerates device-decided query + lengths. Backends that plan off CPU lengths are rejected at startup. +- Full cudagraphs are required: step costs are profiled from captured graphs, + so `--enforce-eager` is rejected at startup. +- Adaptive verification is not supported with LoRA or pipeline parallelism. + DSpark additionally does not support output logprobs. DFlash automatic K + does not compact the current verification logits and supports output + logprobs. ## Tuning the cost profile diff --git a/tests/test_config.py b/tests/test_config.py index 70c8728d25d4..80de08a51509 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -1767,6 +1767,28 @@ def test_draft_sample_method_gumbel_is_rejected(): ) +def test_adaptive_verification_accepts_dflash(): + speculative_config = SpeculativeConfig( + method="ngram", + num_speculative_tokens=3, + ) + speculative_config.method = "dflash" + speculative_config.enable_adaptive_verification = True + + speculative_config._validate_adaptive_verification() + + +def test_adaptive_verification_rejects_unrelated_methods(): + speculative_config = SpeculativeConfig( + method="ngram", + num_speculative_tokens=3, + ) + speculative_config.enable_adaptive_verification = True + + with pytest.raises(ValueError, match="only supported with DSpark and DFlash"): + speculative_config._validate_adaptive_verification() + + def test_ir_op_priority_default(): """Test that IR op priority defaults are set correctly.""" from vllm.config.kernel import IrOpPriorityConfig diff --git a/tests/test_sampling_params.py b/tests/test_sampling_params.py index 65ab0738c964..4f18fe5d3c4b 100644 --- a/tests/test_sampling_params.py +++ b/tests/test_sampling_params.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from dataclasses import dataclass +from types import SimpleNamespace import pytest @@ -49,3 +50,22 @@ def test_diffusion_accepts_top_k_top_p(): def test_non_diffusion_models_unaffected(): params = SamplingParams(temperature=0.7, top_k=10, seed=42) params.verify(MockModelConfig(), None, None, None) + + +def test_dflash_adaptive_k_allows_output_logprobs(): + params = SamplingParams(logprobs=1) + speculative_config = SimpleNamespace( + method="dflash", enable_adaptive_verification=True + ) + + params._validate_spec_decode(speculative_config) + + +def test_dspark_adaptive_verification_rejects_output_logprobs(): + params = SamplingParams(logprobs=1) + speculative_config = SimpleNamespace( + method="dspark", enable_adaptive_verification=True + ) + + with pytest.raises(ValueError, match="DSpark confidence-based"): + params._validate_spec_decode(speculative_config) diff --git a/tests/v1/spec_decode/test_adaptive_verification.py b/tests/v1/spec_decode/test_adaptive_verification.py index 99708f50f8f2..60667d79ef15 100644 --- a/tests/v1/spec_decode/test_adaptive_verification.py +++ b/tests/v1/spec_decode/test_adaptive_verification.py @@ -5,10 +5,16 @@ import numpy as np +from vllm.v1.attention.selector import _uses_device_decided_verification_lengths from vllm.v1.worker.gpu.async_utils import StepTimingSample from vllm.v1.worker.gpu.spec_decode.adaptive_verification import ( AdaptiveVerificationManager, ) +from vllm.v1.worker.gpu.spec_decode.dflash.adaptive_k import ( + DFlashAdaptiveKManager, + DFlashAdaptiveKPolicy, + get_dflash_k_candidates, +) from vllm.v1.worker.gpu.structured_outputs import _build_grammar_mapping @@ -217,3 +223,295 @@ def test_zero_budget_keeps_one_grammar_row_per_scheduled_draft(): # (request, position) keys, so the kernel can mask rows the compacted # device layout no longer has room for. assert mapping == [0, 1, 2, 3, 4, 5, 6] + + +def _make_dflash_policy( + draft_cost_ms: np.ndarray, + verify_cost_ms: np.ndarray, + max_k: int = 3, +) -> DFlashAdaptiveKPolicy: + policy = DFlashAdaptiveKPolicy(max_k=max_k, history_weight=1.0) + shaped_verify = np.full((len(draft_cost_ms), max_k + 1), np.inf) + for batch_size in range(1, len(draft_cost_ms)): + for k in policy.candidates: + total_tokens = batch_size * (k + 1) + if total_tokens < len(verify_cost_ms): + shaped_verify[batch_size, k] = verify_cost_ms[total_tokens] + policy.set_cost_tables(draft_cost_ms, shaped_verify) + return policy + + +def test_dflash_adaptive_k_uses_graph_covered_verify_cost(): + draft = np.full(9, 0.2) + verify = np.ones(33) + policy = _make_dflash_policy(draft, verify) + + assert policy.select_k(batch_size=4) == 3 + + +def test_dflash_adaptive_k_uses_compact_graph_buckets(): + assert get_dflash_k_candidates(15) == [0, 1, 3, 7, 15] + assert DFlashAdaptiveKPolicy._batch_bucket(0) == 0 + + +def test_dflash_profiles_real_batch_and_query_length_shapes(): + manager = DFlashAdaptiveKManager.__new__(DFlashAdaptiveKManager) + manager.req_states = SimpleNamespace(max_num_batched_tokens=2048, max_num_reqs=32) + manager.num_speculative_steps = 15 + manager.policy = SimpleNamespace(candidates=[0, 1, 3, 7, 15]) + + batches = list(manager.batches_to_profile([1, 2, 4, 8, 16, 32, 64, 128, 512])) + + assert { + "num_tokens": 32 * 16, + "uniform_decode_query_len": 16, + "profile_verify": True, + "context_len": 8192, + } in batches + assert { + "num_tokens": 8 * 8, + "uniform_decode_query_len": 8, + "profile_verify": True, + "context_len": 8192, + } in batches + + +def test_dflash_shape_costs_distinguish_equal_total_token_counts(monkeypatch): + manager = DFlashAdaptiveKManager.__new__(DFlashAdaptiveKManager) + manager.req_states = SimpleNamespace(max_num_batched_tokens=2048, max_num_reqs=32) + manager.num_speculative_steps = 15 + manager._capture_sizes = {1, 2, 4, 8, 16, 32, 64, 128, 512} + manager._outcome_buffers = [] + manager._selected_k_by_batch = {} + manager._selection_uses_by_batch = {} + manager._global_k_cap = 15 + manager.current_k = 15 + captured: dict[str, np.ndarray] = {} + manager.policy = SimpleNamespace( + candidates=[0, 1, 3, 7, 15], + set_cost_tables=lambda draft, verify: captured.update( + draft=draft, verify=verify + ), + reset_history=lambda: None, + ) + monkeypatch.setattr( + "vllm.v1.worker.gpu.spec_decode.dflash.adaptive_k.get_tp_group", + lambda: SimpleNamespace(broadcast_object=lambda value, src: value), + ) + samples = [ + # Same total token count, materially different per-request query shape. + StepTimingSample(1.0, 0.1, 32, 32, True), + StepTimingSample(2.0, 0.2, 2, 1, True), + StepTimingSample(3.0, 0.3, 4, 1, True), + StepTimingSample(4.0, 0.4, 8, 1, True), + StepTimingSample(9.0, 0.5, 32, 2, True), + ] + + manager.set_initial_cost_curves(samples) + + assert captured["verify"][32, 0] == 1.0 + assert captured["verify"][2, 15] == 9.0 + + +def test_dflash_adaptive_k_preserves_full_draft_for_serial_decode(): + draft = np.full(17, 100.0) + verify = np.ones(257) + policy = _make_dflash_policy(draft, verify, max_k=15) + + assert policy.select_k(batch_size=1) == 15 + + +def test_dflash_adaptive_k_falls_back_to_baseline_on_graph_miss(): + draft = np.full(65, 0.5) + verify = np.ones(257) + verify[33:] = 10.0 + policy = _make_dflash_policy(draft, verify) + + assert policy.select_k(batch_size=32) == 0 + + +def test_dflash_adaptive_k_uses_observed_accepted_prefix(): + draft = np.full(9, 0.2) + verify = np.ones(33) + verify[5:] = 1.5 + policy = _make_dflash_policy(draft, verify) + assert policy.select_k(batch_size=4) == 3 + + policy.record_outcomes( + num_sampled=np.ones(4, dtype=np.int32), + num_draft_tokens=np.full(4, 3, dtype=np.int32), + ) + + assert policy.select_k(batch_size=4) == 3 + + +def test_dflash_adaptive_k_calibrates_shared_runtime_overhead(): + draft = np.full(65, 0.2) + verify = np.ones(1025) + policy = _make_dflash_policy(draft, verify, max_k=15) + assert policy.select_k(batch_size=32) == 15 + + # The runtime interval also includes scheduler and sampling work that every + # K pays. Treating it as K=15-only cost would falsely disable drafting. + for _ in range(2): + policy.record_runtime(batch_size=31, k=15, num_sampled=32, elapsed_ms=64.0) + + assert policy.select_k(batch_size=32) == 15 + + +def test_dflash_adaptive_k_disables_drafting_after_rejections(): + draft = np.full(65, 0.2) + verify = np.ones(1025) + policy = _make_dflash_policy(draft, verify, max_k=15) + assert policy.select_k(batch_size=32) == 15 + + policy.record_outcomes( + num_sampled=np.ones(32, dtype=np.int32), + num_draft_tokens=np.full(32, 15, dtype=np.int32), + ) + + assert policy.select_k(batch_size=32) == 0 + + +def test_dflash_adaptive_k_keeps_serial_friendly_small_batches(): + draft = np.full(17, 0.2) + verify = np.ones(257) + policy = _make_dflash_policy(draft, verify, max_k=15) + for _ in range(4): + policy.record_runtime(batch_size=8, k=15, num_sampled=8, elapsed_ms=64.0) + + assert policy.select_k(batch_size=8) != 0 + + +def test_dflash_adaptive_k_never_disables_drafting_for_small_batches(): + draft = np.full(17, 100.0) + verify = np.ones(257) + policy = _make_dflash_policy(draft, verify, max_k=15) + policy.record_outcomes( + num_sampled=np.ones(8, dtype=np.int32), + num_draft_tokens=np.full(8, 15, dtype=np.int32), + ) + + assert policy.select_k(batch_size=8) > 0 + + +def test_dflash_profile_outcomes_do_not_seed_runtime_history(): + draft = np.full(9, 0.2) + verify = np.ones(33) + policy = _make_dflash_policy(draft, verify) + policy.record_outcomes( + num_sampled=np.ones(4, dtype=np.int32), + num_draft_tokens=np.full(4, 3, dtype=np.int32), + ) + assert policy.select_k(batch_size=4) == 3 + + policy.reset_history() + + assert policy.select_k(batch_size=4) == 3 + + +def test_dflash_adaptive_k_trims_current_verification_batch(): + manager = DFlashAdaptiveKManager.__new__(DFlashAdaptiveKManager) + manager.num_speculative_steps = 15 + manager.select_k = lambda batch_size: 0 + + num_tokens = manager.get_num_tokens( + {"r0": 16, "r1": 16}, + {"r0": list(range(15)), "r1": list(range(15))}, + ) + + assert num_tokens == 2 + assert manager.batch_query_len == 1 + assert manager._batch_budget == ( + {"r0": 0, "r1": 0}, + {"r0": 1, "r1": 1}, + 0, + ) + assert not manager.consume_unmodified_batch() + + +def test_dflash_full_k_preserves_the_original_verification_path(): + manager = DFlashAdaptiveKManager.__new__(DFlashAdaptiveKManager) + manager.num_speculative_steps = 15 + manager.select_k = lambda batch_size: 15 + manager._write_idx = 0 + manager._runtime_start_events = [SimpleNamespace(record=lambda: None)] + manager._pending_runtime = [None] + manager._pending_draft_counts = [None] + + num_tokens = manager.get_num_tokens( + {"r0": 16, "r1": 16}, + {"r0": list(range(15)), "r1": list(range(15))}, + ) + + assert num_tokens == 32 + assert manager.consume_unmodified_batch() + assert manager._batch_budget is None + + +def test_dflash_adaptive_k_holds_decision_between_updates(): + manager = DFlashAdaptiveKManager.__new__(DFlashAdaptiveKManager) + manager._outcome_buffers = [] + manager._selected_k_by_batch = {} + manager._selection_uses_by_batch = {} + manager._global_k_cap = 15 + manager.current_k = 15 + manager.decision_interval = 2 + choices = iter((15, 0)) + manager.policy = SimpleNamespace(select_k=lambda batch_size: next(choices)) + + assert manager.select_k(32) == 15 + assert manager.select_k(32) == 15 + assert manager.select_k(32) == 0 + + +def test_dflash_empty_batch_does_not_disable_drafting(): + manager = DFlashAdaptiveKManager.__new__(DFlashAdaptiveKManager) + manager._outcome_buffers = [] + manager._global_k_cap = 15 + manager.current_k = 15 + manager.policy = SimpleNamespace(select_k=lambda batch_size: 0) + + assert manager.select_k(0) == 15 + assert manager._global_k_cap == 15 + + +def test_dflash_outcome_poll_does_not_synchronize_the_gpu(): + manager = DFlashAdaptiveKManager.__new__(DFlashAdaptiveKManager) + event = SimpleNamespace( + query=lambda: False, + synchronize=lambda: (_ for _ in ()).throw(AssertionError("GPU sync")), + ) + manager._copy_events = [event] + manager._pending_draft_counts = [np.ones(2, dtype=np.int32)] + manager._pending_runtime = [None] + + assert not manager._consume_outcomes(0, wait=False) + assert manager._pending_draft_counts[0] is not None + + +def test_dflash_adaptive_k_is_monotonic_within_graph_batch_bucket(): + manager = DFlashAdaptiveKManager.__new__(DFlashAdaptiveKManager) + manager._outcome_buffers = [] + manager._selected_k_by_batch = {} + manager._selection_uses_by_batch = {} + manager._global_k_cap = 15 + manager.current_k = 15 + manager.decision_interval = 1 + choices = iter((15, 3, 7, 0)) + manager.policy = SimpleNamespace(select_k=lambda batch_size: next(choices)) + + assert manager.select_k(32) == 15 + assert manager.select_k(31) == 3 + assert manager.select_k(32) == 3 + assert manager.select_k(31) == 0 + assert manager.proposal_k(32) == 0 + assert manager.select_k(8) == 0 + + +def test_only_dspark_uses_device_decided_verification_lengths(): + dspark = SimpleNamespace(method="dspark", enable_adaptive_verification=True) + dflash = SimpleNamespace(method="dflash", enable_adaptive_verification=True) + + assert _uses_device_decided_verification_lengths(dspark) + assert not _uses_device_decided_verification_lengths(dflash) diff --git a/tests/v1/spec_decode/test_dynamic_sd_cug.py b/tests/v1/spec_decode/test_dynamic_sd_cug.py index ae263ffb7810..b65d9eefd6da 100644 --- a/tests/v1/spec_decode/test_dynamic_sd_cug.py +++ b/tests/v1/spec_decode/test_dynamic_sd_cug.py @@ -103,7 +103,6 @@ def test_dynamic_sd_full_cudagraph_covers_all_uniform_decode_shapes(monkeypatch) "get_pp_group", lambda: SimpleNamespace(is_first_rank=True, is_last_rank=True), ) - vllm_config = _create_vllm_config_for_dsd( max_num_seqs=max_num_seqs, max_spec_tokens=max_spec_tokens, @@ -151,6 +150,57 @@ def test_dynamic_sd_full_cudagraph_covers_all_uniform_decode_shapes(monkeypatch) assert desc.num_active_loras == 0 +def test_dflash_adaptive_k_captures_compact_uniform_query_lengths(monkeypatch): + monkeypatch.setattr( + gpu_cudagraph_utils, + "get_pp_group", + lambda: SimpleNamespace(is_first_rank=True, is_last_rank=True), + ) + monkeypatch.setattr( + gpu_cudagraph_utils.current_platform, + "get_global_graph_pool", + lambda: None, + ) + max_num_seqs = 8 + max_spec_tokens = 3 + vllm_config = _create_vllm_config_for_dsd( + max_num_seqs=max_num_seqs, + max_spec_tokens=max_spec_tokens, + use_dynamic_sd=False, + ) + speculative_config = vllm_config.speculative_config + speculative_config.method = "dflash" + speculative_config.enable_adaptive_verification = True + manager = gpu_cudagraph_utils.CudaGraphManager( + vllm_config=vllm_config, + device=torch.device("cpu"), + cudagraph_mode=CUDAGraphMode.FULL_AND_PIECEWISE, + decode_query_len=max_spec_tokens + 1, + ) + manager._graphs_captured = True + + for query_len in (1, 2, 4): + desc = manager.dispatch( + num_reqs=4, + num_tokens=4 * query_len, + uniform_token_count=query_len, + num_active_loras=0, + ) + assert desc.cg_mode == CUDAGraphMode.FULL + assert desc.uniform_token_count == query_len + + baseline_manager = gpu_cudagraph_utils.CudaGraphManager( + vllm_config=vllm_config, + device=torch.device("cpu"), + cudagraph_mode=CUDAGraphMode.FULL, + decode_query_len=1, + decode_query_lens=[1], + ) + baseline_manager._graphs_captured = True + assert baseline_manager.dispatch(4, 4, 1, 0).cg_mode == CUDAGraphMode.FULL + assert baseline_manager.decode_query_lens == [1] + + def test_dynamic_sd_non_uniform_batch_falls_back_to_piecewise(monkeypatch): """DSD should use PIECEWISE when the batch is not a uniform decode batch. diff --git a/tests/v1/worker/test_gpu_batch_ordering.py b/tests/v1/worker/test_gpu_batch_ordering.py index 5037dcc7a96b..78cc64abd7a2 100644 --- a/tests/v1/worker/test_gpu_batch_ordering.py +++ b/tests/v1/worker/test_gpu_batch_ordering.py @@ -21,7 +21,15 @@ from vllm.v1.attention.backend import CommonAttentionMetadata from vllm.v1.attention.backends.utils import split_decodes_and_prefills -from vllm.v1.worker.gpu.model_runner import GPUModelRunner, sort_batch_req_ids +from vllm.v1.worker.gpu.model_runner import ( + GPUModelRunner, + _get_adaptive_batch_metadata, + _is_dflash_baseline_cudagraph, + sort_batch_req_ids, +) +from vllm.v1.worker.gpu.spec_decode.dflash.adaptive_k import ( + DFlashAdaptiveKManager, +) from vllm.v1.worker.utils import get_uniform_decode_token_count @@ -114,6 +122,72 @@ def test_adaptive_verification_sizes_only_batches_with_drafts(): assert uniform_tok_count is None +def test_dflash_manager_trims_without_using_dspark_varlen_slot(): + decodes = {"d0": (16, 16), "d1": (16, 16)} + runner = _make_runner(decodes, decode_query_len=16) + manager = SimpleNamespace( + get_num_tokens=lambda _num_tokens_per_req, _draft_tokens: 2 + ) + runner.adaptive_verification = None + runner.dflash_adaptive_k = manager + scheduler_output = SimpleNamespace( + num_scheduled_tokens={req_id: 16 for req_id in decodes}, + total_num_scheduled_tokens=32, + scheduled_spec_decode_tokens={req_id: list(range(15)) for req_id in decodes}, + ) + + state, uniform_tok_count = runner.gather_batch_req_state(scheduler_output, False) + + assert state is not None + assert state.num_tokens == 2 + assert uniform_tok_count is None + + +def test_dflash_adaptive_k_uses_sampling_batch_and_records_outcomes(): + calls: dict[str, Any] = {} + + def proposal_k(batch_size: int) -> int: + calls["batch_size"] = batch_size + return 0 + + manager = SimpleNamespace( + proposal_k=proposal_k, + record_outcomes=lambda num_sampled, input_batch: calls.update( + num_sampled=num_sampled, input_batch=input_batch + ), + ) + runner: Any = GPUModelRunner.__new__(GPUModelRunner) + runner.num_speculative_steps = 15 + runner.dflash_adaptive_k = manager + input_batch = SimpleNamespace( + num_computed_tokens_np=np.array([16, 8], dtype=np.int32), + num_scheduled_tokens=np.array([1, 8], dtype=np.int32), + prefill_len_np=np.array([16, 40], dtype=np.int32), + ) + num_sampled = torch.tensor([1, 0], dtype=torch.int32) + + assert runner._select_dflash_draft_k(input_batch, num_sampled) == 0 + assert calls == { + "batch_size": 1, + "num_sampled": num_sampled, + "input_batch": input_batch, + } + + +def test_dflash_compacted_lengths_reach_hybrid_model_metadata(): + manager = DFlashAdaptiveKManager.__new__(DFlashAdaptiveKManager) + scheduled = np.array([16, 16], dtype=np.int32) + drafts = np.array([15, 15], dtype=np.int32) + compacted = np.array([1, 1], dtype=np.int32) + + metadata_scheduled, metadata_drafts = _get_adaptive_batch_metadata( + manager, scheduled, drafts, compacted + ) + + np.testing.assert_array_equal(metadata_scheduled, [1, 1]) + np.testing.assert_array_equal(metadata_drafts, [0, 0]) + + def test_prompt_chunk_of_decode_query_len_is_not_uniform_decode(): # Two prompt chunks whose query length coincides with the K+1 spec-decode # query length must reject the batch (issue #49918). @@ -142,6 +216,14 @@ def test_dummy_batches_stay_uniform_decode(): assert uniform_tok_count == 8 +def test_uncaptured_dflash_baseline_does_not_change_profile_output_shape(): + # Before KV-cache initialization both managers are None. Identity alone + # must not classify the ordinary profile run as the K=0 baseline graph. + assert not _is_dflash_baseline_cudagraph(None, None) + baseline = object() + assert _is_dflash_baseline_cudagraph(baseline, baseline) + + def test_sort_batch_req_ids_no_spec(): # decode_query_len == 1: plain ascending order (decodes first). num_tokens_per_req = {"p1": 100, "d1": 1, "p2": 7, "d2": 1} diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index 2b93113b7ed3..b94e7c9aa26e 100644 --- a/vllm/config/speculative.py +++ b/vllm/config/speculative.py @@ -239,8 +239,12 @@ class SpeculativeConfig: Mutually exclusive with synthetic_acceptance_rates.""" enable_adaptive_verification: bool = False - """Whether to adaptively size the draft-verification budget from per-request - confidence. Currently only supported for method="dspark".""" + """Whether to adaptively size speculative decoding from runtime costs. + + DSpark sizes the verification budget from per-request confidence. DFlash + chooses one graph-aware draft length for the next batch from accepted-prefix + history. Only supported for methods ``"dspark"`` and ``"dflash"``. + """ @staticmethod def _acceptance_length_to_rates(length: float, n: int) -> list[float]: @@ -1138,11 +1142,19 @@ def __post_init__(self): ) ) - if self.method != "dspark" and self.enable_adaptive_verification: - raise ValueError("Adaptive verification only supported with DSpark") + self._validate_adaptive_verification() return self + def _validate_adaptive_verification(self) -> None: + if self.enable_adaptive_verification and self.method not in ( + "dspark", + "dflash", + ): + raise ValueError( + "Adaptive verification is only supported with DSpark and DFlash" + ) + def _validate_suffix_decoding(self): if not has_arctic_inference(): raise ImportError( diff --git a/vllm/model_executor/models/qwen3_5.py b/vllm/model_executor/models/qwen3_5.py index 28c6a1189937..30410eb20ba1 100644 --- a/vllm/model_executor/models/qwen3_5.py +++ b/vllm/model_executor/models/qwen3_5.py @@ -366,10 +366,15 @@ def forward( positions: torch.Tensor, intermediate_tensors: IntermediateTensors | None = None, inputs_embeds: torch.Tensor | None = None, + return_aux_hidden_states: bool = True, **kwargs: object, ): hidden_states = self.model( - input_ids, positions, intermediate_tensors, inputs_embeds + input_ids, + positions, + intermediate_tensors, + inputs_embeds, + return_aux_hidden_states=return_aux_hidden_states, ) return hidden_states @@ -541,6 +546,7 @@ def forward( positions: torch.Tensor, intermediate_tensors: IntermediateTensors | None = None, inputs_embeds: torch.Tensor | None = None, + return_aux_hidden_states: bool = True, **kwargs: object, ) -> torch.Tensor | IntermediateTensors: """Run forward pass for Qwen3.5. @@ -575,6 +581,7 @@ def forward( positions=positions, intermediate_tensors=intermediate_tensors, inputs_embeds=inputs_embeds, + return_aux_hidden_states=return_aux_hidden_states, ) return hidden_states diff --git a/vllm/model_executor/models/qwen3_next.py b/vllm/model_executor/models/qwen3_next.py index b241baca6cbc..adb3982b2793 100644 --- a/vllm/model_executor/models/qwen3_next.py +++ b/vllm/model_executor/models/qwen3_next.py @@ -623,6 +623,7 @@ def forward( positions: torch.Tensor, intermediate_tensors: IntermediateTensors | None = None, inputs_embeds: torch.Tensor | None = None, + return_aux_hidden_states: bool = True, ) -> torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]: if get_pp_group().is_first_rank: if inputs_embeds is not None: @@ -640,7 +641,11 @@ def forward( hidden_states = sequence_parallel_chunk(hidden_states) assert residual is None - aux_hidden_states = self._maybe_add_hidden_state([], 0, hidden_states, residual) + aux_hidden_states = ( + self._maybe_add_hidden_state([], 0, hidden_states, residual) + if return_aux_hidden_states + else [] + ) for layer_idx, layer in enumerate( islice(self.layers, self.start_layer, self.end_layer), start=self.start_layer, @@ -650,9 +655,10 @@ def forward( hidden_states=hidden_states, residual=residual, ) - self._maybe_add_hidden_state( - aux_hidden_states, layer_idx + 1, hidden_states, residual - ) + if return_aux_hidden_states: + self._maybe_add_hidden_state( + aux_hidden_states, layer_idx + 1, hidden_states, residual + ) if not get_pp_group().is_last_rank: return IntermediateTensors( @@ -791,10 +797,15 @@ def forward( positions: torch.Tensor, intermediate_tensors: IntermediateTensors | None = None, inputs_embeds: torch.Tensor | None = None, + return_aux_hidden_states: bool = True, **kwargs: object, ): hidden_states = self.model( - input_ids, positions, intermediate_tensors, inputs_embeds + input_ids, + positions, + intermediate_tensors, + inputs_embeds, + return_aux_hidden_states=return_aux_hidden_states, ) return hidden_states diff --git a/vllm/sampling_params.py b/vllm/sampling_params.py index 2f0af9e1b43b..0ca722c364b0 100644 --- a/vllm/sampling_params.py +++ b/vllm/sampling_params.py @@ -917,6 +917,7 @@ def _validate_spec_decode( # the logprobs D2H, so it could ride along on that copy. if ( speculative_config.enable_adaptive_verification + and speculative_config.method == "dspark" and self.num_logprobs is not None ): raise ValueError( diff --git a/vllm/v1/attention/selector.py b/vllm/v1/attention/selector.py index 00f2aed5c506..dbe6e95ee8e7 100644 --- a/vllm/v1/attention/selector.py +++ b/vllm/v1/attention/selector.py @@ -21,6 +21,14 @@ logger = init_logger(__name__) +def _uses_device_decided_verification_lengths(speculative_config: object) -> bool: + return bool( + speculative_config is not None + and getattr(speculative_config, "method", None) == "dspark" + and getattr(speculative_config, "enable_adaptive_verification", False) + ) + + class AttentionSelectorConfig(NamedTuple): head_size: int dtype: torch.dtype @@ -139,9 +147,8 @@ def get_attn_backend( ) speculative_config = vllm_config.speculative_config - use_adaptive_verification = ( - speculative_config is not None - and speculative_config.enable_adaptive_verification + use_adaptive_verification = _uses_device_decided_verification_lengths( + speculative_config ) if use_adaptive_verification: from vllm.compilation.backends import model_tag diff --git a/vllm/v1/worker/gpu/cudagraph_utils.py b/vllm/v1/worker/gpu/cudagraph_utils.py index 8e24a0ee0259..084baac8ef2e 100644 --- a/vllm/v1/worker/gpu/cudagraph_utils.py +++ b/vllm/v1/worker/gpu/cudagraph_utils.py @@ -108,6 +108,7 @@ def __init__( decode_query_len: int, lora_capture_cases: list[int] | None = None, varlen_decode: bool = False, + decode_query_lens: list[int] | None = None, ): self.vllm_config = vllm_config self.device = device @@ -116,6 +117,7 @@ def __init__( assert self.compilation_config is not None self.cudagraph_mode = cudagraph_mode self.decode_query_len = decode_query_len + self.decode_query_lens = decode_query_lens self.varlen_decode = varlen_decode self.dp_size = vllm_config.parallel_config.data_parallel_size @@ -194,7 +196,9 @@ def _init_candidates(self) -> None: # draft tokens. The scheduler might use a smaller number so we need # to capture graphs for all possible values during decode. speculative_config = self.vllm_config.speculative_config - if ( + if self.decode_query_lens is not None: + decode_query_lens = self.decode_query_lens + elif ( speculative_config and speculative_config.uses_dynamic_speculative_decoding() ): @@ -213,6 +217,18 @@ def _init_candidates(self) -> None: decode_query_lens = [ x[2] + num_new_sampled_tokens_per_step for x in num_spec_per_batch_size ] + elif ( + speculative_config + and speculative_config.method == "dflash" + and speculative_config.enable_adaptive_verification + ): + from vllm.v1.worker.gpu.spec_decode.dflash.adaptive_k import ( + get_dflash_k_candidates, + ) + + decode_query_lens = [ + k + 1 for k in get_dflash_k_candidates(self.decode_query_len - 1) + ] else: decode_query_lens = [self.decode_query_len] @@ -449,6 +465,7 @@ def __init__( decode_query_len: int, lora_capture_cases: list[int] | None = None, varlen_decode: bool = False, + decode_query_lens: list[int] | None = None, ): super().__init__( vllm_config, @@ -457,6 +474,7 @@ def __init__( decode_query_len, lora_capture_cases=lora_capture_cases, varlen_decode=varlen_decode, + decode_query_lens=decode_query_lens, ) self.hidden_states: torch.Tensor | None = None self.aux_hidden_states: list[torch.Tensor] = [] @@ -474,6 +492,7 @@ def capture( kv_cache_config: KVCacheConfig, has_lora: bool = False, use_aux_hidden_state_outputs: bool = False, + discard_aux_hidden_state_outputs: bool = False, lora_capture_hook: Callable[[int, int, int], None] | None = None, progress_bar_desc: str = "Capturing CUDA graphs", ) -> None: @@ -504,6 +523,8 @@ def create_forward_fn( "positions": input_buffers.positions[:num_tokens], **model_state.prepare_dummy_inputs(num_reqs, num_tokens), } + if discard_aux_hidden_state_outputs: + model_inputs["return_aux_hidden_states"] = False if not self.is_first_pp_rank: # Update for non-first PP ranks. model_inputs["input_ids"] = None @@ -560,6 +581,11 @@ def forward_fn(cg_mode: CUDAGraphMode) -> None: # Last PP rank (common case). if self.use_aux_hidden_state_outputs: hidden_states, aux_hidden_states = model_output + elif discard_aux_hidden_state_outputs and isinstance( + model_output, tuple + ): + hidden_states, _ = model_output + aux_hidden_states = [] else: hidden_states = model_output aux_hidden_states = [] diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index d68666cf4a0c..a7874ab262c5 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -21,7 +21,7 @@ import gc import time from copy import deepcopy -from typing import Any, NamedTuple +from typing import Any, NamedTuple, cast import numpy as np import torch @@ -46,6 +46,7 @@ initialize_mamba_ssu_backend, ) from vllm.model_executor.model_loader import get_model_loader +from vllm.model_executor.models.interfaces import SupportsEagle3 from vllm.model_executor.offloader import ( create_offloader, get_offloader, @@ -133,6 +134,9 @@ AdaptiveVerificationManager, maybe_create_adaptive_verification_manager, ) +from vllm.v1.worker.gpu.spec_decode.dflash.adaptive_k import ( + DFlashAdaptiveKManager, +) from vllm.v1.worker.gpu.spec_decode.eagle.eagle3_utils import ( set_eagle3_aux_hidden_state_layers, ) @@ -155,6 +159,14 @@ logger = init_logger(__name__) +def _is_dflash_baseline_cudagraph( + selected: ModelCudaGraphManager | None, + baseline: ModelCudaGraphManager | None, +) -> bool: + """Return whether an initialized K=0 graph manager was selected.""" + return baseline is not None and selected is baseline + + class GPUModelRunner(LoRAModelRunnerMixin): def __init__(self, vllm_config: VllmConfig, device: torch.device): self.vllm_config = vllm_config @@ -283,6 +295,8 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): num_prefill_lookahead=num_prefill_lookahead, ) self.adaptive_verification: AdaptiveVerificationManager | None = None + self.dflash_adaptive_k: DFlashAdaptiveKManager | None = None + self.dflash_baseline_cudagraph_manager: ModelCudaGraphManager | None = None self.input_buffers = InputBuffers( max_num_reqs=self.max_num_reqs, max_num_tokens=self.max_num_tokens, @@ -548,6 +562,13 @@ def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None: num_bonus_tokens=self.model_state.num_new_sampled_tokens_per_step, max_total_logits=get_max_chunk_logits(self.vocab_size), ) + if getattr(self.speculator, "enable_adaptive_k", False): + self.dflash_adaptive_k = DFlashAdaptiveKManager( + self.req_states, + self.input_buffers.query_start_loc, + self.model_state.num_new_sampled_tokens_per_step, + get_max_chunk_logits(self.vocab_size), + ) self.block_tables = BlockTables( block_sizes=block_sizes, @@ -571,7 +592,7 @@ def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None: initialize_mamba_ssu_backend( self.vllm_config.mamba_config, self.kv_cache_config ) - if self.adaptive_verification is not None: + if self.adaptive_verification is not None or self.dflash_adaptive_k is not None: self.compilation_config.cudagraph_mode = CUDAGraphMode.FULL_AND_PIECEWISE cudagraph_mode = self.compilation_config.resolve_cudagraph_mode_and_sizes( attn_cg_support.min_cg_support, @@ -590,6 +611,15 @@ def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None: lora_capture_cases=self.lora_capture_cases, varlen_decode=self.adaptive_verification is not None, ) + if self.dflash_adaptive_k is not None: + self.dflash_baseline_cudagraph_manager = ModelCudaGraphManager( + self.vllm_config, + self.device, + CUDAGraphMode.FULL, + decode_query_len=1, + lora_capture_cases=self.lora_capture_cases, + decode_query_lens=[1], + ) check_attention_cp_compatibility(self.vllm_config) if isinstance(self.speculator, DraftModelSpeculator): # HACK(woosuk) @@ -636,6 +666,8 @@ def _dummy_run( *args, skip_attn: bool = False, uniform_decode: bool = False, + uniform_decode_query_len: int | None = None, + profile_verify: bool = False, context_len: int = 0, skip_eplb: bool = False, is_profile: bool = False, @@ -651,14 +683,15 @@ def _dummy_run( # Create a dummy scheduler output. num_reqs = min(num_tokens, self.max_num_reqs) - if uniform_decode: + if uniform_decode or uniform_decode_query_len is not None: # HACK(lucas): for now since the worker is shared between MRV1 and MRV2, # and for spec-decode with MTP we want to make sure the dummy runs use # 1+num_speculative_tokens we use max here, this will likely be eventually # changed in the worker: https://github.com/vllm-project/vllm/pull/35243 - num_tokens = max(num_tokens, self.decode_query_len) - num_reqs = num_tokens // self.decode_query_len - assert num_tokens % self.decode_query_len == 0 + query_len = uniform_decode_query_len or self.decode_query_len + num_tokens = max(num_tokens, query_len) + num_reqs = num_tokens // query_len + assert num_tokens % query_len == 0 # Distribute the remainder evenly so no dummy request exceeds # ceil(num_tokens / num_reqs) <= max_model_len tokens. num_tokens_per_request = [ @@ -714,6 +747,9 @@ def _dummy_run( aux_hidden_states = self.execute_model_state.aux_hidden_states self.execute_model_state = None + if profile_verify: + assert hidden_states is not None + self._dummy_sampler_run(hidden_states[input_batch.logits_indices]) self.step_timing.forward_end() # dummy run the eagle speculator's propose to ensure DP/EP sync. @@ -855,7 +891,11 @@ def capture_model(self) -> int: and self.model_state.encoder_runner.has_cudagraph() ) capture_decoder = self.cudagraph_manager.needs_capture() - if not capture_encoder and not capture_decoder: + capture_dflash_baseline = ( + self.dflash_baseline_cudagraph_manager is not None + and self.dflash_baseline_cudagraph_manager.needs_capture() + ) + if not capture_encoder and not capture_decoder and not capture_dflash_baseline: logger.warning( "Skipping encoder and decoder CUDA graph capture. To enable " "encoder capture, ensure `cudagraph_mm_encoder` is enabled; " @@ -874,6 +914,37 @@ def capture_model(self) -> int: if capture_encoder: self.model_state.encoder_runner.capture() + # Capture the ordinary decode specialization before the DFlash + # target specialization. torch.compile guards the auxiliary-layer + # tuple, so doing this second would replay the already-specialized + # target output shape instead of compiling a no-auxiliary path. + if capture_dflash_baseline: + assert self.dflash_baseline_cudagraph_manager is not None + assert self.speculative_config is not None + dflash_target = cast(SupportsEagle3, self.model) + dflash_target.set_aux_hidden_state_layers(()) + try: + self.dflash_baseline_cudagraph_manager.capture( + self.model, + self.model_state, + self.input_buffers, + self.intermediate_tensors, + self.block_tables, + self.attn_groups, + self.kv_cache_config, + has_lora=self.lora_config is not None, + use_aux_hidden_state_outputs=False, + discard_aux_hidden_state_outputs=True, + lora_capture_hook=create_lora_capture_hook( + self.lora_config, self + ), + progress_bar_desc="Capturing DFlash K=0 CUDA graphs", + ) + finally: + set_eagle3_aux_hidden_state_layers( + self.model, self.speculative_config + ) + if capture_decoder: self.cudagraph_manager.capture( self.model, @@ -890,13 +961,14 @@ def capture_model(self) -> int: if self.speculator is not None: with use_workspace_lane(self._draft_workspace_lane): self.speculator.capture() - if self.adaptive_verification is not None: + cost_manager = self.adaptive_verification or self.dflash_adaptive_k + if cost_manager is not None: with self.step_timing.collect() as timings: - for batch in self.adaptive_verification.batches_to_profile( + for batch in cost_manager.batches_to_profile( self.cudagraph_manager.captured_token_counts() ): self._dummy_run(**batch) - self.adaptive_verification.set_initial_cost_curves(timings) + cost_manager.set_initial_cost_curves(timings) end_time = time.perf_counter() end_free_gpu_memory = torch.accelerator.get_memory_info()[0] @@ -1080,10 +1152,13 @@ def gather_batch_req_state( ] is_prefilling_np = num_computed_prefill_tokens_np < prefill_len_np - if self.adaptive_verification is not None and draft_tokens: - num_toks = self.adaptive_verification.get_num_tokens( + verification_manager = self.adaptive_verification or self.dflash_adaptive_k + if verification_manager is not None and draft_tokens: + num_toks = verification_manager.get_num_tokens( num_tokens_per_req, draft_tokens ) + if isinstance(verification_manager, DFlashAdaptiveKManager): + max_query_len = verification_manager.batch_query_len batch_state = BatchReqState( req_ids=req_ids, @@ -1152,7 +1227,9 @@ def prepare_inputs( cu_num_logits = async_copy_to_gpu(cu_num_logits_np, device=self.device) adaptive_verification = ( - self.adaptive_verification if num_draft_tokens_per_req is not None else None + self.adaptive_verification or self.dflash_adaptive_k + if num_draft_tokens_per_req is not None + else None ) num_scheduled_tokens_upper_bound = num_scheduled_tokens_np if adaptive_verification is not None: @@ -1166,6 +1243,21 @@ def prepare_inputs( cu_num_logits_np, ) ) + if ( + isinstance(adaptive_verification, DFlashAdaptiveKManager) + and adaptive_verification.consume_unmodified_batch() + ): + adaptive_verification = None + else: + ( + num_scheduled_tokens_upper_bound, + num_draft_tokens_per_req, + ) = _get_adaptive_batch_metadata( + adaptive_verification, + num_scheduled_tokens_upper_bound, + num_draft_tokens_per_req, + num_scheduled_tokens_np, + ) # Get query_start_loc. # num_reqs_padded is None for PIECEWISE graphs (no request padding needed) @@ -1395,6 +1487,23 @@ def postprocess_sampled( idx_mapping, num_sampled, self.req_states.num_computed_tokens.gpu ) + def _select_dflash_draft_k( + self, + input_batch: InputBatch, + num_sampled: torch.Tensor, + ) -> int: + manager = self.dflash_adaptive_k + if manager is None: + return self.num_speculative_steps + num_sampling_reqs = int( + np.count_nonzero( + input_batch.num_computed_tokens_np + input_batch.num_scheduled_tokens + >= input_batch.prefill_len_np + ) + ) + manager.record_outcomes(num_sampled, input_batch) + return manager.proposal_k(num_sampling_reqs) + def _merge_ec_connector_no_forward( self, scheduler_output: SchedulerOutput, output: ModelRunnerOutput ) -> ModelRunnerOutput: @@ -1453,8 +1562,22 @@ def execute_model( # cross-attention cache with dynamic encoder outputs. skip_compiled = True + cudagraph_manager = self.cudagraph_manager + if ( + self.dflash_baseline_cudagraph_manager is not None + and self.dflash_adaptive_k is not None + and self.dflash_adaptive_k.current_k == 0 + and batch_req_state is not None + and not batch_req_state.has_prefill + and uniform_tok_count == 1 + ): + cudagraph_manager = self.dflash_baseline_cudagraph_manager + use_dflash_baseline = _is_dflash_baseline_cudagraph( + cudagraph_manager, self.dflash_baseline_cudagraph_manager + ) + batch_desc, num_tokens_across_dp = dispatch_cg_and_sync_dp( - self.cudagraph_manager, + cudagraph_manager, num_reqs, num_toks, uniform_tok_count, @@ -1599,6 +1722,8 @@ def execute_model( # values above. **self.model_state.prepare_inputs(input_batch, self.req_states), } + if use_dflash_baseline: + model_inputs["return_aux_hidden_states"] = False if not self.is_first_pp_rank: # Update for non-first PP ranks. model_inputs["input_ids"] = None @@ -1630,9 +1755,9 @@ def execute_model( # Use explicit cudagraph replay for FULL mode. # NOTE(woosuk): Here, we don't need to pass the input tensors, # because they are already copied to the CUDA graph input buffers. - assert self.cudagraph_manager is not None + assert cudagraph_manager is not None self.kv_connector.pre_forward(scheduler_output) - model_output = self.cudagraph_manager.run_fullgraph(batch_desc) + model_output = cudagraph_manager.run_fullgraph(batch_desc) else: # For piecewise and eager mode, just call model(). batch_descriptor = BatchDescriptor( @@ -1666,10 +1791,15 @@ def execute_model( model_output = self.model(**model_inputs) if self.is_last_pp_rank: - if self.use_aux_hidden_state_outputs: + use_aux_hidden_state_outputs = ( + self.use_aux_hidden_state_outputs and not use_dflash_baseline + ) + if use_aux_hidden_state_outputs: assert isinstance(model_output, tuple) hidden_states, aux_hidden_states = model_output else: + if isinstance(model_output, tuple): + model_output, _ = model_output assert isinstance(model_output, torch.Tensor) hidden_states = model_output aux_hidden_states = None @@ -1814,43 +1944,49 @@ def sample_tokens( input_batch.query_start_loc, ) + published_draft_tokens = self.req_states.draft_tokens[input_batch.idx_mapping] if self.speculator is not None: assert self.sampler is not None - # Let the target override the hidden state fed to the drafter - # (e.g. DeepSeek V4 MTP needs the pre-hc_head residual). The - # target returns a persistent buffer sized at max_num_batched_tokens; - # slice to the active token count that propose() expects. - spec_hidden_states = hidden_states - if hasattr(self.model, "get_mtp_target_hidden_states"): - pre_hc_hidden_states = self.model.get_mtp_target_hidden_states() - spec_hidden_states = pre_hc_hidden_states[: hidden_states.shape[0]] # type: ignore[union-attr] - with use_workspace_lane(self._draft_workspace_lane): - draft_tokens = self.speculator.propose( - input_batch, - attn_metadata, - slot_mappings_by_layer, - spec_hidden_states, - aux_hidden_states, - num_sampled, - num_rejected, - self.req_states.last_sampled_tokens, - self.req_states.next_prefill_tokens, - self.sampler.sampling_states.temperature.gpu, - self.sampler.sampling_states.seeds.gpu, - mm_inputs=mm_inputs, - ) - self.req_states.draft_tokens[input_batch.idx_mapping] = draft_tokens - if self.adaptive_verification is not None: - self.adaptive_verification.record_confidences( - self.speculator.draft_token_confidence_probs, input_batch - ) + draft_k = self._select_dflash_draft_k(input_batch, num_sampled) + if draft_k > 0: + # Let the target override the hidden state fed to the drafter + # (e.g. DeepSeek V4 MTP needs the pre-hc_head residual). The + # target returns a persistent buffer sized at max_num_batched_tokens; + # slice to the active token count that propose() expects. + spec_hidden_states = hidden_states + if hasattr(self.model, "get_mtp_target_hidden_states"): + pre_hc_hidden_states = self.model.get_mtp_target_hidden_states() + spec_hidden_states = pre_hc_hidden_states[: hidden_states.shape[0]] + with use_workspace_lane(self._draft_workspace_lane): + draft_tokens = self.speculator.propose( + input_batch, + attn_metadata, + slot_mappings_by_layer, + spec_hidden_states, + aux_hidden_states, + num_sampled, + num_rejected, + self.req_states.last_sampled_tokens, + self.req_states.next_prefill_tokens, + self.sampler.sampling_states.temperature.gpu, + self.sampler.sampling_states.seeds.gpu, + mm_inputs=mm_inputs, + ) + self.req_states.draft_tokens[input_batch.idx_mapping] = draft_tokens + published_draft_tokens = draft_tokens[:, :draft_k] + if self.adaptive_verification is not None: + self.adaptive_verification.record_confidences( + self.speculator.draft_token_confidence_probs, input_batch + ) + else: + published_draft_tokens = published_draft_tokens[:, :0] if self.num_speculative_steps > 0: # Spec-decode and diffusion LLMs both use draft tokens but the latter does # not have a speculator (i.e. self.speculator is None) self.draft_tokens_handler.set_draft_tokens( input_batch, - self.req_states.draft_tokens[input_batch.idx_mapping], + published_draft_tokens, ) # Post-step KV connector related operations. @@ -1920,6 +2056,7 @@ def shutdown(self) -> None: memory is reclaimable when running in the same process.""" torch.accelerator.synchronize() self.cudagraph_manager = None + self.dflash_baseline_cudagraph_manager = None if hasattr(self, "kv_caches"): self.kv_caches.clear() if hasattr(self, "attn_groups"): @@ -2003,6 +2140,19 @@ class BatchReqState(NamedTuple): has_prefill: bool +def _get_adaptive_batch_metadata( + manager: AdaptiveVerificationManager, + scheduled_tokens: np.ndarray, + scheduled_drafts: np.ndarray, + compacted_tokens: np.ndarray, +) -> tuple[np.ndarray, np.ndarray]: + """Expose DFlash's exact uniform K to hybrid model metadata.""" + if isinstance(manager, DFlashAdaptiveKManager): + non_draft_tokens = scheduled_tokens - scheduled_drafts + return compacted_tokens, compacted_tokens - non_draft_tokens + return scheduled_tokens, scheduled_drafts + + def sort_batch_req_ids( num_tokens_per_req: dict[str, int], draft_tokens: dict[str, list[int]], diff --git a/vllm/v1/worker/gpu/spec_decode/adaptive_verification.py b/vllm/v1/worker/gpu/spec_decode/adaptive_verification.py index 12dc8c7c656e..42a03306b190 100644 --- a/vllm/v1/worker/gpu/spec_decode/adaptive_verification.py +++ b/vllm/v1/worker/gpu/spec_decode/adaptive_verification.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Adaptive verification for DSpark speculative decoding.""" +"""Cost profiling and adaptive verification for speculative decoding.""" from collections import defaultdict from collections.abc import Iterable, Iterator @@ -234,7 +234,7 @@ def set_cost_curves( self.req_states.max_num_batched_tokens, self._cudagraph_limit, ) - logger.debug("DSpark cost tables: %s", self.cost_tables) + logger.debug("Adaptive speculative decoding cost tables: %s", self.cost_tables) def record_confidences( self, diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/adaptive_k.py b/vllm/v1/worker/gpu/spec_decode/dflash/adaptive_k.py new file mode 100644 index 000000000000..adb68a3eae52 --- /dev/null +++ b/vllm/v1/worker/gpu/spec_decode/dflash/adaptive_k.py @@ -0,0 +1,446 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from collections import defaultdict +from collections.abc import Iterator +from typing import TYPE_CHECKING + +import numpy as np +import torch + +import vllm.envs as envs +from vllm.distributed.parallel_state import get_tp_group +from vllm.logger import init_logger +from vllm.utils.gpu_sync_debug import gpu_sync_allowed +from vllm.v1.utils import CpuGpuBuffer +from vllm.v1.worker.gpu.async_utils import stream +from vllm.v1.worker.gpu.spec_decode.adaptive_verification import ( + AdaptiveVerificationManager, + build_cost_tables_from_curves, +) + +if TYPE_CHECKING: + from vllm.v1.worker.gpu.input_batch import InputBatch + from vllm.v1.worker.gpu.states import RequestState + +logger = init_logger(__name__) +_PROFILE_REPLAYS = 5 +_RUNTIME_CALIBRATION_SAMPLES = 2 + + +def get_dflash_k_candidates(max_k: int) -> list[int]: + """Return a compact set of K values whose query lengths are powers of two.""" + candidates = [0] + k = 1 + while k < max_k: + candidates.append(k) + k = 2 * k + 1 + if max_k > 0 and candidates[-1] != max_k: + candidates.append(max_k) + return candidates + + +class DFlashAdaptiveKPolicy: + def __init__(self, max_k: int, history_weight: float): + if max_k < 1: + raise ValueError("max_k must be positive") + if not 0.0 < history_weight <= 1.0: + raise ValueError("history_weight must be in (0, 1]") + self.max_k = max_k + self.candidates = get_dflash_k_candidates(max_k) + self.history_weight = history_weight + self._survival = np.ones(max_k, dtype=np.float64) + self._draft_cost_ms: np.ndarray | None = None + self._verify_cost_ms: np.ndarray | None = None + self._runtime_overhead_ms: dict[int, float] = {} + self._runtime_samples: dict[int, int] = {} + + def set_cost_tables( + self, draft_cost_ms: np.ndarray, verify_cost_ms: np.ndarray + ) -> None: + if verify_cost_ms.ndim != 2: + raise ValueError("DFlash verify costs must be indexed by batch and K") + self._draft_cost_ms = draft_cost_ms + self._verify_cost_ms = verify_cost_ms + + def reset_history(self) -> None: + self._survival.fill(1.0) + self._runtime_overhead_ms.clear() + self._runtime_samples.clear() + + @staticmethod + def _batch_bucket(batch_size: int) -> int: + if batch_size < 1: + return 0 + return 1 << (batch_size - 1).bit_length() + + def record_runtime( + self, batch_size: int, k: int, num_sampled: int, elapsed_ms: float + ) -> None: + del num_sampled + if ( + batch_size < 1 + or k == 0 + or elapsed_ms <= 0.0 + or self._verify_cost_ms is None + or batch_size >= len(self._verify_cost_ms) + ): + return + profiled_verify_ms = self._verify_cost_ms[batch_size, k] + if not np.isfinite(profiled_verify_ms): + return + bucket = self._batch_bucket(batch_size) + observed_overhead_ms = max(elapsed_ms - profiled_verify_ms, 0.0) + previous = self._runtime_overhead_ms.get(bucket) + self._runtime_overhead_ms[bucket] = ( + observed_overhead_ms + if previous is None + else (1.0 - self.history_weight) * previous + + self.history_weight * observed_overhead_ms + ) + self._runtime_samples[bucket] = self._runtime_samples.get(bucket, 0) + 1 + + def select_k(self, batch_size: int) -> int: + if batch_size < 1: + return 0 + # At small batch sizes DFlash's full draft is the measured fast path: + # the target is under-occupied and the longer accepted prefix amortizes + # the drafter. Startup microbenchmarks do not include that end-to-end + # under-occupancy benefit, so keep the stable full-K graph there. + if batch_size < 16: + return self.max_k + if self._draft_cost_ms is None or self._verify_cost_ms is None: + return self.max_k + if batch_size >= len(self._draft_cost_ms): + return 0 + + bucket = self._batch_bucket(batch_size) + shared_overhead_ms = 0.0 + if self._runtime_samples.get(bucket, 0) >= _RUNTIME_CALIBRATION_SAMPLES: + shared_overhead_ms = self._runtime_overhead_ms[bucket] + expected_tokens = 1.0 + np.concatenate((np.zeros(1), np.cumsum(self._survival))) + scores = np.full(self.max_k + 1, -np.inf) + for k in self.candidates: + draft_cost = 0.0 if k == 0 else self._draft_cost_ms[batch_size] + cost = draft_cost + self._verify_cost_ms[batch_size, k] + shared_overhead_ms + scores[k] = batch_size * expected_tokens[k] / max(cost, 1e-6) + return int(np.argmax(scores)) + + def record_outcomes( + self, num_sampled: np.ndarray, num_draft_tokens: np.ndarray + ) -> None: + accepted = np.maximum(num_sampled.astype(np.int64) - 1, 0) + drafted = num_draft_tokens.astype(np.int64) + for position in range(self.max_k): + eligible = drafted > position + if not eligible.any(): + continue + observed = np.mean(accepted[eligible] > position) + self._survival[position] = (1.0 - self.history_weight) * self._survival[ + position + ] + self.history_weight * observed + np.minimum.accumulate(self._survival, out=self._survival) + + +class DFlashAdaptiveKManager(AdaptiveVerificationManager): + """Graph-aware batch-level draft-length control for DFlash.""" + + def __init__( + self, + req_states: "RequestState", + query_start_loc: torch.Tensor, + num_bonus_tokens: int, + max_total_logits: int, + history_weight: float = 0.2, + decision_interval: int = 2, + ) -> None: + super().__init__( + req_states, + query_start_loc, + num_bonus_tokens, + max_total_logits, + ) + self.policy = DFlashAdaptiveKPolicy( + self.num_speculative_steps, history_weight=history_weight + ) + if decision_interval < 1: + raise ValueError("decision_interval must be positive") + self.decision_interval = decision_interval + + device = req_states.device + self._outcome_buffers = [ + CpuGpuBuffer( + req_states.max_num_reqs, + dtype=torch.int32, + device=device, + ) + for _ in range(2) + ] + self._copy_events = [torch.cuda.Event(blocking=True) for _ in range(2)] + self._runtime_start_events = [ + torch.cuda.Event(enable_timing=True) for _ in range(2) + ] + self._runtime_end_events = [ + torch.cuda.Event(enable_timing=True) for _ in range(2) + ] + self._pending_runtime: list[tuple[int, int] | None] = [None, None] + self._pending_draft_counts: list[np.ndarray | None] = [None, None] + self._write_idx = 0 + self._selected_k_by_batch: dict[int, int] = {} + self._selection_uses_by_batch: dict[int, int] = {} + self._global_k_cap = self.num_speculative_steps + self.current_k = self.num_speculative_steps + self._batch_is_unmodified = False + + def batches_to_profile(self, capture_sizes: list[int]) -> Iterator[dict[str, int]]: + """Profile real ``(batch, K + 1)`` verification shapes. + + Equal total token counts can have very different hybrid-attention cost + (for example, 32 requests x 16 tokens versus 128 x 4). The generic + one-dimensional profiler intentionally cannot distinguish them. + """ + self._capture_sizes = set(capture_sizes) + max_reqs = self.req_states.max_num_reqs + max_tokens = self.req_states.max_num_batched_tokens + query_lens = [k + 1 for k in self.policy.candidates] + batch_sizes = {1, max_reqs} + for query_len in query_lens: + batch_sizes.update( + size // query_len + for size in capture_sizes + if size % query_len == 0 and 0 < size // query_len <= max_reqs + ) + for batch_size in sorted(batch_sizes): + for query_len in query_lens: + num_tokens = batch_size * query_len + if num_tokens > max_tokens: + continue + for _ in range(_PROFILE_REPLAYS): + yield { + "num_tokens": num_tokens, + "uniform_decode_query_len": query_len, + "profile_verify": True, + "context_len": ( + envs.VLLM_ADAPTIVE_VERIFICATION_PROFILE_CONTEXT_LEN + ), + } + + def set_initial_cost_curves(self, samples: list) -> None: + grouped: defaultdict[tuple[int, int], list[float]] = defaultdict(list) + draft_grouped: defaultdict[int, list[float]] = defaultdict(list) + max_query_len = self.num_speculative_steps + 1 + for sample in samples: + if sample.num_reqs < 1 or sample.num_target_tokens % sample.num_reqs: + continue + query_len = sample.num_target_tokens // sample.num_reqs + grouped[(query_len, sample.num_reqs)].append(sample.forward_ms) + if query_len == max_query_len and sample.full_cudagraph: + draft_grouped[sample.num_reqs].append(sample.drafter_ms) + + curves = { + query_len: [ + (batch_size, float(np.median(values))) + for (shape_query_len, batch_size), values in sorted(grouped.items()) + if shape_query_len == query_len + ] + for query_len in (k + 1 for k in self.policy.candidates) + } + draft_curve = [ + (batch_size, float(np.median(values))) + for batch_size, values in sorted(draft_grouped.items()) + ] + draft_curve, curves = get_tp_group().broadcast_object( + (draft_curve, curves), src=0 + ) + if not draft_curve or any(not curve for curve in curves.values()): + raise RuntimeError( + "DFlash adaptive K could not profile every verification shape. " + "Pass enable_adaptive_verification=false to use a fixed K." + ) + + max_reqs = self.req_states.max_num_reqs + draft_table, _ = build_cost_tables_from_curves( + draft_curve, + [(1, 1.0)], + max_reqs, + max_reqs, + cudagraph_limit=max_reqs, + ) + verify_table = np.full( + (max_reqs + 1, self.num_speculative_steps + 1), + np.inf, + dtype=np.float64, + ) + for k in self.policy.candidates: + query_len = k + 1 + captured_batches = [ + size // query_len + for size in self._capture_sizes + if size % query_len == 0 + ] + capture_limit = min(max(captured_batches, default=0), max_reqs) + _, costs = build_cost_tables_from_curves( + [(1, 0.0)], + curves[query_len], + max_reqs, + max_reqs, + cudagraph_limit=capture_limit, + ) + verify_table[:, k] = costs + + self.cost_tables = (draft_table, verify_table) + self.policy.set_cost_tables(draft_table, verify_table) + for idx in range(len(self._outcome_buffers)): + self._consume_outcomes(idx) + self.policy.reset_history() + self._selected_k_by_batch.clear() + self._selection_uses_by_batch.clear() + self._global_k_cap = self.num_speculative_steps + self.current_k = self.num_speculative_steps + + def get_num_tokens( + self, + num_tokens_per_req: dict[str, int], + draft_tokens: dict[str, list[int]], + ) -> int: + """Trim the current target verification batch to the selected K.""" + req_ids = list(num_tokens_per_req) + scheduled_drafts = np.fromiter( + (len(draft_tokens.get(req_id, ())) for req_id in req_ids), + dtype=np.int32, + count=len(req_ids), + ) + num_non_draft_tokens = np.fromiter( + ( + num_tokens_per_req[req_id] - scheduled_drafts[idx] + for idx, req_id in enumerate(req_ids) + ), + dtype=np.int32, + count=len(req_ids), + ) + k = self.select_k(int(np.count_nonzero(scheduled_drafts))) + batch_size = int(np.count_nonzero(scheduled_drafts)) + if k > 0 and batch_size > 0: + idx = self._write_idx + # Do not overwrite timing metadata if the double-buffer slot is + # still carrying a result from two steps ago. This wait is only on + # slot reuse; the normal policy poll above remains non-blocking. + self._consume_outcomes(idx) + self._runtime_start_events[idx].record() + self._pending_runtime[idx] = (batch_size, k) + admitted_drafts = np.minimum(scheduled_drafts, k) + self.batch_query_len = int( + np.max(num_non_draft_tokens + admitted_drafts, initial=1) + ) + num_drafts_per_req = { + req_id: int(num_drafts) + for req_id, num_drafts in zip(req_ids, admitted_drafts, strict=True) + } + num_non_draft_tokens_per_req = { + req_id: int(num_tokens) + for req_id, num_tokens in zip(req_ids, num_non_draft_tokens, strict=True) + } + draft_budget = int(admitted_drafts.sum()) + self._batch_budget = ( + num_drafts_per_req, + num_non_draft_tokens_per_req, + draft_budget, + ) + self._batch_is_unmodified = draft_budget == int(scheduled_drafts.sum()) + return int(num_non_draft_tokens.sum()) + draft_budget + + def consume_unmodified_batch(self) -> bool: + if not self._batch_is_unmodified: + return False + self._batch_is_unmodified = False + self._batch_budget = None + return True + + def _consume_outcomes(self, idx: int, *, wait: bool = True) -> bool: + draft_counts = self._pending_draft_counts[idx] + if draft_counts is None: + return True + if not wait and not self._copy_events[idx].query(): + return False + if wait: + with gpu_sync_allowed(): + self._copy_events[idx].synchronize() + self.policy.record_outcomes( + self._outcome_buffers[idx].np[: len(draft_counts)], draft_counts + ) + runtime = self._pending_runtime[idx] + if runtime is not None: + self.policy.record_runtime( + *runtime, + int(self._outcome_buffers[idx].np[: len(draft_counts)].sum()), + self._runtime_start_events[idx].elapsed_time( + self._runtime_end_events[idx] + ), + ) + self._pending_runtime[idx] = None + self._pending_draft_counts[idx] = None + return True + + def select_k(self, batch_size: int) -> int: + for idx in range(len(self._outcome_buffers)): + self._consume_outcomes(idx, wait=False) + if batch_size < 1: + self.current_k = self._global_k_cap + return self.current_k + if self._global_k_cap == 0: + self.current_k = 0 + return 0 + bucket = DFlashAdaptiveKPolicy._batch_bucket(batch_size) + previous = self._selected_k_by_batch.get(bucket) + selection_uses = self._selection_uses_by_batch.get(bucket, 0) + if previous is not None and selection_uses < self.decision_interval: + self._selection_uses_by_batch[bucket] = selection_uses + 1 + self.current_k = min(previous, self._global_k_cap) + return self.current_k + k = self.policy.select_k(batch_size) + # Requests shrink within a decode batch. Re-enabling a longer draft + # after shortening it creates shape churn and can reuse stale drafter + # state after K=0. A graph bucket therefore only moves toward shorter + # verification shapes for the lifetime of the manager. + if previous is not None: + k = min(previous, k) + k = min(k, self._global_k_cap) + self._global_k_cap = min(self._global_k_cap, k) + self.current_k = k + if previous != k: + logger.info("DFlash adaptive K: batch_size=%d, K=%d", batch_size, k) + self._selected_k_by_batch[bucket] = k + self._selection_uses_by_batch[bucket] = 1 + return k + + def proposal_k(self, batch_size: int) -> int: + bucket = DFlashAdaptiveKPolicy._batch_bucket(batch_size) + if bucket not in self._selected_k_by_batch: + return self.select_k(batch_size) + self.current_k = min(self._selected_k_by_batch[bucket], self._global_k_cap) + return self.current_k + + def record_outcomes( + self, + num_sampled: torch.Tensor, + input_batch: "InputBatch", + ) -> None: + draft_counts = input_batch.num_draft_tokens_per_req + if draft_counts is None or not np.any(draft_counts): + return + + idx = self._write_idx + self._consume_outcomes(idx) + draft_counts = draft_counts.copy() + num_reqs = len(draft_counts) + slot = self._outcome_buffers[idx] + slot.gpu[:num_reqs].copy_(num_sampled[:num_reqs]) + self._runtime_end_events[idx].record() + + current_stream = torch.cuda.current_stream(self.req_states.device) + self._copy_stream.wait_stream(current_stream) + with stream(self._copy_stream, current_stream): + slot.copy_to_cpu(num_reqs) + self._copy_events[idx].record() + self._pending_draft_counts[idx] = draft_counts + self._write_idx = (idx + 1) % len(self._outcome_buffers) diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py index eb5dc470b26c..c68fcfc34dfa 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py @@ -44,6 +44,7 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): # Each request emits exactly (bonus + N mask) query tokens per step. self.num_query_per_req = 1 + self.num_speculative_steps + self.enable_adaptive_k = self.speculative_config.enable_adaptive_verification self.parallel_drafting_token_id = get_parallel_drafting_token_id( self.draft_model_config.hf_config