diff --git a/tests/v1/spec_decode/test_lilicorr.py b/tests/v1/spec_decode/test_lilicorr.py index 5a42f3826a..46dbaf41af 100644 --- a/tests/v1/spec_decode/test_lilicorr.py +++ b/tests/v1/spec_decode/test_lilicorr.py @@ -290,18 +290,49 @@ def test_context_anchor_is_last_committed_normalized_feature(): anchor_hidden=torch.full((4, 4), 123.0), anchor_valid=torch.ones(4, dtype=torch.bool), ) - batch = SimpleNamespace(num_reqs=2, query_start_loc=torch.tensor([0, 2, 6])) - LiLiCorrSpeculator.prepare_context_anchor(spec, batch, torch.tensor([0, 2])) + LiLiCorrSpeculator.prepare_context_anchor( + spec, 2, torch.tensor([0, 2, 6]), torch.tensor([0, 2]) + ) torch.testing.assert_close(spec.anchor_hidden[:2], norm(hidden[[1, 3]])) assert spec.anchor_valid.tolist() == [True, True, False, False] assert not spec.anchor_hidden[2:].any() # A smaller reordered batch must not retain an old request's anchor. - batch = SimpleNamespace(num_reqs=1, query_start_loc=torch.tensor([0, 3])) - LiLiCorrSpeculator.prepare_context_anchor(spec, batch, torch.tensor([1])) + LiLiCorrSpeculator.prepare_context_anchor( + spec, 1, torch.tensor([0, 3]), torch.tensor([1]) + ) torch.testing.assert_close(spec.anchor_hidden[0], norm(hidden[1])) assert not spec.anchor_hidden[1:].any() +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_captured_context_anchor_rereads_fixed_inputs(): + """The anchor runs in the draft's FULL graph: each replay must re-read its + fixed-address inputs and ignore a padded request's stale rejected count.""" + dev = torch.device("cuda") + query_start_loc = torch.zeros(3, dtype=torch.int32, device=dev) + num_rejected = torch.zeros(2, dtype=torch.int32, device=dev) + spec = SimpleNamespace( + hidden_states=torch.randn(8, 2, device=dev), + model=SimpleNamespace(model=SimpleNamespace(hidden_norm=lambda x: 2 * x)), + anchor_hidden=torch.zeros(2, 2, device=dev), + anchor_valid=torch.zeros(2, dtype=torch.bool, device=dev), + ) + args = (spec, 2, query_start_loc, num_rejected) + LiLiCorrSpeculator.prepare_context_anchor(*args) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + LiLiCorrSpeculator.prepare_context_anchor(*args) + # One live request padded to two; the second case anchors on a prefill row + # past the graph's context rows, stored eagerly before replay. + for loc, rejected, row in (([0, 2, 2], [1, 5], 0), ([0, 7, 7], [0, 0], 6)): + query_start_loc.copy_(torch.tensor(loc)) + num_rejected.copy_(torch.tensor(rejected)) + graph.replay() + torch.testing.assert_close(spec.anchor_hidden[0], 2 * spec.hidden_states[row]) + assert spec.anchor_valid.tolist() == [True, False] + assert not spec.anchor_hidden[1].any() + + @pytest.mark.parametrize("lilicorr", [False, True]) def test_candidate_generation_routes_scores_and_adaptive_inputs(monkeypatch, lilicorr): """Both heads share the walk, including a padded request and adaptive scores.""" diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py index 790193dac5..a75e84100e 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py @@ -25,6 +25,9 @@ from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers from vllm.v1.worker.gpu.model_states.interface import ModelState from vllm.v1.worker.gpu.spec_decode.dflash.cudagraph import DFlashCudaGraphManager from vllm.v1.worker.gpu.spec_decode.dflash.utils import load_dflash_model +from vllm.v1.worker.gpu.spec_decode.eagle.eagle3_utils import ( + get_eagle3_aux_layers_from_config, +) from vllm.v1.worker.gpu.spec_decode.speculator import DraftModelSpeculator from vllm.v1.worker.gpu.spec_decode.utils import get_parallel_drafting_token_id from vllm.v1.worker.utils import AttentionGroup @@ -90,6 +93,20 @@ class DFlashSpeculator(DraftModelSpeculator): self.max_num_tokens, dtype=torch.int64, device=device ) + # Fixed-address inputs of the FULL graph's context step: the target's aux + # hidden states and, for an anchor that reads them, the rejected counts. + self._aux_staging: torch.Tensor | None = None + self._num_rejected: torch.Tensor | None = None + aux_layers = get_eagle3_aux_layers_from_config(speculative_config) + cudagraph_mode = vllm_config.compilation_config.cudagraph_mode + if aux_layers and cudagraph_mode.decode_mode() == CUDAGraphMode.FULL: + self._aux_staging = torch.zeros( + self._num_graph_context_tokens(self.max_num_reqs), + len(aux_layers) * vllm_config.model_config.get_hidden_size(), + dtype=self.dtype, + device=device, + ) + # Per-mask-token sampling buffers. Flattened from (num_reqs, num_spec_tokens). max_num_sampled_tokens = self.max_num_reqs * self.num_speculative_steps self.sample_indices = torch.zeros( @@ -173,11 +190,52 @@ class DFlashSpeculator(DraftModelSpeculator): def _prepare_graph_forward( self, num_reqs: int, attn_metadata: dict[str, Any] | None ) -> None: - self._precompute_context_kv(0, self._num_graph_context_tokens(num_reqs)) + """The FULL graph's context step: combine the staged target hidden + states, publish the context anchor and store the context K/V.""" + num_context = self._num_graph_context_tokens(num_reqs) + if self._aux_staging is not None: + self.hidden_states[:num_context].copy_( + self.model.combine_hidden_states(self._aux_staging[:num_context]) + ) + self.prepare_context_anchor( + num_reqs, self.target_input_buffers.query_start_loc, self._num_rejected + ) + self._precompute_context_kv(0, num_context) # Recorded in the FULL graph so replays skip the eager metadata build. if attn_metadata is not None and self.supports_decode_metadata_update: self._update_draft_decode_metadata(attn_metadata, num_reqs) + def _prepare_graph_replay( + self, + input_batch: InputBatch, + num_graph_reqs: int, + last_hidden_states: torch.Tensor, + aux_hidden_states: list[torch.Tensor] | None, + num_rejected: torch.Tensor, + dummy_run: bool, + ) -> None: + """Eager work before a FULL replay: stage the inputs _prepare_graph_forward + reads and store the context rows past the graph's num_context rows.""" + num_reqs = input_batch.num_reqs + num_target_tokens = input_batch.num_tokens + num_context = self._num_graph_context_tokens(num_graph_reqs) + self._store_hidden_states( + num_target_tokens, last_hidden_states, aux_hidden_states, num_context + ) + if self._num_rejected is not None: + self._num_rejected[:num_reqs].copy_(num_rejected[:num_reqs]) + if dummy_run: + # Dummy block tables are placeholders: write no context K/V. + self._context_slot_mappings[:, :num_context].fill_(PAD_SLOT_ID) + elif num_target_tokens <= num_context: + # Rows past the batch keep stale positions but write no K/V. + self._context_slot_mappings[:, num_target_tokens:num_context].fill_( + PAD_SLOT_ID + ) + else: + # Prefill context beyond the graph's rows is stored before replay. + self._precompute_context_kv(num_context, num_target_tokens) + def load_draft_model( self, target_model: nn.Module, @@ -335,10 +393,42 @@ class DFlashSpeculator(DraftModelSpeculator): context_slots, ) - def prepare_context_anchor( - self, input_batch: InputBatch, num_rejected: torch.Tensor + def _store_hidden_states( + self, + num_tokens: int, + last_hidden_states: torch.Tensor, + aux_hidden_states: list[torch.Tensor] | None, + num_graph_tokens: int = 0, ) -> None: - """Publish context features required by a draft's candidate head.""" + # A FULL graph combines its first num_graph_tokens rows from the staged + # aux hidden states; the remaining rows are combined here. + start = 0 + if num_graph_tokens and self._aux_staging is not None: + assert aux_hidden_states + start = min(num_tokens, num_graph_tokens) + torch.cat( + [h[:start] for h in aux_hidden_states], + dim=-1, + out=self._aux_staging[:start], + ) + if start == num_tokens: + return + if aux_hidden_states: + hidden_states = self.model.combine_hidden_states( + torch.cat([h[start:num_tokens] for h in aux_hidden_states], dim=-1) + ) + else: + hidden_states = last_hidden_states[start:num_tokens] + self.hidden_states[start:num_tokens].copy_(hidden_states) + + def prepare_context_anchor( + self, + num_reqs: int, + query_start_loc: torch.Tensor, + num_rejected: torch.Tensor | None, + ) -> None: + """Publish context features required by a draft's candidate head. Runs + inside the FULL graph, so it must not sync with the host.""" @torch.inference_mode() def propose( @@ -381,19 +471,16 @@ class DFlashSpeculator(DraftModelSpeculator): # number of rejected tokens, we maintain the size of input_ids and # hidden_states the same as the target model's. This means, we pad each # request's query length to include any rejected positions. - if aux_hidden_states: - hidden_states = self.model.combine_hidden_states( - torch.cat(aux_hidden_states, dim=-1) - ) - else: - hidden_states = last_hidden_states - self.hidden_states[:num_target_tokens].copy_(hidden_states[:num_target_tokens]) - self.prepare_context_anchor(input_batch, num_rejected) - if dummy_run and skip_attn_for_dummy_run: # Memory profiling path: block_tables / kv_cache_config are not initialized. # Since DFlash needs to build its own attention metadata, we must skip the # preparation in this path and run a minimal forward pass. + self._store_hidden_states( + num_target_tokens, last_hidden_states, aux_hidden_states + ) + self.prepare_context_anchor( + num_reqs, input_batch.query_start_loc, num_rejected + ) self.model.precompute_and_store_context_kv( self.hidden_states[:num_target_tokens], self.context_positions[:num_target_tokens], @@ -476,21 +563,22 @@ class DFlashSpeculator(DraftModelSpeculator): ) if batch_desc.cg_mode == CUDAGraphMode.FULL: - # The graph stores the first num_context context rows. assert batch_desc.num_reqs is not None - num_context = self._num_graph_context_tokens(batch_desc.num_reqs) - if dummy_run: - # Dummy block tables are placeholders: write no context K/V. - self._context_slot_mappings[:, :num_context].fill_(PAD_SLOT_ID) - elif num_target_tokens <= num_context: - # Rows past the batch keep stale positions but write no K/V. - self._context_slot_mappings[:, num_target_tokens:num_context].fill_( - PAD_SLOT_ID - ) - else: - # Prefill context beyond the graph's rows is stored before replay. - self._precompute_context_kv(num_context, num_target_tokens) + self._prepare_graph_replay( + input_batch, + batch_desc.num_reqs, + last_hidden_states, + aux_hidden_states, + num_rejected, + dummy_run, + ) else: + self._store_hidden_states( + num_target_tokens, last_hidden_states, aux_hidden_states + ) + self.prepare_context_anchor( + num_reqs, input_batch.query_start_loc, num_rejected + ) self._precompute_context_kv(0, num_target_tokens, dummy_run) is_full_cudagraph = batch_desc.cg_mode == CUDAGraphMode.FULL diff --git a/vllm/v1/worker/gpu/spec_decode/lilicorr/speculator.py b/vllm/v1/worker/gpu/spec_decode/lilicorr/speculator.py index 213dd962c6..e7177ead9e 100644 --- a/vllm/v1/worker/gpu/spec_decode/lilicorr/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/lilicorr/speculator.py @@ -6,7 +6,6 @@ from torch import nn from vllm.config import VllmConfig from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding -from vllm.v1.worker.gpu.input_batch import InputBatch from vllm.v1.worker.gpu.spec_decode.dflash2.speculator import DFlash2Speculator from vllm.v1.worker.gpu.spec_decode.eagle.utils import get_target_lm_head @@ -24,6 +23,9 @@ class LiLiCorrSpeculator(DFlash2Speculator): self.anchor_valid = torch.zeros( self.max_num_reqs, dtype=torch.bool, device=device ) + self._num_rejected = torch.zeros( + self.max_num_reqs, dtype=torch.int32, device=device + ) def load_draft_model( self, target_model: nn.Module, target_attn_layer_names: set[str] @@ -67,11 +69,14 @@ class LiLiCorrSpeculator(DFlash2Speculator): return model def prepare_context_anchor( - self, input_batch: InputBatch, num_rejected: torch.Tensor + self, + num_reqs: int, + query_start_loc: torch.Tensor, + num_rejected: torch.Tensor | None, ) -> None: - num_reqs = input_batch.num_reqs - starts = input_batch.query_start_loc[:num_reqs] - ends = input_batch.query_start_loc[1 : num_reqs + 1] - num_rejected[:num_reqs] + assert num_rejected is not None + starts = query_start_loc[:num_reqs] + ends = query_start_loc[1 : num_reqs + 1] - num_rejected[:num_reqs] valid = ends > starts indices = (ends - 1).clamp_min(0).long() self.anchor_hidden.zero_()