iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Perf][DFlash] Capture context combine and anchor in the draft CUDA graph (#59511)

Browse Source 
Signed-off-by: Juntian Liu <Juntianl777@gmail.com>
Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
This commit is contained in:
Juntian Liu
2026-10-14 15:40:52 -04:00
committed by GitHub
co-authored by Claude Opus 5.5
parent 4975fda882
commit 388a610dfb
3 changed files with 159 additions and 35 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/v1/spec_decode/test_lilicorr.py
+35 -4
View File
@@ -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."""
vllm/v1/worker/gpu/spec_decode/dflash/speculator.py
+114 -26
View File
@@ -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
vllm/v1/worker/gpu/spec_decode/lilicorr/speculator.py
+10 -5
View File
@@ -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_()