diff --git a/tests/test_config.py b/tests/test_config.py index 4d39b56d265f..e3dfa303c3a8 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -538,12 +538,13 @@ def test_dsa_models_select_matching_mtp(model_type, expected_architecture): assert hf_config.architectures == [expected_architecture] -def test_v2_model_runner_supports_extract_hidden_states(): +@pytest.mark.parametrize("method", ["extract_hidden_states", "ngram", "ngram_gpu"]) +def test_v2_model_runner_supports_speculative_method(method): config = VllmConfig() config.speculative_config = cast( SpeculativeConfig, SimpleNamespace( - method="extract_hidden_states", + method=method, parallel_drafting=False, enable_adaptive_verification=False, ), diff --git a/tests/v1/spec_decode/test_max_len.py b/tests/v1/spec_decode/test_max_len.py index 81f842419375..17dd318563b7 100644 --- a/tests/v1/spec_decode/test_max_len.py +++ b/tests/v1/spec_decode/test_max_len.py @@ -36,6 +36,39 @@ def test_ngram_max_len(num_speculative_tokens: int, vllm_runner): runner.llm.generate(_PROMPTS, sampling_params) +@pytest.mark.parametrize("num_speculative_tokens", [1, 3, 10]) +@pytest.mark.parametrize("method", ["ngram", "ngram_gpu"]) +def test_ngram_gpu_max_len( + method: str, + num_speculative_tokens: int, + vllm_runner, + monkeypatch: pytest.MonkeyPatch, +): + """V2 n-gram decoding stops at max_model_len.""" + monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1") + with vllm_runner( + "facebook/opt-125m", + trust_remote_code=False, + max_model_len=100, + enable_chunked_prefill=None, + enforce_eager=True, # For faster initialization. + speculative_config={ + "method": method, + "prompt_lookup_max": 5, + "prompt_lookup_min": 3, + "num_speculative_tokens": num_speculative_tokens, + }, + ) as runner: + assert runner.llm.llm_engine.vllm_config.use_v2_model_runner + sampling_params = SamplingParams(max_tokens=100, ignore_eos=True) + outputs = runner.llm.generate(_PROMPTS, sampling_params) + for output in outputs: + assert output.prompt_token_ids is not None + assert ( + len(output.prompt_token_ids) + len(output.outputs[0].token_ids) == 100 + ) + + @pytest.mark.parametrize("num_speculative_tokens", [1, 3, 10]) @pytest.mark.parametrize("attn_backend", get_attn_backend_list_based_on_platform()) def test_eagle_max_len( diff --git a/tests/v1/spec_decode/test_ngram_gpu.py b/tests/v1/spec_decode/test_ngram_gpu.py new file mode 100644 index 000000000000..69393c403f98 --- /dev/null +++ b/tests/v1/spec_decode/test_ngram_gpu.py @@ -0,0 +1,437 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for the V2 GPU-accelerated n-gram speculator. + +These tests target the Triton proposer in +``vllm.v1.worker.gpu.spec_decode.ngram.speculator`` and complement the CPU +``NgramProposer`` tests in ``test_ngram.py``. The GPU speculator follows a +slightly different policy than the CPU one: when multiple n-gram matches of +the same length exist, the GPU kernel picks the right-most (most recent) +match inside the active context, whereas the CPU implementation returns the +left-most. The expectations below reflect the GPU behavior. + +Also covers the GPU draft-trimming layout helpers in +``adaptive_verification`` that ngram_gpu shares with DSpark. +""" + +from __future__ import annotations + +from types import SimpleNamespace + +import numpy as np +import pytest +import torch + +from vllm.config import ( + ModelConfig, + SchedulerConfig, + SpeculativeConfig, + VllmConfig, +) +from vllm.v1.worker.gpu.attn_utils import AttentionCGSupportInfo +from vllm.v1.worker.gpu.spec_decode.adaptive_verification import ( + VariableDraftTrimmer, + build_verification_layout, + maybe_create_draft_trimmer, +) +from vllm.v1.worker.gpu.spec_decode.ngram.speculator import NgramGPUSpeculator +from vllm.v1.worker.gpu.states import RequestState + +if not torch.cuda.is_available(): + pytest.skip( + "CUDA required for NgramGPUSpeculator tests", + allow_module_level=True, + ) + +DEVICE = torch.device("cuda") + + +def _make_vllm_config( + min_n: int, + max_n: int, + k: int, + max_num_seqs: int = 8, + max_model_len: int = 64, + method: str = "ngram_gpu", +) -> VllmConfig: + model_config = ModelConfig( + model="facebook/opt-125m", + max_model_len=max_model_len, + enforce_eager=True, + ) + scheduler_config = SchedulerConfig.default_factory( + max_num_seqs=max_num_seqs, + max_model_len=max_model_len, + ) + speculative_config = SpeculativeConfig( + method=method, + prompt_lookup_min=min_n, + prompt_lookup_max=max_n, + num_speculative_tokens=k, + ) + return VllmConfig( + model_config=model_config, + scheduler_config=scheduler_config, + speculative_config=speculative_config, + ) + + +def _make_request_state(cfg: VllmConfig) -> RequestState: + return RequestState( + max_num_reqs=cfg.scheduler_config.max_num_seqs, + max_model_len=cfg.model_config.max_model_len, + max_num_batched_tokens=cfg.scheduler_config.max_num_batched_tokens, + num_speculative_steps=cfg.speculative_config.num_speculative_tokens, + vocab_size=cfg.model_config.get_vocab_size(), + device=DEVICE, + use_dense_all_token_ids=True, + ) + + +def _make_speculator( + min_n: int, + max_n: int, + k: int, + max_num_seqs: int = 8, + max_model_len: int = 32, +) -> NgramGPUSpeculator: + cfg = _make_vllm_config( + min_n=min_n, + max_n=max_n, + k=k, + max_num_seqs=max_num_seqs, + max_model_len=max_model_len, + ) + return NgramGPUSpeculator(cfg, DEVICE, _make_request_state(cfg)) + + +def _propose( + spec: NgramGPUSpeculator, + rows: list[list[int]], + seq_lens: list[int] | None = None, + num_sampled: list[int] | None = None, + last_sampled: list[int] | None = None, + slots: list[int] | None = None, +) -> tuple[list[list[int]], list[int]]: + """Place each batch row at a request slot and run propose(). + + Returns (drafts, num_valid) as python lists in batch order. + """ + B = len(rows) + if seq_lens is None: + seq_lens = [len(r) for r in rows] + if num_sampled is None: + num_sampled = [1] * B + if last_sampled is None: + last_sampled = [0] * B + if slots is None: + slots = list(range(B)) + + max_num_reqs = spec.max_num_reqs + all_token_ids = spec.req_states.all_token_ids.gpu + total_len = spec.req_states.total_len.gpu + all_token_ids.zero_() + total_len.zero_() + last_sampled_t = torch.zeros((max_num_reqs, 1), dtype=torch.int64, device=DEVICE) + for row, slot, seq_len, last in zip(rows, slots, seq_lens, last_sampled): + if row: + all_token_ids[slot, : len(row)] = torch.tensor( + row, dtype=torch.int32, device=DEVICE + ) + total_len[slot] = seq_len + last_sampled_t[slot, 0] = last + + idx_mapping = torch.tensor(slots, dtype=torch.int64, device=DEVICE) + input_batch = SimpleNamespace(num_reqs=B, idx_mapping=idx_mapping) + + drafts = spec.propose( + input_batch=input_batch, + attn_metadata=None, + slot_mappings=None, + last_hidden_states=torch.empty(0, device=DEVICE), + aux_hidden_states=None, + num_sampled=torch.tensor(num_sampled, dtype=torch.int32, device=DEVICE), + num_rejected=torch.zeros(B, dtype=torch.int32, device=DEVICE), + last_sampled=last_sampled_t, + next_prefill_tokens=torch.zeros(B, dtype=torch.int32, device=DEVICE), + temperature=torch.zeros(B, dtype=torch.float32, device=DEVICE), + seeds=torch.zeros(B, dtype=torch.int64, device=DEVICE), + dp_sync=None, + ) + num_valid = spec.num_valid_drafts_for_trim[idx_mapping] + return drafts.cpu().tolist(), num_valid.cpu().tolist() + + +# --------------------------------------------------------------------------- +# Proposal behavior +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("max_model_len", [32, 300]) +def test_no_match_clears_previous_proposal(max_model_len): + spec = _make_speculator(min_n=2, max_n=2, k=2, max_model_len=max_model_len) + row = [0] * (max_model_len - 5) + [1, 2, 3, 1, 2] + assert _propose(spec, [row]) == ([[3, 1]], [2]) + drafts, num_valid = _propose(spec, [[1, 2, 3, 4, 5]], last_sampled=[42]) + assert num_valid == [0] + assert drafts == [[42, 42]] + + +def test_no_4gram_match_only(): + """No 4-gram match in [1,2,3,4,1,2,3] → 0 valid drafts.""" + spec = _make_speculator(min_n=4, max_n=4, k=2) + drafts, num_valid = _propose(spec, [[1, 2, 3, 4, 1, 2, 3]], last_sampled=[7]) + assert num_valid == [0] + assert drafts == [[7, 7]] + + +def test_falls_back_to_3gram_when_4gram_missing(): + """No 4-gram match but a 3-gram match exists → propose [4, 1].""" + spec = _make_speculator(min_n=3, max_n=4, k=2) + drafts, num_valid = _propose(spec, [[1, 2, 3, 4, 1, 2, 3]]) + assert num_valid == [2] + assert drafts == [[4, 1]] + + +def test_prefers_longer_ngram(): + """Prefer a 4-gram match over a more recent 3-gram match.""" + spec = _make_speculator(min_n=3, max_n=4, k=2) + drafts, num_valid = _propose( + spec, [[1, 2, 3, 4, 50, 51, 2, 3, 4, 60, 61, 1, 2, 3, 4]] + ) + assert num_valid == [2] + assert drafts == [[50, 51]] + + +def test_picks_longest_match_among_2_3_4_grams(): + """Prefer a 3-gram match over a more recent 2-gram match.""" + spec = _make_speculator(min_n=2, max_n=4, k=2) + drafts, num_valid = _propose(spec, [[2, 3, 4, 50, 51, 3, 4, 60, 61, 1, 2, 3, 4]]) + assert num_valid == [2] + assert drafts == [[50, 51]] + + +@pytest.mark.parametrize("max_model_len", [32, 128, 257, 1025]) +def test_picks_rightmost_when_multiple_matches(max_model_len): + """Pick the last valid match across blocks, ignoring trailing tokens.""" + spec = _make_speculator(min_n=3, max_n=3, k=2, max_model_len=max_model_len) + padding = [0] * (spec.block_l - 5) if spec.n_blocks > 1 else [] + row = [1, 2, 3, 100] + padding + [1, 2, 3, 200, 1, 2, 3, 300, 1, 2, 3] + drafts, num_valid = _propose( + spec, [row + [1, 2, 3, 999, 1, 2, 3]], seq_lens=[len(row)] + ) + assert num_valid == [2] + assert drafts == [[300, 1]] + + +def test_short_context_yields_zero_valid(): + """The only length-2 window overlaps the suffix itself → no match.""" + spec = _make_speculator(min_n=2, max_n=2, k=2) + drafts, num_valid = _propose(spec, [[5, 6]], last_sampled=[99]) + assert num_valid == [0] + assert drafts == [[99, 99]] + + +def test_zero_sampled_disables_proposal(): + """num_sampled==0 disables proposals for that request regardless of match.""" + spec = _make_speculator(min_n=2, max_n=2, k=2) + drafts, num_valid = _propose( + spec, [[1, 2, 3, 1, 2]], num_sampled=[0], last_sampled=[77] + ) + assert num_valid == [0] + assert drafts == [[77, 77]] + + +def test_truncates_num_valid_when_few_tokens_after_match(): + """Fewer than k tokens after the match → num_valid < k, tail falls back. + + Tokens: [1, 2, 1, 2] (seq_len=4). Suffix (1, 2) matches at position 0 + (the match at position 2 is the suffix itself and is excluded). With + k=3, only 2 slots map to tokens inside the context. + """ + spec = _make_speculator(min_n=2, max_n=2, k=3) + drafts, num_valid = _propose(spec, [[1, 2, 1, 2]], last_sampled=[55]) + assert num_valid == [2] + assert drafts[0][:2] == [1, 2] + assert drafts[0][2] == 55 + + +def test_multibatch_mixed(): + """Mixed batch: row 0 matches, row 1 has no match.""" + spec = _make_speculator(min_n=2, max_n=2, k=2) + drafts, num_valid = _propose( + spec, + [[1, 2, 3, 1, 2], [4, 5, 6]], + last_sampled=[10, 20], + ) + assert num_valid == [2, 0] + assert drafts[0] == [3, 1] + assert drafts[1] == [20, 20] + + +def test_multibatch_independent_choice_of_n(): + """Each row independently picks its longest matched n.""" + spec = _make_speculator(min_n=2, max_n=3, k=2) + drafts, num_valid = _propose( + spec, + [ + [9, 1, 2, 3, 8, 1, 2, 3], # 3-gram (1,2,3) at idx 1 → [8, 1] + [7, 1, 2, 9, 1, 2], # 2-gram (1,2) at idx 1 → [9, 1] + ], + ) + assert num_valid == [2, 2] + assert drafts[0] == [8, 1] + assert drafts[1] == [9, 1] + + +def test_min_n_eq_1(): + """min_n=max_n=1 — single-token n-grams always match if context > 1.""" + spec = _make_speculator(min_n=1, max_n=1, k=2) + drafts, num_valid = _propose(spec, [[1, 2, 3, 4, 1]]) + assert num_valid == [2] + assert drafts == [[2, 3]] + + +def test_noncontiguous_idx_mapping(): + """propose() reads token rows in place via idx_mapping (non-contiguous).""" + spec = _make_speculator(min_n=2, max_n=2, k=2) + drafts, num_valid = _propose( + spec, + [[7, 8, 9, 7, 8], [1, 2, 3, 1, 2]], + slots=[3, 0], + ) + assert drafts == [[9, 7], [3, 1]] + assert num_valid == [2, 2] + + +def test_num_valid_written_to_request_slots(): + """num_valid_drafts_for_trim is req-slot indexed for the draft trimmer.""" + spec = _make_speculator(min_n=2, max_n=2, k=2) + _propose( + spec, + [[7, 8, 9, 7, 8], [1, 2, 3, 4, 5]], + slots=[5, 2], + ) + nv = spec.num_valid_drafts_for_trim.cpu() + assert nv[5].item() == 2 # match + assert nv[2].item() == 0 # no match + + +def test_dummy_run_does_not_touch_state(): + """Dummy runs must not mutate persistent request or drafter state.""" + spec = _make_speculator(min_n=2, max_n=2, k=2) + _propose(spec, [[1, 2, 3, 1, 2]], slots=[1]) + before = spec.num_valid_drafts_for_trim.clone() + + input_batch = SimpleNamespace( + num_reqs=1, + idx_mapping=torch.tensor([1], dtype=torch.int64, device=DEVICE), + ) + drafts = spec.propose( + input_batch=input_batch, + attn_metadata=None, + slot_mappings=None, + last_hidden_states=torch.empty(0, device=DEVICE), + aux_hidden_states=None, + num_sampled=torch.ones(1, dtype=torch.int32, device=DEVICE), + num_rejected=torch.zeros(1, dtype=torch.int32, device=DEVICE), + last_sampled=torch.zeros((8, 1), dtype=torch.int64, device=DEVICE), + next_prefill_tokens=torch.zeros(1, dtype=torch.int32, device=DEVICE), + temperature=torch.zeros(1, dtype=torch.float32, device=DEVICE), + seeds=torch.zeros(1, dtype=torch.int64, device=DEVICE), + dp_sync=None, + dummy_run=True, + ) + assert drafts.shape == (1, 2) + assert torch.equal(spec.num_valid_drafts_for_trim.cpu(), before.cpu()) + + +def test_construction_validates_speculative_config(): + spec = _make_speculator(min_n=2, max_n=3, k=2) + assert spec.min_n == 2 + assert spec.max_n == 3 + assert spec.num_speculative_steps == 2 + # Inherited no-op hooks must not raise. + spec.init_cudagraph_manager(None) + spec.capture() + + +# --------------------------------------------------------------------------- +# GPU draft trimming (shared verification-layout machinery) +# --------------------------------------------------------------------------- + + +def test_build_verification_layout_exact_and_gpu_tail(): + """Layout cumsums match a numpy reference; padding tail equals the total.""" + capacities = torch.tensor([2, 0, 1], dtype=torch.int32, device=DEVICE) + non_draft = torch.tensor([1, 5, 1], dtype=torch.int32, device=DEVICE) + num_bonus = 1 + max_num_reqs = 6 + cu_num_logits = torch.empty(max_num_reqs + 1, dtype=torch.int32, device=DEVICE) + qsl = torch.empty(max_num_reqs + 1, dtype=torch.int32, device=DEVICE) + + for num_tokens in (10, None): # exact CPU total vs GPU cumsum tail + cnl, out_qsl = build_verification_layout( + capacities, non_draft, num_bonus, cu_num_logits, qsl, num_tokens + ) + assert cnl.cpu().tolist() == [0, 3, 4, 6] + assert out_qsl.cpu().tolist()[:4] == [0, 3, 8, 10] + # Trailing (padding) entries hold the batch total. + assert out_qsl.cpu().tolist()[4:] == [10, 10, 10] + + +def test_variable_draft_trimmer_clamps_to_num_valid(): + """Scheduled draft slots are clamped per request to the drafter's counts.""" + max_num_reqs = 8 + num_valid_drafts = torch.zeros(max_num_reqs, dtype=torch.int32, device=DEVICE) + num_valid_drafts[4] = 1 # drafter produced 1 valid draft for slot 4 + num_valid_drafts[2] = 3 # more than scheduled for slot 2 + qsl_buf = torch.empty(max_num_reqs + 1, dtype=torch.int32, device=DEVICE) + + trimmer = VariableDraftTrimmer( + num_valid_drafts, + qsl_buf, + num_bonus_tokens=1, + max_num_reqs=max_num_reqs, + max_total_logits=1024, + device=DEVICE, + ) + # Batch: [slot 4 (2 drafts scheduled), slot 2 (2 drafts), slot 0 (prefill)]. + idx_mapping = torch.tensor([4, 2, 0], dtype=torch.int64, device=DEVICE) + num_draft_tokens_per_req = np.array([2, 2, 0], dtype=np.int32) + num_scheduled_tokens = np.array([3, 3, 7], dtype=np.int32) + + cu_num_logits, qsl = trimmer.trim( + idx_mapping, num_draft_tokens_per_req, num_scheduled_tokens + ) + # capacities = min(scheduled, num_valid) = [1, 2, 0] + assert cu_num_logits.cpu().tolist() == [0, 2, 5, 6] + # query lens = non-draft + capacities = [1+1, 1+2, 7+0] + assert qsl.cpu().tolist()[:4] == [0, 2, 5, 12] + # Padding tail equals the (GPU) batch total. + assert qsl.cpu().tolist()[4:] == [12] * (max_num_reqs - 3) + + +@pytest.mark.parametrize("batch_sharded_sampling", [False, True]) +def test_draft_trimmer_disabled_with_batch_sharded_sampling(batch_sharded_sampling): + """The sharder plans from CPU logits boundaries, so GPU trimming must stay off.""" + cfg = _make_vllm_config(min_n=2, max_n=2, k=2) + cfg.parallel_config.enable_batch_sharded_sampling = batch_sharded_sampling + backend = SimpleNamespace( + __name__="FakeBackend", + supports_device_cpu_query_lens_mismatch=lambda: True, + ) + trimmer = maybe_create_draft_trimmer( + vllm_config=cfg, + speculator=SimpleNamespace( + num_valid_drafts_for_trim=torch.zeros(8, dtype=torch.int32, device=DEVICE) + ), + attn_groups=[[SimpleNamespace(backend=backend, layer_names=set())]], + attn_cg_support=AttentionCGSupportInfo(), + req_states=SimpleNamespace(max_num_reqs=8, vocab_size=32, device=DEVICE), + query_start_loc=torch.empty(9, dtype=torch.int32, device=DEVICE), + num_bonus_tokens=1, + ) + assert (trimmer is None) == batch_sharded_sampling + if trimmer is not None: + assert isinstance(trimmer, VariableDraftTrimmer) diff --git a/tests/v1/worker/test_gpu_model_runner_v2.py b/tests/v1/worker/test_gpu_model_runner_v2.py index 86ff1a074538..c029db73983c 100644 --- a/tests/v1/worker/test_gpu_model_runner_v2.py +++ b/tests/v1/worker/test_gpu_model_runner_v2.py @@ -298,3 +298,23 @@ def test_capture_model_profile_only_skips_lock(monkeypatch): runner.capture_model(profile_only=True) assert lock_calls == [] + + +@pytest.mark.parametrize("target_buffer", ["absent", "none", "tensor"]) +def test_get_drafter_hidden_states_tolerates_missing_target_buffer(target_buffer): + """Targets allocate the MTP hidden buffer only for hidden-state drafters.""" + runner = GPUModelRunner.__new__(GPUModelRunner) + hidden_states = torch.zeros(4, 8) + buffer = torch.arange(16 * 8, dtype=torch.float32).view(16, 8) + if target_buffer == "absent": + runner.model = SimpleNamespace() + else: + returned = buffer if target_buffer == "tensor" else None + runner.model = SimpleNamespace(get_mtp_target_hidden_states=lambda: returned) + + out = runner._get_drafter_hidden_states(hidden_states) + + if target_buffer == "tensor": + assert torch.equal(out, buffer[:4]) + else: + assert out is hidden_states diff --git a/tests/v1/worker/test_gpu_rejection_sampler_chunking.py b/tests/v1/worker/test_gpu_rejection_sampler_chunking.py index 2e5adfabd13c..3ab54be3db6f 100644 --- a/tests/v1/worker/test_gpu_rejection_sampler_chunking.py +++ b/tests/v1/worker/test_gpu_rejection_sampler_chunking.py @@ -28,15 +28,19 @@ def test_iter_request_chunks_preserves_request_boundaries(): @pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA") @pytest.mark.parametrize("logprobs_mode", get_args(LogprobsMode)) -def test_chunked_scores_match_full_batch(logprobs_mode: str): +@pytest.mark.parametrize("trim_drafts", [False, True]) +def test_chunked_scores_match_full_batch(logprobs_mode: str, trim_drafts: bool): device = torch.device("cuda") cu_num_logits_np = np.array([0, 3, 4, 8, 10], dtype=np.int32) - num_logits_per_req = np.diff(cu_num_logits_np) + expected_offsets = cu_num_logits_np.copy() + if trim_drafts: + expected_offsets[1:] = [1, 2, 5, 6] + num_logits_per_req = np.diff(expected_offsets) idx_mapping_np = np.array([7, 2, 9, 1], dtype=np.int32) input_batch = SimpleNamespace( num_reqs=4, cu_num_logits_np=cu_num_logits_np, - cu_num_logits=torch.from_numpy(cu_num_logits_np).to(device), + cu_num_logits=torch.from_numpy(expected_offsets).to(device), idx_mapping_np=idx_mapping_np, idx_mapping=torch.from_numpy(idx_mapping_np).to(device), expanded_idx_mapping=torch.from_numpy( @@ -78,7 +82,7 @@ def fake_verify( draft_logits=None, draft_sampled=torch.arange(10, device=device), pos=torch.arange(10, device=device), - max_chunk_logits=5, + max_chunk_logits=10 if trim_drafts else 5, max_num_logprobs=2, ) score_logits = logits + 1 if logprobs_mode in PROCESSED_LOGPROBS_MODES else logits @@ -87,7 +91,7 @@ def fake_verify( num_sampled, score_logits, input_batch.cu_num_logits, - input_batch.cu_num_logits_np, + expected_offsets, max_num_logprobs=2, ) @@ -105,6 +109,7 @@ def fake_verify( full_logprobs.selected_token_ranks, ) assert ( - chunked_logprobs.cu_num_generated_tokens - == full_logprobs.cu_num_generated_tokens + chunked_logprobs.tolists().cu_num_generated_tokens + == full_logprobs.tolists().cu_num_generated_tokens + == expected_offsets.tolist() ) diff --git a/vllm/compilation/backends.py b/vllm/compilation/backends.py index 439a87e1522a..48eeb6abb8cf 100644 --- a/vllm/compilation/backends.py +++ b/vllm/compilation/backends.py @@ -1081,10 +1081,13 @@ def __call__(self, graph: fx.GraphModule, example_inputs: Sequence[Any]) -> Any: # Honors opt-outs such as CompilationMode.NONE or VLLM_DISABLE_COMPILE_CACHE. disable_cache = not is_compile_cache_enabled(self.inductor_config) - # TODO(patchy): ngram gpu kernel will cause vllm torch compile cache errors. + # TODO(patchy): the V1 torch.compile ngram-gpu kernel causes vllm + # torch compile cache errors. The V2 implementation is pure Triton and + # does not need the cache disabled. is_ngram_gpu_enabled = ( vllm_config.speculative_config is not None and vllm_config.speculative_config.use_ngram_gpu() + and not vllm_config.use_v2_model_runner ) disable_cache = disable_cache or is_ngram_gpu_enabled diff --git a/vllm/config/compilation.py b/vllm/config/compilation.py index 2cd5fae77554..3935c560dabc 100644 --- a/vllm/config/compilation.py +++ b/vllm/config/compilation.py @@ -1390,6 +1390,7 @@ def resolve_cudagraph_mode_and_sizes( max_num_reqs: int | None = None, is_profiling: bool = False, piecewise_capture_available: bool = True, + varlen_decode: bool = False, ) -> CUDAGraphMode: from vllm.v1.attention.backend import AttentionCGSupport @@ -1398,6 +1399,23 @@ def resolve_cudagraph_mode_and_sizes( self.cudagraph_mode = CUDAGraphMode.NONE return CUDAGraphMode.NONE + # Decode batches whose per-request query lengths are decided on device + # (adaptive verification, variable-length drafters) are captured as + # varlen decode graphs, which requires a separate decode routine. + # Modes without one would replay such a batch on a mixed graph. + if ( + varlen_decode + and cudagraph_mode.has_full_cudagraphs() + and not cudagraph_mode.separate_routine() + ): + logger.warning( + "CUDAGraphMode.%s cannot capture decode batches with varying " + "per-request query lengths; setting " + "cudagraph_mode=FULL_AND_PIECEWISE", + cudagraph_mode.name, + ) + cudagraph_mode = CUDAGraphMode.FULL_AND_PIECEWISE + # Check cudagraph for mixed batch is supported if ( cudagraph_mode.mixed_mode() == CUDAGraphMode.FULL diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index 2396b26cb01a..caa2c8c24390 100644 --- a/vllm/config/speculative.py +++ b/vllm/config/speculative.py @@ -1941,6 +1941,9 @@ def uses_extract_hidden_states(self) -> bool: def use_ngram_gpu(self) -> bool: return self.method == "ngram_gpu" + def use_ngram(self) -> bool: + return self.method in ("ngram", "ngram_gpu") + def use_multi_module_mtp(self) -> bool: if self.method != "mtp" or self.draft_model_config is None: return False diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index ad8f077148c3..919506ba20c3 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -2830,9 +2830,6 @@ def _get_v2_model_runner_unsupported_features(self) -> list[str]: if speculative_config is not None: if speculative_config.method in ( - # https://github.com/vllm-project/vllm/pull/40704 - "ngram", - "ngram_gpu", # https://github.com/vllm-project/vllm/pull/43091 "draft_model", "suffix", diff --git a/vllm/v1/worker/gpu/input_batch.py b/vllm/v1/worker/gpu/input_batch.py index 5cbd642d2529..5a7e8a2c951a 100644 --- a/vllm/v1/worker/gpu/input_batch.py +++ b/vllm/v1/worker/gpu/input_batch.py @@ -461,6 +461,9 @@ def combine_sampled_and_draft_tokens( cu_num_logits: torch.Tensor, num_logits: int, num_new_sampled_tokens: int = 1, # excl accepted draft tokens, a.k.a bonus tokens + # Set when num_logits is only an upper bound (GPU draft trimming), so + # unwritten trailing entries hold benign in-bounds indices. + zero_init_logits_indices: bool = False, ) -> torch.Tensor: assert num_new_sampled_tokens in (0, 1), ( f"num_new_sampled_tokens must be 0 or 1, got {num_new_sampled_tokens}" @@ -469,7 +472,8 @@ def combine_sampled_and_draft_tokens( num_reqs = idx_mapping.shape[0] num_speculative_steps = draft_tokens.shape[-1] - logits_indices = torch.empty( + alloc = torch.zeros if zero_init_logits_indices else torch.empty + logits_indices = alloc( num_logits, dtype=torch.int64, device=input_ids.device, @@ -703,12 +707,21 @@ def expand_idx_mapping( total_num_logits: int, cu_num_logits: torch.Tensor, max_expand_len: int, + # Set when total_num_logits is only an upper bound (GPU draft trimming), + # so unwritten trailing entries hold benign in-bounds values. + zero_init: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: num_reqs = idx_mapping.shape[0] - expanded_idx_mapping = idx_mapping.new_empty(total_num_logits) - expanded_local_pos = torch.empty( - total_num_logits, dtype=torch.int32, device=idx_mapping.device - ) + if zero_init: + expanded_idx_mapping = idx_mapping.new_zeros(total_num_logits) + expanded_local_pos = torch.zeros( + total_num_logits, dtype=torch.int32, device=idx_mapping.device + ) + else: + expanded_idx_mapping = idx_mapping.new_empty(total_num_logits) + expanded_local_pos = torch.empty( + total_num_logits, dtype=torch.int32, device=idx_mapping.device + ) _expand_idx_mapping_kernel[(num_reqs,)]( idx_mapping, expanded_idx_mapping, diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 761893303a5c..2a8d9842efe4 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -155,7 +155,9 @@ from vllm.v1.worker.gpu.spec_decode import init_speculator from vllm.v1.worker.gpu.spec_decode.adaptive_verification import ( AdaptiveVerificationManager, + VariableDraftTrimmer, maybe_create_adaptive_verification_manager, + maybe_create_draft_trimmer, resolve_adaptive_cudagraph_mode, ) from vllm.v1.worker.gpu.spec_decode.eagle.eagle3_utils import ( @@ -273,13 +275,41 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): self.encoder_cache = EncoderCache() self.ec_connector = get_ec_connector(vllm_config, self.encoder_cache) + self.num_speculative_steps = vllm_config.num_speculative_tokens + + # Multi-module MTP feeds its modules the next num_speculative_steps prefill + # tokens during chunked prefill. Other speculators only read the immediate + # next one. + num_prefill_lookahead = ( + self.num_speculative_steps + if self.speculative_config is not None + and self.speculative_config.use_multi_module_mtp() + else 1 + ) + use_dense_all_token_ids = ( + self.speculative_config is not None and self.speculative_config.use_ngram() + ) + + # General request states. + self.req_states = RequestState( + max_num_reqs=self.max_num_reqs, + max_model_len=self.max_model_len, + max_num_batched_tokens=self.max_num_tokens, + num_speculative_steps=self.num_speculative_steps, + vocab_size=self.vocab_size, + device=self.device, + num_prefill_lookahead=num_prefill_lookahead, + use_dense_all_token_ids=use_dense_all_token_ids, + ) + # Speculative decoding. self.speculator = None self.use_aux_hidden_state_outputs = False - self.num_speculative_steps = vllm_config.num_speculative_tokens if self.speculative_config is not None: if self.is_last_pp_rank: - self.speculator = init_speculator(self.vllm_config, self.device) + self.speculator = init_speculator( + self.vllm_config, self.device, self.req_states + ) if self.speculative_config.method in ( "eagle3", @@ -299,29 +329,9 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): self.is_pooling_model = self.model_config.runner_type == "pooling" self.pooling_runner: PoolingRunner | None = None - # Multi-module MTP feeds its modules the next num_speculative_steps prefill - # tokens during chunked prefill. Other speculators only read the immediate - # next one. - num_prefill_lookahead = ( - self.num_speculative_steps - if self.speculative_config is not None - and self.speculative_config.use_multi_module_mtp() - else 1 - ) - self.step_timing = StepTimingCollector() - - # General request states. - self.req_states = RequestState( - max_num_reqs=self.max_num_reqs, - max_model_len=self.max_model_len, - max_num_batched_tokens=self.max_num_tokens, - num_speculative_steps=self.num_speculative_steps, - vocab_size=self.vocab_size, - device=self.device, - num_prefill_lookahead=num_prefill_lookahead, - ) self.adaptive_verification: AdaptiveVerificationManager | None = None + self.draft_trimmer: VariableDraftTrimmer | None = None self.input_buffers = InputBuffers( max_num_reqs=self.max_num_reqs, max_num_tokens=self.max_num_tokens, @@ -662,6 +672,19 @@ def initialize_kv_cache( target_layer_names=target_attn_layer_names, additional_attn_cg_support=additional_attn_cg_support, ) + # Variable-length drafters (ngram_gpu) trim scheduled draft slots to + # the drafter's valid counts on GPU, when supported. + self.draft_trimmer = None + if self.adaptive_verification is None: + self.draft_trimmer = maybe_create_draft_trimmer( + vllm_config=self.vllm_config, + speculator=self.speculator, + attn_groups=self.attn_groups, + attn_cg_support=attn_cg_support, + req_states=self.req_states, + query_start_loc=self.input_buffers.query_start_loc, + num_bonus_tokens=self.model_state.num_new_sampled_tokens_per_step, + ) self.block_tables = BlockTables( block_sizes=block_sizes, @@ -700,6 +723,9 @@ def initialize_kv_cache( piecewise_capture_available = bool( envs.VLLM_USE_BREAKABLE_CUDAGRAPH or has_compiled_submodule(self.model) ) + varlen_decode = ( + self.adaptive_verification is not None or self.draft_trimmer is not None + ) if self.adaptive_verification is not None: self.compilation_config.cudagraph_mode = resolve_adaptive_cudagraph_mode( self.compilation_config.cudagraph_mode, @@ -715,6 +741,7 @@ def initialize_kv_cache( max_num_reqs=self.max_num_reqs, is_profiling=is_profiling, piecewise_capture_available=piecewise_capture_available, + varlen_decode=varlen_decode, ) self.cudagraph_manager = ModelCudaGraphManager( self.vllm_config, @@ -722,7 +749,7 @@ def initialize_kv_cache( cudagraph_mode, decode_query_len=self.decode_query_len, lora_capture_cases=self.lora_capture_cases, - varlen_decode=self.adaptive_verification is not None, + varlen_decode=varlen_decode, ubatch_runner=self.ubatch_runner, ) if self.cache_config.kv_sharing_fast_prefill and self.pcp_manager is None: @@ -876,14 +903,7 @@ def _dummy_run( ), ) - # 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] + spec_hidden_states = self._get_drafter_hidden_states(hidden_states) if isinstance(self.sampler, GPUWatermarkSampler): self.speculator.prepare_watermarking( self.sampler._get_contexts(input_batch.idx_mapping), @@ -1329,6 +1349,16 @@ def prepare_inputs( adaptive_verification = ( self.adaptive_verification if num_draft_tokens_per_req is not None else None ) + draft_trimmer = None + if ( + adaptive_verification is None + and num_draft_tokens_per_req is not None + and self.draft_trimmer is not None + # The chunked logits path indexes by the CPU (untrimmed) offsets, + # which cannot address the trimmed layout. + and total_num_logits <= self.draft_trimmer.max_total_logits + ): + draft_trimmer = self.draft_trimmer num_scheduled_tokens_upper_bound = num_scheduled_tokens_np if adaptive_verification is not None: # num_scheduled_tokens represents the draft budget evenly distributed across @@ -1358,9 +1388,21 @@ def prepare_inputs( adaptive_verification.reallocate_drafts(req_ids, idx_mapping) ) total_num_logits = num_reqs * num_bonus_tokens + total_num_draft_tokens + elif draft_trimmer is not None: + # Clamp scheduled draft slots to the drafter's valid counts on GPU. + # CPU-side totals remain upper bounds; trimmed gap treated as padding. + cu_num_logits, query_start_loc = draft_trimmer.trim( + idx_mapping, num_draft_tokens_per_req, num_scheduled_tokens_np + ) if draft_tokens: expanded_idx_mapping, expanded_local_pos = expand_idx_mapping( - idx_mapping, total_num_logits, cu_num_logits, self.decode_query_len + idx_mapping, + total_num_logits, + cu_num_logits, + self.decode_query_len, + # With GPU trimming, total_num_logits is an upper bound; the + # gap must hold benign (in-bounds) values. + zero_init=draft_trimmer is not None, ) query_start_loc_np = query_start_loc_np[: num_reqs_padded + 1] query_start_loc = query_start_loc[: num_reqs_padded + 1] @@ -1400,6 +1442,7 @@ def prepare_inputs( cu_num_logits, total_num_logits, self.model_state.num_new_sampled_tokens_per_step, + zero_init_logits_indices=draft_trimmer is not None, ) fast_prefill = None @@ -1461,7 +1504,7 @@ def prepare_inputs( fast_prefill=fast_prefill, max_query_len=( int(num_scheduled_tokens_upper_bound.max()) - if adaptive_verification is not None + if adaptive_verification is not None or draft_trimmer is not None else None ), ) @@ -1990,6 +2033,24 @@ def execute_model( ) return None + def _get_drafter_hidden_states(self, hidden_states: torch.Tensor) -> torch.Tensor: + """Hidden states fed to the drafter. + + Targets such as DeepSeek V4 expose the pre-hc_head residual through + get_mtp_target_hidden_states(). The buffer is sized at + max_num_batched_tokens and only allocated for drafters that consume + target hidden states, so None means "use the regular hidden states". + """ + get_target_hidden_states = getattr( + self.model, "get_mtp_target_hidden_states", None + ) + if get_target_hidden_states is None: + return hidden_states + target_hidden_states = get_target_hidden_states() + if target_hidden_states is None: + return hidden_states + return target_hidden_states[: hidden_states.shape[0]] + @torch.inference_mode() @step_eplb_after() def sample_tokens( @@ -2124,14 +2185,7 @@ def sample_tokens( self.speculator.observe_verification( input_batch.idx_mapping, num_sampled, num_rejected ) - # 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 = draft_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[: draft_hidden_states.size(0)] + spec_hidden_states = self._get_drafter_hidden_states(draft_hidden_states) if isinstance(self.sampler, GPUWatermarkSampler): self.speculator.prepare_watermarking( self.sampler._get_contexts(input_batch.idx_mapping), @@ -2160,11 +2214,9 @@ def sample_tokens( ) 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) + # Spec-decode and diffusion LLMs both use draft tokens. self.draft_tokens_handler.set_draft_tokens( - input_batch, - self.req_states.draft_tokens[input_batch.idx_mapping], + input_batch, self.req_states.draft_tokens[input_batch.idx_mapping] ) if self.pp_handler is not None: self.pp_handler.broadcast_drafts( diff --git a/vllm/v1/worker/gpu/spec_decode/__init__.py b/vllm/v1/worker/gpu/spec_decode/__init__.py index b527c9f21bdf..3720a7533738 100644 --- a/vllm/v1/worker/gpu/spec_decode/__init__.py +++ b/vllm/v1/worker/gpu/spec_decode/__init__.py @@ -1,11 +1,21 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from typing import TYPE_CHECKING + import torch from vllm.config import VllmConfig +if TYPE_CHECKING: + from vllm.v1.worker.gpu.states import RequestState + -def init_speculator(vllm_config: VllmConfig, device: torch.device): +def init_speculator( + vllm_config: VllmConfig, + device: torch.device, + req_states: "RequestState", +): + """Build the speculator for this config.""" speculative_config = vllm_config.speculative_config assert speculative_config is not None if speculative_config.method == "extract_hidden_states": @@ -54,5 +64,11 @@ def init_speculator(vllm_config: VllmConfig, device: torch.device): ) return EagleSpeculator(vllm_config, device) + elif speculative_config.use_ngram(): + from vllm.v1.worker.gpu.spec_decode.ngram.speculator import ( + NgramGPUSpeculator, + ) + + return NgramGPUSpeculator(vllm_config, device, req_states) else: raise NotImplementedError(f"{speculative_config.method} is not supported yet.") diff --git a/vllm/v1/worker/gpu/spec_decode/adaptive_verification.py b/vllm/v1/worker/gpu/spec_decode/adaptive_verification.py index af58b8993bfa..1840a4eda7c7 100644 --- a/vllm/v1/worker/gpu/spec_decode/adaptive_verification.py +++ b/vllm/v1/worker/gpu/spec_decode/adaptive_verification.py @@ -30,6 +30,7 @@ from vllm.config import VllmConfig from vllm.v1.worker.gpu.attn_utils import AttentionCGSupportInfo from vllm.v1.worker.gpu.input_batch import InputBatch + from vllm.v1.worker.gpu.spec_decode.speculator import BaseSpeculator from vllm.v1.worker.gpu.states import RequestState from vllm.v1.worker.utils import AttentionGroup @@ -78,6 +79,161 @@ def _assign_draft_token_budget( ) +def build_verification_layout( + capacities: torch.Tensor, + num_non_draft_tokens: torch.Tensor, + num_bonus_tokens: int, + cu_num_logits: torch.Tensor, + query_start_loc: torch.Tensor, + num_tokens: int | None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Build GPU cu_num_logits / query_start_loc from per-request admitted + draft counts. + + Trailing (padding) query_start_loc entries are filled with the batch + total: the exact CPU value when known (`num_tokens`), otherwise the GPU + cumsum tail, so downstream kernels treat everything past the real tokens + as padding. + """ + num_reqs = capacities.shape[0] + cu_num_logits[:1].zero_() + torch.cumsum( + capacities + num_bonus_tokens, dim=0, out=cu_num_logits[1 : num_reqs + 1] + ) + query_start_loc[:1].zero_() + torch.cumsum( + capacities + num_non_draft_tokens, dim=0, out=query_start_loc[1 : num_reqs + 1] + ) + tail = num_tokens if num_tokens is not None else query_start_loc[num_reqs] + query_start_loc[num_reqs + 1 :] = tail + return cu_num_logits[: num_reqs + 1], query_start_loc + + +class VariableDraftTrimmer: + """GPU-side verification trimming for variable-length drafters (ngram). + + The drafter records per-request valid draft counts on GPU in + `num_valid_drafts_for_trim`. The scheduler still schedules the full + num_speculative_tokens per request; at the next step this trimmer clamps + each request's scheduled draft slots to the recorded count and rebuilds + cu_num_logits / query_start_loc on device, so the CPU keeps only upper + bounds. Trimmed slots surface as ordinary rejections through the + existing num_rejected accounting — no scheduler round-trip and no + CPU<->GPU synchronization. + """ + + def __init__( + self, + num_valid_drafts: torch.Tensor, + query_start_loc: torch.Tensor, + num_bonus_tokens: int, + max_num_reqs: int, + max_total_logits: int, + device: torch.device, + ): + self.num_valid_drafts = num_valid_drafts + self.query_start_loc = query_start_loc + self.num_bonus_tokens = num_bonus_tokens + # Rejection sampling chunks logits by the CPU (untrimmed) offsets, + # which cannot address the compacted layout; skip trimming for + # batches that would not fit in one chunk. + self.max_total_logits = max_total_logits + self._capacities = torch.empty(max_num_reqs, dtype=torch.int32, device=device) + self._num_non_draft_tokens = torch.empty( + max_num_reqs, dtype=torch.int32, device=device + ) + self._cu_num_logits = torch.empty( + max_num_reqs + 1, dtype=torch.int32, device=device + ) + + def trim( + self, + idx_mapping: torch.Tensor, + num_draft_tokens_per_req: np.ndarray, + num_scheduled_tokens_np: np.ndarray, + ) -> tuple[torch.Tensor, torch.Tensor]: + num_reqs = idx_mapping.shape[0] + capacities = self._capacities[:num_reqs] + async_copy_to_gpu(num_draft_tokens_per_req, out=capacities) + torch.minimum(capacities, self.num_valid_drafts[idx_mapping], out=capacities) + num_non_draft_tokens = self._num_non_draft_tokens[:num_reqs] + async_copy_to_gpu( + num_scheduled_tokens_np - num_draft_tokens_per_req, + out=num_non_draft_tokens, + ) + return build_verification_layout( + capacities, + num_non_draft_tokens, + self.num_bonus_tokens, + self._cu_num_logits, + self.query_start_loc, + num_tokens=None, + ) + + +def maybe_create_draft_trimmer( + *, + vllm_config: "VllmConfig", + speculator: "BaseSpeculator | None", + attn_groups: list[list["AttentionGroup"]], + attn_cg_support: "AttentionCGSupportInfo", + req_states: "RequestState", + query_start_loc: torch.Tensor, + num_bonus_tokens: int, +) -> VariableDraftTrimmer | None: + """Create a VariableDraftTrimmer when the drafter and environment support + GPU-side trimming; otherwise fall back (with a log) to verifying the full + padded drafts, which is correct but wastes verification compute.""" + from vllm.v1.worker.gpu.spec_decode.rejection_sampler import get_max_chunk_logits + + if speculator is None or speculator.num_valid_drafts_for_trim is None: + return None + + parallel_config = vllm_config.parallel_config + cudagraph_mode = vllm_config.compilation_config.cudagraph_mode + + reason = None + backend = get_query_lens_mismatch_unsupported_backend(attn_groups) + if backend is not None: + reason = f"the {backend} attention backend" + elif ( + cudagraph_mode.has_full_cudagraphs() + and attn_cg_support.min_cg_support != AttentionCGSupport.ALWAYS + ): + reason = f"varlen decode cudagraphs with {attn_cg_support.min_cg_attn_backend}" + elif vllm_config.lora_config is not None: + reason = "LoRA" + elif parallel_config.enable_batch_sharded_sampling: + # The sharder plans its all-to-all splits and local buffers from the + # CPU (untrimmed) logits boundaries. + reason = "batch-sharded sampling" + elif parallel_config.pipeline_parallel_size > 1: + reason = "pipeline parallelism" + elif ( + parallel_config.decode_context_parallel_size > 1 + or parallel_config.prefill_context_parallel_size > 1 + ): + reason = "context parallelism" + + if reason is not None: + logger.info( + "GPU draft trimming is not supported with %s; invalid draft " + "slots will be verified (and rejected) instead of trimmed.", + reason, + ) + return None + + logger.info("GPU draft trimming enabled for variable-length drafts.") + return VariableDraftTrimmer( + speculator.num_valid_drafts_for_trim, + query_start_loc, + num_bonus_tokens, + req_states.max_num_reqs, + get_max_chunk_logits(req_states.vocab_size), + req_states.device, + ) + + def build_cost_tables_from_curves( draft_curve: list[tuple[int, float]], verify_curve: list[tuple[int, float]], @@ -428,24 +584,15 @@ def reallocate_drafts( num_non_draft_tokens, out=num_non_draft_tokens_gpu, ) - self._cu_num_logits[:1].zero_() - torch.cumsum( - capacities + self.num_bonus_tokens, - dim=0, - out=self._cu_num_logits[1 : num_reqs + 1], - ) - self.query_start_loc[:1].zero_() - torch.cumsum( - capacities + num_non_draft_tokens_gpu, - dim=0, - out=self.query_start_loc[1 : num_reqs + 1], - ) - self.query_start_loc[num_reqs + 1 :].fill_(num_tokens) - return ( - self._cu_num_logits[: num_reqs + 1], + cu_num_logits, query_start_loc = build_verification_layout( + capacities, + num_non_draft_tokens_gpu, + self.num_bonus_tokens, + self._cu_num_logits, self.query_start_loc, - draft_budget, + num_tokens, ) + return cu_num_logits, query_start_loc, draft_budget def maybe_create_adaptive_verification_manager( diff --git a/vllm/v1/worker/gpu/spec_decode/ngram/__init__.py b/vllm/v1/worker/gpu/spec_decode/ngram/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/vllm/v1/worker/gpu/spec_decode/ngram/speculator.py b/vllm/v1/worker/gpu/spec_decode/ngram/speculator.py new file mode 100644 index 000000000000..3fe3c392bd8a --- /dev/null +++ b/vllm/v1/worker/gpu/spec_decode/ngram/speculator.py @@ -0,0 +1,282 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +import torch + +from vllm.config import VllmConfig +from vllm.triton_utils import HAS_TRITON, tl, triton +from vllm.v1.worker.gpu.input_batch import InputBatch +from vllm.v1.worker.gpu.spec_decode.speculator import BaseSpeculator + +if TYPE_CHECKING: + from vllm.v1.worker.gpu.dp_utils import DPSyncState + from vllm.v1.worker.gpu.states import RequestState + + +@triton.jit +def _ngram_scan_kernel( + token_ids_ptr, # *int32 [max_num_reqs, token_ids_stride] + token_ids_stride, + idx_mapping_ptr, # *int64 [B] batch_idx -> req_state_idx + total_len_ptr, # *int32 [max_num_reqs] + num_sampled_ptr, # *int32 [B] + scratch_ptr, # *int64 [B, scratch_stride] (output) + scratch_stride, + L, # int64 scalar (= max_model_len) + MIN_N: tl.constexpr, + MAX_N: tl.constexpr, + MAX_N_PO2: tl.constexpr, + BLOCK_L: tl.constexpr, +): + b = tl.program_id(0).to(tl.int64) + blk = tl.program_id(1).to(tl.int64) + Lp1 = tl.cast(L, tl.int64) + 1 + + req_state_idx = tl.load(idx_mapping_ptr + b).to(tl.int64) + seq_len = tl.load(total_len_ptr + req_state_idx).to(tl.int64) + num_sampled = tl.load(num_sampled_ptr + b) + eligible_row = (num_sampled > 0) & (seq_len >= MIN_N) + + scratch_off = b * scratch_stride + blk + + # Ineligible rows, and blocks fully past the last candidate match + # position, write 0 and exit. + if not (eligible_row & (blk * BLOCK_L <= seq_len - MIN_N - 1)): + tl.store(scratch_ptr + scratch_off, tl.zeros((), tl.int64)) + return + + row_off = req_state_idx * token_ids_stride + + # Load the length-MAX_N suffix once into registers. + suf_iota = tl.arange(0, MAX_N_PO2).to(tl.int64) + suf_pos = seq_len - MAX_N + suf_iota + suf_in_range = (suf_iota < MAX_N) & (suf_pos >= 0) & (suf_pos < seq_len) + suffix = tl.load( + token_ids_ptr + row_off + suf_pos, + mask=suf_in_range, + other=-1, + ).to(tl.int32) + + pos_iota = tl.arange(0, BLOCK_L).to(tl.int64) + pos = blk * BLOCK_L + pos_iota # ascending + + best_score = tl.zeros([BLOCK_L], dtype=tl.int64) + + for n_iter in tl.static_range(MIN_N, MAX_N + 1): + max_pos_n = seq_len - n_iter - 1 + match = (pos >= 0) & (pos <= max_pos_n) + for j in tl.static_range(0, n_iter): + tok = tl.load( + token_ids_ptr + row_off + (pos + j), + mask=match, + other=0, + ).to(tl.int32) + suf_idx = (MAX_N - n_iter) + j + suf_val = tl.sum(tl.where(suf_iota == suf_idx, suffix, 0)) + match = match & (tok == suf_val) + + # Pack (n, pos) so a single max yields longest-n, rightmost-pos. + cand = n_iter * Lp1 + pos + 1 + best_score = tl.where(match, cand, best_score) + + block_best = tl.max(best_score, axis=0) + tl.store(scratch_ptr + scratch_off, block_best) + + +@triton.jit +def _ngram_finalize_kernel( + token_ids_ptr, # *int32 [max_num_reqs, token_ids_stride] + token_ids_stride, + idx_mapping_ptr, # *int64 [B] + total_len_ptr, # *int32 [max_num_reqs] + num_sampled_ptr, # *int32 [B] + last_sampled_ptr, # *int64 [max_num_reqs] + scratch_ptr, # *int64 [B, scratch_stride] + scratch_stride, + drafts_ptr, # *int64 [B, K] (output, batch indexed) + num_valid_ptr, # *int32 [max_num_reqs] (output, req-slot indexed) + L, + N_BLOCKS, + K: tl.constexpr, + K_PO2: tl.constexpr, + N_BLOCKS_PO2: tl.constexpr, +): + b = tl.program_id(0).to(tl.int64) + Lp1 = tl.cast(L, tl.int64) + 1 + NB = tl.cast(N_BLOCKS, tl.int64) + + req_state_idx = tl.load(idx_mapping_ptr + b).to(tl.int64) + + nb_iota = tl.arange(0, N_BLOCKS_PO2).to(tl.int64) + nb_in_range = nb_iota < NB + block_scores = tl.load( + scratch_ptr + b * scratch_stride + nb_iota, + mask=nb_in_range, + other=0, + ) + score = tl.max(block_scores, axis=0) + + seq_len = tl.load(total_len_ptr + req_state_idx).to(tl.int64) + num_sampled = tl.load(num_sampled_ptr + b) + last_tok = tl.load(last_sampled_ptr + req_state_idx) + + has_match = score > 0 + s1 = score - 1 + best_n = tl.where(has_match, s1 // Lp1, tl.zeros_like(s1)) + best_pos = tl.where(has_match, s1 - best_n * Lp1, tl.zeros_like(s1)) + draft_start = tl.where(has_match, best_pos + best_n, tl.zeros_like(s1)) + + tokens_avail = tl.maximum(seq_len - draft_start, 0) + write_ok = (num_sampled > 0) & has_match + nv = tl.where(write_ok, tl.minimum(tl.cast(K, tl.int64), tokens_avail), 0) + tl.store(num_valid_ptr + req_state_idx, nv.to(tl.int32)) + + row_off = req_state_idx * token_ids_stride + k_iota = tl.arange(0, K_PO2).to(tl.int64) + k_in_range = k_iota < K + gather_idx = tl.minimum(draft_start + k_iota, tl.cast(L, tl.int64) - 1) + slot_valid = (k_iota < tokens_avail) & write_ok & k_in_range + gathered = tl.load( + token_ids_ptr + row_off + gather_idx, + mask=slot_valid, + other=0, + ).to(tl.int64) + # Invalid slots fall back to the last sampled token; they are either + # trimmed from the verification batch on GPU or verified as ordinary + # (rejectable) drafts, so the fill value only affects efficiency. + out = tl.where(slot_valid, gathered, last_tok) + tl.store(drafts_ptr + b * K + k_iota, out, mask=k_in_range) + + +class NgramGPUSpeculator(BaseSpeculator): + """V2-compatible GPU n-gram speculator.""" + + supports_mm_inputs = False + draft_logits = None + + def __init__( + self, + vllm_config: VllmConfig, + device: torch.device, + req_states: RequestState, + ): + if not HAS_TRITON: + raise RuntimeError("ngram_gpu speculative decoding requires Triton.") + spec = vllm_config.speculative_config + assert spec is not None + assert spec.prompt_lookup_min is not None, ( + "prompt_lookup_min must be configured for ngram_gpu" + ) + assert spec.prompt_lookup_max is not None, ( + "prompt_lookup_max must be configured for ngram_gpu" + ) + assert 1 <= spec.prompt_lookup_min <= spec.prompt_lookup_max + + self.vllm_config = vllm_config + self.device = device + self.req_states = req_states + self.speculative_config = spec + self.num_speculative_steps: int = spec.num_speculative_tokens + + self.min_n: int = spec.prompt_lookup_min + self.max_n: int = spec.prompt_lookup_max + + self.max_num_reqs: int = vllm_config.scheduler_config.max_num_seqs + self.max_model_len: int = vllm_config.model_config.max_model_len + + L = self.max_model_len + if L >= 1024: + self.block_l = 256 + elif L >= 256: + self.block_l = 128 + elif L >= 64: + self.block_l = 64 + else: + self.block_l = max(16, triton.next_power_of_2(max(L, 1))) + self.n_blocks = triton.cdiv(L, self.block_l) + + self.scratch = torch.zeros( + (self.max_num_reqs, self.n_blocks), dtype=torch.int64, device=device + ) + # Per request-slot count of usable drafts from the latest proposal, + # consumed by the model runner's GPU draft trimmer. + self.num_valid_drafts_for_trim = torch.zeros( + self.max_num_reqs, dtype=torch.int32, device=device + ) + # Batch-ordered draft output, scattered into RequestState.draft_tokens + # by the model runner (same contract as the model-based speculators). + self.drafts = torch.zeros( + (self.max_num_reqs, self.num_speculative_steps), + dtype=torch.int64, + device=device, + ) + + @torch.inference_mode() + def propose( + self, + input_batch: InputBatch, + attn_metadata: Any, + slot_mappings: Any, + last_hidden_states: torch.Tensor, + aux_hidden_states: list[torch.Tensor] | None, + num_sampled: torch.Tensor, + num_rejected: torch.Tensor, + last_sampled: torch.Tensor, + next_prefill_tokens: torch.Tensor, + temperature: torch.Tensor, + seeds: torch.Tensor, + dp_sync: DPSyncState | None = None, + dummy_run: bool = False, + skip_attn_for_dummy_run: bool = False, + mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None, + is_profile: bool = False, + ) -> torch.Tensor: + num_reqs = input_batch.num_reqs + if dummy_run: + # No persistent request state may be touched during dummy runs. + return self.drafts[:num_reqs] + + req_states = self.req_states + token_ids = req_states.all_token_ids.gpu + idx_mapping = input_batch.idx_mapping + + _ngram_scan_kernel[(num_reqs, self.n_blocks)]( + token_ids, + token_ids.stride(0), + idx_mapping, + req_states.total_len.gpu, + num_sampled, + self.scratch, + self.scratch.stride(0), + self.max_model_len, + self.min_n, + self.max_n, + max(1, triton.next_power_of_2(self.max_n)), + self.block_l, + num_warps=4, + num_stages=2, + ) + + _ngram_finalize_kernel[(num_reqs,)]( + token_ids, + token_ids.stride(0), + idx_mapping, + req_states.total_len.gpu, + num_sampled, + last_sampled.view(-1), + self.scratch, + self.scratch.stride(0), + self.drafts, + self.num_valid_drafts_for_trim, + self.max_model_len, + self.n_blocks, + self.num_speculative_steps, + max(1, triton.next_power_of_2(self.num_speculative_steps)), + max(1, triton.next_power_of_2(self.n_blocks)), + num_warps=2, + num_stages=1, + ) + return self.drafts[:num_reqs] diff --git a/vllm/v1/worker/gpu/spec_decode/rejection_sampler.py b/vllm/v1/worker/gpu/spec_decode/rejection_sampler.py index 609924e87fd2..3d20ebd21573 100644 --- a/vllm/v1/worker/gpu/spec_decode/rejection_sampler.py +++ b/vllm/v1/worker/gpu/spec_decode/rejection_sampler.py @@ -110,7 +110,7 @@ def _get_logprobs_tensors( num_sampled: torch.Tensor, logits: torch.Tensor, cu_num_logits: torch.Tensor, - cu_num_logits_np: np.ndarray, + cu_num_logits_np: np.ndarray | None, max_num_logprobs: int, ) -> LogprobsTensors | None: if max_num_logprobs == NO_LOGPROBS: @@ -132,10 +132,7 @@ def _get_logprobs_tensors( expanded_logits = num_logits != num_reqs cu_num_generated_tokens: list[int] | torch.Tensor | None = None if expanded_logits: - if self.enable_adaptive_verification: - # Adaptive verification keeps the true per-request boundaries - # on device only; cu_num_logits_np holds the pre-compacted - # layout. + if cu_num_logits_np is None: cu_num_generated_tokens = cu_num_logits.clone() else: cu_num_generated_tokens = cu_num_logits_np.tolist() @@ -222,10 +219,8 @@ def _verify_in_chunks( num_reqs = input_batch.num_reqs if logits.shape[0] <= max_chunk_logits: - # One chunk covers the batch. Adaptive verification compacts the logits - # without updating cu_num_logits_np (it keeps the pre-compacted layout), - # so the stale sums must not pick chunk boundaries; its budget cap - # guarantees the compacted batch always lands here. + # GPU trimming can leave the CPU boundaries stale. + # Trimmed batches fit in one chunk. request_chunks: Iterable[tuple[int, int]] = ((0, num_reqs),) else: assert not self.enable_adaptive_verification @@ -257,7 +252,11 @@ def _verify_in_chunks( num_sampled, processed_logits if use_processed_logits else logits[lo:hi], chunk_cu_num_logits, - chunk_cu_num_logits_np, + ( + chunk_cu_num_logits_np + if logits.shape[0] > max_chunk_logits + else None + ), max_num_logprobs, ) if chunk_logprobs is not None: diff --git a/vllm/v1/worker/gpu/spec_decode/speculator.py b/vllm/v1/worker/gpu/spec_decode/speculator.py index 299cac641620..f75645dc69a2 100644 --- a/vllm/v1/worker/gpu/spec_decode/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/speculator.py @@ -58,13 +58,17 @@ def _target_feeds_hc_residual(vllm_config: VllmConfig) -> bool: class BaseSpeculator(ABC): - @abstractmethod + # Variable-length drafters publish per-request counts of usable drafts + # here, [max_num_reqs] int32 indexed by request slot. Leaving it None + # means every scheduled draft is verified; setting it opts the drafter + # into device-side trimming by the model runner's draft trimmer. + num_valid_drafts_for_trim: torch.Tensor | None = None + def init_cudagraph_manager(self, cudagraph_mode: CUDAGraphMode) -> None: - pass + return None - @abstractmethod def capture(self) -> None: - pass + return None @abstractmethod def propose( diff --git a/vllm/v1/worker/gpu/states.py b/vllm/v1/worker/gpu/states.py index 025c8300fcd2..f56cd6895126 100644 --- a/vllm/v1/worker/gpu/states.py +++ b/vllm/v1/worker/gpu/states.py @@ -16,6 +16,7 @@ def __init__( vocab_size: int, device: torch.device, num_prefill_lookahead: int = 1, + use_dense_all_token_ids: bool = False, ): self.max_num_reqs = max_num_reqs self.max_model_len = max_model_len @@ -30,12 +31,14 @@ def __init__( # NOTE(woosuk): This tensor can be extremely large (e.g., several GBs) # depending on the configured max_num_reqs and max_model_len. - # To save GPU memory, we use UVA instead of GPU for this tensor. + # To save GPU memory, we use UVA instead of GPU by default, but + # ngram_gpu benefits from dense device residency because it scans + # active rows repeatedly during proposal. self.all_token_ids = StagedWriteTensor( (self.max_num_reqs, self.max_model_len), dtype=torch.int32, device=device, - uva_instead_of_gpu=True, + uva_instead_of_gpu=not use_dense_all_token_ids, ) # NOTE(woosuk): Distinguish clearly between prompt_len and prefill_len: # - prompt_len: Number of tokens in the user-provided prompt.