@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / aa7012eb6d
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 14, 05:30 UTC

Add tree attention backend for v1 (part 1) (#20401)

Signed-off-by: Giancarlo Delfin <gdelfin@meta.com>

Giancarlo Delfin committed Aug 3, 2025 at 22:13 UTC aa7012eb6db69baab57c80ac596d088eb81e090f
12 files changed +1098 -25
tests/v1/attention/test_attention_backends.py
+1 -1
@@ -17,7 +17,7 @@ from vllm.v1.kv_cache_interface import FullAttentionSpec
17
18 BACKENDS_TO_TEST = [
19 _Backend.FLASH_ATTN_VLLM_V1, _Backend.FLASHINFER_VLLM_V1,
20 - _Backend.FLEX_ATTENTION, _Backend.TRITON_ATTN_VLLM_V1
20 + _Backend.FLEX_ATTENTION, _Backend.TRITON_ATTN_VLLM_V1, _Backend.TREE_ATTN
21 ]
22
23 # Remove flashinfer from the list if it's not available
tests/v1/attention/utils.py
+4 -2
@@ -109,11 +109,11 @@ def create_common_attn_metadata(
109
110 def get_attention_backend(backend_name: _Backend):
111 """Set up attention backend classes for testing.
112 -
112 +
113 Args:
114 backend_name: Name of the backend ("flash_attn", "flashinfer", etc.)
115 vllm_config: VllmConfig instance
116 -
116 +
117 Returns:
118 Tuple of (backend_builder_class, backend_impl_class)
119 """
@@ -126,6 +126,8 @@ def get_attention_backend(backend_name: _Backend):
126 "vllm.v1.attention.backends.flex_attention.FlexAttentionBackend",
127 _Backend.TRITON_ATTN_VLLM_V1:
128 "vllm.v1.attention.backends.triton_attn.TritonAttentionBackend",
129 + _Backend.TREE_ATTN:
130 + "vllm.v1.attention.backends.tree_attn.TreeAttentionBackend",
131 }
132
133 if backend_name not in backend_map:
tests/v1/spec_decode/test_eagle.py
+4 -3
@@ -202,7 +202,9 @@ def test_load_model(mock_get_model, mock_get_layers, mock_get_pp_group, method,
202
203
204 @pytest.mark.parametrize("num_speculative_tokens", [1, 3, 8])
205 -def test_propose(num_speculative_tokens):
205 +@pytest.mark.parametrize("backend",
206 + [_Backend.FLASH_ATTN_VLLM_V1, _Backend.TREE_ATTN])
207 +def test_propose(num_speculative_tokens, backend):
208 # Use GPU device
209 device = torch.device(current_platform.device_type)
210
@@ -301,8 +303,7 @@ def test_propose(num_speculative_tokens):
303 device=device)
304 sampling_metadata = mock.MagicMock()
305
304 - attn_metadata_builder_cls, _ = get_attention_backend(
305 - _Backend.FLASH_ATTN_VLLM_V1)
306 + attn_metadata_builder_cls, _ = get_attention_backend(backend)
307 attn_metadata_builder = attn_metadata_builder_cls(
308 kv_cache_spec=create_standard_kv_cache_spec(proposer.vllm_config),
309 layer_names=proposer.attn_layer_names,
tests/v1/spec_decode/test_tree_attention.py new
+299
@@ -0,0 +1,299 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import math
5 +from typing import Optional
6 +
7 +import torch
8 +
9 +from tests.v1.attention.utils import (_Backend, create_standard_kv_cache_spec,
10 + create_vllm_config,
11 + get_attention_backend)
12 +from vllm.config import ParallelConfig, SpeculativeConfig
13 +from vllm.v1.attention.backends.utils import CommonAttentionMetadata
14 +
15 +
16 +class MockAttentionLayer(torch.nn.Module):
17 + _q_scale = torch.tensor(1.0, dtype=torch.float32, device="cuda")
18 + _k_scale = torch.tensor(1.0, dtype=torch.float32, device="cuda")
19 + _v_scale = torch.tensor(1.0, dtype=torch.float32, device="cuda")
20 +
21 + def __init__(self):
22 + super().__init__()
23 +
24 + def forward(self, x):
25 + return x
26 +
27 +
28 +def forward_attention(
29 + q: torch.Tensor,
30 + k: torch.Tensor,
31 + v: torch.Tensor,
32 + kv_cache: torch.Tensor,
33 + block_table: torch.Tensor,
34 + slot_mapping: torch.Tensor,
35 + seqlen_k: int,
36 + backend: _Backend,
37 + spec_token_tree: Optional[str] = None,
38 + num_spec_tokens: int = 0,
39 +) -> torch.Tensor:
40 + batch_size, q_len, num_heads, dim_per_head = q.shape
41 + num_kv_heads = k.shape[-2]
42 + # Initialize the query and KV sequence lengths.
43 + query_start_loc = q_len * torch.arange(
44 + batch_size + 1, device=q.device, dtype=torch.int32)
45 + query_lens = torch.diff(query_start_loc)
46 + seq_lens = torch.full(
47 + (batch_size, ),
48 + seqlen_k,
49 + device=q.device,
50 + dtype=torch.int32,
51 + )
52 + context_lens = seq_lens - query_lens
53 + max_query_len = q_len
54 + num_actual_tokens = query_start_loc[-1]
55 +
56 + softmax_scale = q.shape[-1]**(-0.5)
57 + layer = MockAttentionLayer()
58 +
59 + # Build common metadata.
60 + model_name = "meta-llama/Meta-Llama-3-8B"
61 + builder_cls, impl_cls = get_attention_backend(backend)
62 + vllm_config = create_vllm_config(model_name=model_name,
63 + max_model_len=max(seq_lens))
64 + if spec_token_tree is not None:
65 + # Create speculative config if token tree is specified.
66 + vllm_config.speculative_config = SpeculativeConfig(
67 + target_model_config=vllm_config.model_config,
68 + target_parallel_config=ParallelConfig(),
69 + model=model_name,
70 + method="eagle",
71 + num_speculative_tokens=num_spec_tokens,
72 + speculative_token_tree=spec_token_tree)
73 + kv_cache_spec = create_standard_kv_cache_spec(vllm_config)
74 + builder = builder_cls(kv_cache_spec, [], vllm_config, q.device)
75 + common_attn_metadata = CommonAttentionMetadata(
76 + query_start_loc=query_start_loc,
77 + query_start_loc_cpu=query_start_loc.cpu(),
78 + seq_lens=seq_lens,
79 + seq_lens_cpu=seq_lens.cpu(),
80 + num_computed_tokens_cpu=context_lens.cpu(),
81 + num_reqs=batch_size,
82 + num_actual_tokens=num_actual_tokens,
83 + max_query_len=max_query_len,
84 + block_table_tensor=block_table,
85 + slot_mapping=slot_mapping,
86 + )
87 +
88 + # Build attention metadata.
89 + attn_metadata = builder.build(
90 + common_prefix_len=0,
91 + common_attn_metadata=common_attn_metadata,
92 + )
93 +
94 + # Initialize the backend implementation.
95 + instance = impl_cls(
96 + num_heads=num_heads,
97 + head_size=dim_per_head,
98 + scale=softmax_scale,
99 + num_kv_heads=num_kv_heads,
100 + alibi_slopes=None,
101 + sliding_window=None,
102 + kv_cache_dtype="auto",
103 + )
104 +
105 + # Run forward pass and return output.
106 + query = q.view(-1, num_heads, dim_per_head)
107 + key = k.view(-1, num_kv_heads, dim_per_head)
108 + value = v.view(-1, num_kv_heads, dim_per_head)
109 + output = torch.empty_like(query)
110 + return instance.forward(
111 + layer=layer,
112 + query=query,
113 + key=key,
114 + value=value,
115 + kv_cache=kv_cache.clone(),
116 + attn_metadata=attn_metadata,
117 + output=output,
118 + )
119 +
120 +
121 +def test_tree_attn_correctness() -> None:
122 + torch.manual_seed(42)
123 + torch.cuda.manual_seed_all(42)
124 +
125 + device = "cuda"
126 + tree_attn_masks = {
127 + # Chain.
128 + "[(0,), (0, 0), (0, 0, 0)]":
129 + torch.tensor(
130 + [
131 + [1, 0, 0, 0],
132 + [1, 1, 0, 0],
133 + [1, 1, 1, 0],
134 + [1, 1, 1, 1],
135 + ],
136 + device=device,
137 + dtype=torch.int32,
138 + ),
139 + # Tree.
140 + "[(0,), (1,), (0, 0), (0, 1), (1, 0), (1, 1)]":
141 + torch.tensor(
142 + [
143 + [1, 0, 0, 0, 0, 0, 0],
144 + [1, 1, 0, 0, 0, 0, 0],
145 + [1, 0, 1, 0, 0, 0, 0],
146 + [1, 1, 0, 1, 0, 0, 0],
147 + [1, 1, 0, 0, 1, 0, 0],
148 + [1, 0, 1, 0, 0, 1, 0],
149 + [1, 0, 1, 0, 0, 0, 1],
150 + ],
151 + device=device,
152 + dtype=torch.int32,
153 + ),
154 + }
155 +
156 + dim_per_head = 128
157 + num_kv_heads = 2
158 + block_size = 128
159 + max_sequence_length = 8192
160 + randomize_blocks = True
161 + for batch_size in [1, 16, 32]:
162 + for num_heads in [2, 4]:
163 + for sequence_position in [16, 1024, 2048]:
164 + for spec_token_tree, tree_attn_mask in tree_attn_masks.items():
165 + # Assert that the number of heads is divisible
166 + # by the number of KV heads.
167 + assert num_heads % num_kv_heads == 0
168 +
169 + # Initialize q, k, and v.
170 + tree_size_q = tree_attn_mask.shape[0]
171 + seqlen_k = sequence_position + tree_size_q
172 + q = torch.randn(
173 + (batch_size, tree_size_q, num_heads, dim_per_head),
174 + device=device,
175 + dtype=torch.bfloat16,
176 + )
177 + k = torch.randn(
178 + (batch_size, tree_size_q, num_kv_heads, dim_per_head),
179 + device=device,
180 + dtype=torch.bfloat16,
181 + )
182 + v = torch.randn(
183 + (batch_size, tree_size_q, num_kv_heads, dim_per_head),
184 + device=device,
185 + dtype=torch.bfloat16,
186 + )
187 +
188 + # Setup the block table and KV cache for paged KV.
189 + assert max_sequence_length % block_size == 0
190 + max_blocks_per_batch = max_sequence_length // block_size
191 + kv_cache = torch.randn(
192 + (
193 + 2,
194 + batch_size * max_blocks_per_batch,
195 + block_size,
196 + num_kv_heads,
197 + dim_per_head,
198 + ),
199 + device=q.device,
200 + dtype=torch.bfloat16,
201 + )
202 + num_alloc_blocks_per_batch = math.ceil(seqlen_k /
203 + block_size)
204 + block_table = torch.zeros(
205 + (batch_size, max_blocks_per_batch),
206 + device=q.device,
207 + dtype=torch.int32,
208 + )
209 + block_ids = torch.arange(
210 + 0,
211 + batch_size * num_alloc_blocks_per_batch,
212 + device=q.device,
213 + dtype=torch.int32,
214 + )
215 + if randomize_blocks:
216 + # Randomize the block ids.
217 + block_ids = block_ids[torch.randperm(
218 + block_ids.numel())]
219 + block_table[:, :
220 + num_alloc_blocks_per_batch] = block_ids.view(
221 + -1, num_alloc_blocks_per_batch)
222 +
223 + # Setup the slot mapping for the input KVs.
224 + tree_positions = sequence_position + torch.arange(
225 + 0,
226 + tree_size_q,
227 + device=q.device,
228 + dtype=torch.int64,
229 + ).repeat(batch_size, 1)
230 + tree_slot_mapping = _gen_slot_mapping(
231 + tree_positions, block_table, block_size)
232 +
233 + # Compute attention for the tree.
234 + tree_attn_output = forward_attention(
235 + q=q,
236 + k=k,
237 + v=v,
238 + kv_cache=kv_cache,
239 + block_table=block_table,
240 + slot_mapping=tree_slot_mapping,
241 + seqlen_k=seqlen_k,
242 + backend=_Backend.TREE_ATTN,
243 + spec_token_tree=spec_token_tree,
244 + num_spec_tokens=tree_size_q - 1,
245 + ).view(batch_size, -1, num_heads, dim_per_head)
246 +
247 + # Verify that the chain attention output for each
248 + # branch of the tree (computed using FA3) matches
249 + # the tree attention output.
250 + for q_index in range(tree_size_q):
251 + # Get the q, k, and v for the branch.
252 + branch_mask = tree_attn_mask[q_index, :]
253 + branch_indices = torch.nonzero(branch_mask,
254 + as_tuple=True)[0]
255 + q_len = branch_indices.shape[0]
256 + q_branch = q[:, branch_indices]
257 + k_branch = k[:, branch_indices]
258 + v_branch = v[:, branch_indices]
259 +
260 + # Setup slot mapping for the branch.
261 + branch_positions = sequence_position + torch.arange(
262 + 0,
263 + q_len,
264 + device=q.device,
265 + dtype=torch.int64,
266 + ).repeat(batch_size, 1)
267 + branch_slot_mapping = _gen_slot_mapping(
268 + branch_positions, block_table, block_size)
269 +
270 + # Compute flash attention for the branch.
271 + flash_attn_output = forward_attention(
272 + q=q_branch,
273 + k=k_branch,
274 + v=v_branch,
275 + kv_cache=kv_cache,
276 + block_table=block_table,
277 + slot_mapping=branch_slot_mapping,
278 + seqlen_k=sequence_position + q_len,
279 + backend=_Backend.FLASH_ATTN_VLLM_V1,
280 + ).view(batch_size, -1, num_heads, dim_per_head)
281 +
282 + # Compare the outputs.
283 + assert torch.allclose(
284 + tree_attn_output[:, branch_indices],
285 + flash_attn_output,
286 + atol=7.81e-3,
287 + ), (f"outputs are not close for "
288 + f"batch_size: {batch_size}, "
289 + f"num_heads: {num_heads}, "
290 + f"sequence_position: {sequence_position}, "
291 + f"tree_attn_mask: {tree_attn_mask}, "
292 + f"q_index: {q_index}.")
293 +
294 +
295 +def _gen_slot_mapping(positions: torch.Tensor, block_table: torch.Tensor,
296 + block_size: int):
297 + block_indices = positions // block_size
298 + blocks = block_table.gather(dim=1, index=block_indices)
299 + return (blocks * block_size + positions % block_size).view(-1)
vllm/attention/ops/triton_unified_attention.py
+48
@@ -55,6 +55,7 @@ def kernel_unified_attention_2d(
55 block_tables_ptr, # [num_seqs, max_num_blocks_per_seq]
56 seq_lens_ptr, # [num_seqs]
57 alibi_slopes_ptr, # [num_query_heads]
58 + qq_bias_ptr, # [num_query_tokens, num_query_tokens]
59 scale, # float32
60 k_scale, # float32
61 v_scale, # float32
@@ -66,10 +67,12 @@ def kernel_unified_attention_2d(
67 query_stride_1: tl.int64, # int, should be equal to head_size
68 output_stride_0: tl.int64, # int
69 output_stride_1: tl.int64, # int, should be equal to head_size
70 + qq_bias_stride_0: tl.int64, # int
71 BLOCK_SIZE: tl.constexpr, # int
72 HEAD_SIZE: tl.constexpr, # int
73 HEAD_SIZE_PADDED: tl.constexpr, # int, must be power of 2
74 USE_ALIBI_SLOPES: tl.constexpr, # bool
75 + USE_QQ_BIAS: tl.constexpr, # bool
76 USE_SOFTCAP: tl.constexpr, # bool
77 SLIDING_WINDOW: tl.constexpr, # int
78 stride_k_cache_0: tl.int64, # int
@@ -144,6 +147,11 @@ def kernel_unified_attention_2d(
147 mask=query_mask_1,
148 other=0.0)
149
150 + # query-query attention bias
151 + if USE_QQ_BIAS:
152 + qq_bias_row_ptrs = (qq_bias_ptr + query_pos[:, None] * qq_bias_stride_0
153 + ) # shape: [BLOCK_M]
154 +
155 # compute the length of the longest sequence prefix spanned by any
156 # query token in the current q_block (q_block_local_idx)
157 max_seq_prefix_len = context_len + q_block_local_idx * BLOCK_Q + (
@@ -223,6 +231,18 @@ def kernel_unified_attention_2d(
231 if USE_ALIBI_SLOPES:
232 S += alibi_slope[:, None] * (seq_offset - context_len)
233
234 + if USE_QQ_BIAS:
235 + # compute key positions relative to query section
236 + key_rel_pos = seq_offset - context_len # shape: [BLOCK_SIZE]
237 + # load bias only for keys that correspond to queries
238 + is_query_key = key_rel_pos >= 0 and key_rel_pos < qq_bias_stride_0
239 + qq_bias = tl.load(
240 + qq_bias_row_ptrs + key_rel_pos[None, :],
241 + mask=is_query_key[None, :], # avoid OOB for context keys
242 + other=0.0,
243 + )
244 + S += qq_bias
245 +
246 # compute running maximum
247 # m_j : (BLOCK_M,)
248 m_j = tl.maximum(M, tl.max(S, axis=1))
@@ -275,6 +295,7 @@ def kernel_unified_attention_3d(
295 block_tables_ptr, # [num_seqs, max_num_blocks_per_seq]
296 seq_lens_ptr, # [num_seqs]
297 alibi_slopes_ptr, # [num_query_heads]
298 + qq_bias_ptr, # [num_query_tokens, num_query_tokens]
299 scale, # float32
300 k_scale, # float32
301 v_scale, # float32
@@ -284,10 +305,12 @@ def kernel_unified_attention_3d(
305 block_table_stride: tl.int64, # int
306 query_stride_0: tl.int64, # int
307 query_stride_1: tl.int64, # int, should be equal to head_size
308 + qq_bias_stride_0: tl.int64, # int
309 BLOCK_SIZE: tl.constexpr, # int
310 HEAD_SIZE: tl.constexpr, # int
311 HEAD_SIZE_PADDED: tl.constexpr, # int, must be power of 2
312 USE_ALIBI_SLOPES: tl.constexpr, # bool
313 + USE_QQ_BIAS: tl.constexpr, # bool
314 USE_SOFTCAP: tl.constexpr, # bool
315 SLIDING_WINDOW: tl.constexpr, # int
316 stride_k_cache_0: tl.int64, # int
@@ -373,6 +396,11 @@ def kernel_unified_attention_3d(
396 mask=query_mask_1,
397 other=0.0)
398
399 + # query-query attention bias
400 + if USE_QQ_BIAS:
401 + qq_bias_row_ptrs = (qq_bias_ptr + query_pos[:, None] * qq_bias_stride_0
402 + ) # shape: [BLOCK_M]
403 +
404 num_blocks = cdiv_fn(seq_len, BLOCK_SIZE)
405
406 # iterate through tiles within current segment
@@ -442,6 +470,18 @@ def kernel_unified_attention_3d(
470 if USE_ALIBI_SLOPES:
471 S += alibi_slope[:, None] * (seq_offset - context_len)
472
473 + if USE_QQ_BIAS:
474 + # compute key positions relative to query section
475 + key_rel_pos = seq_offset - context_len # shape: [BLOCK_SIZE]
476 + # load bias only for keys that correspond to queries
477 + is_query_key = key_rel_pos >= 0 and key_rel_pos < qq_bias_stride_0
478 + qq_bias = tl.load(
479 + qq_bias_row_ptrs + key_rel_pos[None, :],
480 + mask=is_query_key[None, :], # avoid OOB for context keys
481 + other=0.0,
482 + )
483 + S += qq_bias
484 +
485 # compute running maximum
486 # m_j : (BLOCK_M,)
487 m_j = tl.maximum(M, tl.max(S, axis=1))
@@ -586,6 +626,7 @@ def unified_attention(
626 k_descale,
627 v_descale,
628 alibi_slopes=None,
629 + qq_bias=None,
630 ):
631 assert causal, "Only causal attention is supported"
632 assert q_descale is None, "Q scales not supported"
@@ -595,6 +636,7 @@ def unified_attention(
636 "Block size must be at least 32 for fp8"
637
638 use_alibi_slopes = alibi_slopes is not None
639 + use_qq_bias = qq_bias is not None
640
641 block_size = v.shape[1]
642 num_seqs = len(seqused_k)
@@ -630,6 +672,7 @@ def unified_attention(
672 block_tables_ptr=block_table,
673 seq_lens_ptr=seqused_k,
674 alibi_slopes_ptr=alibi_slopes,
675 + qq_bias_ptr=qq_bias,
676 scale=softmax_scale,
677 k_scale=k_descale,
678 v_scale=v_descale,
@@ -641,10 +684,12 @@ def unified_attention(
684 query_stride_1=q.stride(1),
685 output_stride_0=out.stride(0),
686 output_stride_1=out.stride(1),
687 + qq_bias_stride_0=qq_bias.stride(0) if use_qq_bias else 0,
688 BLOCK_SIZE=block_size,
689 HEAD_SIZE=head_size,
690 HEAD_SIZE_PADDED=triton.next_power_of_2(head_size),
691 USE_ALIBI_SLOPES=use_alibi_slopes,
692 + USE_QQ_BIAS=use_qq_bias,
693 USE_SOFTCAP=(softcap > 0),
694 SLIDING_WINDOW=(1 + window_size[0]),
695 stride_k_cache_0=k.stride(0),
@@ -699,6 +744,7 @@ def unified_attention(
744 block_tables_ptr=block_table,
745 seq_lens_ptr=seqused_k,
746 alibi_slopes_ptr=alibi_slopes,
747 + qq_bias_ptr=qq_bias,
748 scale=softmax_scale,
749 k_scale=k_descale,
750 v_scale=v_descale,
@@ -708,10 +754,12 @@ def unified_attention(
754 block_table_stride=block_table.stride(0),
755 query_stride_0=q.stride(0),
756 query_stride_1=q.stride(1),
757 + qq_bias_stride_0=qq_bias.stride(0) if use_qq_bias else 0,
758 BLOCK_SIZE=block_size,
759 HEAD_SIZE=head_size,
760 HEAD_SIZE_PADDED=triton.next_power_of_2(head_size),
761 USE_ALIBI_SLOPES=use_alibi_slopes,
762 + USE_QQ_BIAS=use_qq_bias,
763 USE_SOFTCAP=(softcap > 0),
764 SLIDING_WINDOW=(1 + window_size[0]),
765 stride_k_cache_0=k.stride(0),
vllm/config.py
+13
@@ -3049,6 +3049,19 @@ class SpeculativeConfig:
3049 f"num_speculative_tokens:{self.num_speculative_tokens}"
3050 f" must be divisible by {n_predict=}")
3051
3052 + if self.speculative_token_tree is None:
3053 + # Generate chain of tokens.
3054 + self.speculative_token_tree = str([
3055 + (i + 1) * (0, )
3056 + for i in range(self.num_speculative_tokens)
3057 + ])
3058 + else:
3059 + # Sort the token tree breadth-first.
3060 + tree_choices = ast.literal_eval(
3061 + self.speculative_token_tree)
3062 + self.speculative_token_tree = str(
3063 + sorted(tree_choices, key=lambda t: (len(t), t)))
3064 +
3065 self.draft_tensor_parallel_size = \
3066 SpeculativeConfig._verify_and_get_draft_tp(
3067 self.target_parallel_config,
vllm/engine/arg_utils.py
+1 -1
@@ -1454,7 +1454,6 @@ class EngineArgs:
1454 "Please consider using other speculative decoding methods "
1455 "such as ngram, medusa, eagle, or deepseek_mtp.")
1456
1457 - # No XFormers so far.
1457 V1_BACKENDS = [
1458 "FLASH_ATTN_VLLM_V1",
1459 "FLASH_ATTN",
@@ -1469,6 +1468,7 @@ class EngineArgs:
1468 "ROCM_AITER_MLA",
1469 "TORCH_SDPA_VLLM_V1",
1470 "FLEX_ATTENTION",
1471 + "TREE_ATTN",
1472 ]
1473 if (envs.is_set("VLLM_ATTENTION_BACKEND")
1474 and envs.VLLM_ATTENTION_BACKEND not in V1_BACKENDS):
vllm/platforms/cuda.py
+4
@@ -270,6 +270,7 @@ class CudaPlatformBase(Platform):
270 FLEX_ATTENTION_V1 = "vllm.v1.attention.backends.flex_attention.FlexAttentionBackend" # noqa: E501
271 TRITON_ATTN_VLLM_V1 = "vllm.v1.attention.backends.triton_attn.TritonAttentionBackend" # noqa: E501
272 FLASH_ATTN_V1 = "vllm.v1.attention.backends.flash_attn.FlashAttentionBackend" # noqa: E501
273 + TREE_ATTN_V1 = "vllm.v1.attention.backends.tree_attn.TreeAttentionBackend" # noqa: E501
274
275 if selected_backend == _Backend.FLASHINFER:
276 logger.info_once("Using FlashInfer backend on V1 engine.")
@@ -287,6 +288,9 @@ class CudaPlatformBase(Platform):
288 elif selected_backend == _Backend.FLASH_ATTN:
289 logger.info_once("Using Flash Attention backend on V1 engine.")
290 return FLASH_ATTN_V1
291 + elif selected_backend == _Backend.TREE_ATTN:
292 + logger.info_once("Using Tree Attention backend on V1 engine.")
293 + return TREE_ATTN_V1
294
295 from vllm.attention.selector import is_attn_backend_supported
296
vllm/platforms/interface.py
+1
@@ -62,6 +62,7 @@ class _Backend(enum.Enum):
62 DIFFERENTIAL_FLASH_ATTN = enum.auto()
63 NO_ATTENTION = enum.auto()
64 FLEX_ATTENTION = enum.auto()
65 + TREE_ATTN = enum.auto()
66
67
68 class PlatformEnum(enum.Enum):
vllm/v1/attention/backends/tree_attn.py new
+452
@@ -0,0 +1,452 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Attention layer with TreeAttention."""
4 +
5 +import ast
6 +from dataclasses import dataclass
7 +from typing import TYPE_CHECKING, Any, Optional
8 +
9 +import torch
10 +
11 +from vllm.attention.backends.abstract import (AttentionBackend, AttentionImpl,
12 + AttentionMetadata, AttentionType)
13 +from vllm.attention.ops.triton_unified_attention import unified_attention
14 +from vllm.config import VllmConfig
15 +from vllm.logger import init_logger
16 +from vllm.v1.attention.backends.utils import (
17 + AttentionMetadataBuilder, CommonAttentionMetadata,
18 + reorder_batch_to_split_decodes_and_prefills, split_decodes_and_prefills)
19 +from vllm.v1.kv_cache_interface import AttentionSpec
20 +
21 +if TYPE_CHECKING:
22 + from vllm.v1.core.sched.output import SchedulerOutput
23 + from vllm.v1.worker.gpu_input_batch import InputBatch
24 +
25 +from vllm import _custom_ops as ops
26 +
27 +logger = init_logger(__name__)
28 +
29 +
30 +class TreeAttentionBackend(AttentionBackend):
31 +
32 + accept_output_buffer: bool = True
33 +
34 + @classmethod
35 + def get_supported_dtypes(cls) -> list[torch.dtype]:
36 + return [torch.float16, torch.bfloat16]
37 +
38 + @classmethod
39 + def get_supported_head_sizes(cls) -> list[int]:
40 + return [32, 64, 96, 128, 160, 192, 224, 256]
41 +
42 + @classmethod
43 + def validate_head_size(cls, head_size: int) -> None:
44 + supported_head_sizes = cls.get_supported_head_sizes()
45 + if head_size not in supported_head_sizes:
46 + attn_type = cls.__name__.removesuffix("Backend")
47 + raise ValueError(
48 + f"Head size {head_size} is not supported by {attn_type}. "
49 + f"Supported head sizes are: {supported_head_sizes}. "
50 + "Set VLLM_ATTENTION_BACKEND=FLEX_ATTENTION to use "
51 + "FlexAttention backend which supports all head sizes.")
52 +
53 + @staticmethod
54 + def get_name() -> str:
55 + return "TREE_ATTN_VLLM_V1"
56 +
57 + @staticmethod
58 + def get_impl_cls() -> type["TreeAttentionImpl"]:
59 + return TreeAttentionImpl
60 +
61 + @staticmethod
62 + def get_metadata_cls() -> type["AttentionMetadata"]:
63 + return TreeAttentionMetadata
64 +
65 + @staticmethod
66 + def get_kv_cache_shape(
67 + num_blocks: int,
68 + block_size: int,
69 + num_kv_heads: int,
70 + head_size: int,
71 + ) -> tuple[int, ...]:
72 + if block_size % 16 != 0:
73 + raise ValueError("Block size must be a multiple of 16.")
74 + return (2, num_blocks, block_size, num_kv_heads, head_size)
75 +
76 + @staticmethod
77 + def get_builder_cls() -> type["TreeAttentionMetadataBuilder"]:
78 + return TreeAttentionMetadataBuilder
79 +
80 + @staticmethod
81 + def use_cascade_attention(*args, **kwargs) -> bool:
82 + return False
83 +
84 +
85 +@dataclass
86 +class TreeAttentionMetadata:
87 + num_actual_tokens: int # Number of tokens excluding padding.
88 + max_query_len: int
89 + query_start_loc: torch.Tensor
90 + max_seq_len: int
91 + seq_lens: torch.Tensor
92 + block_table: torch.Tensor
93 + slot_mapping: torch.Tensor
94 +
95 + num_prefill_tokens: int = 0
96 + num_decode_tokens: int = 0
97 + num_prefills: int = 0
98 + num_decodes: int = 0
99 +
100 + tree_attn_bias: Optional[torch.Tensor] = None
101 +
102 + # Cached Prefill/decode metadata.
103 + _cached_prefill_metadata: Optional["TreeAttentionMetadata"] = None
104 + _cached_decode_metadata: Optional["TreeAttentionMetadata"] = None
105 +
106 + @property
107 + def prefill_metadata(self) -> Optional["TreeAttentionMetadata"]:
108 + if self.num_prefills == 0:
109 + return None
110 +
111 + if self._cached_prefill_metadata is not None:
112 + # Recover cached prefill-phase attention
113 + # metadata structure
114 + return self._cached_prefill_metadata
115 +
116 + q_start_loc = self.query_start_loc[self.num_decodes:]
117 + q_seqlens = torch.diff(q_start_loc)
118 + kv_seqlens = self.seq_lens[self.num_decodes:]
119 + # Construct & cache prefill-phase attention metadata structure
120 + self._cached_prefill_metadata = TreeAttentionMetadata(
121 + num_actual_tokens=self.num_prefill_tokens,
122 + max_query_len=int(q_seqlens.max().item()),
123 + query_start_loc=q_start_loc - q_start_loc[0],
124 + max_seq_len=int(kv_seqlens.max().item()),
125 + seq_lens=kv_seqlens,
126 + block_table=self.block_table[self.num_decodes:],
127 + slot_mapping=self.slot_mapping[self.num_decode_tokens:],
128 + )
129 + return self._cached_prefill_metadata
130 +
131 + @property
132 + def decode_metadata(self) -> Optional["TreeAttentionMetadata"]:
133 + if self.num_decode_tokens == 0:
134 + return None
135 +
136 + if self._cached_decode_metadata is not None:
137 + # Recover cached decode-phase attention
138 + # metadata structure
139 + return self._cached_decode_metadata
140 +
141 + q_start_loc = self.query_start_loc[:self.num_decodes + 1]
142 + q_seqlens = torch.diff(q_start_loc)
143 + kv_seqlens = self.seq_lens[:self.num_decodes]
144 + # Construct & cache decode-phase attention metadata structure
145 + self._cached_decode_metadata = TreeAttentionMetadata(
146 + num_actual_tokens=self.num_decode_tokens,
147 + max_query_len=int(q_seqlens.max().item()),
148 + query_start_loc=q_start_loc,
149 + max_seq_len=int(kv_seqlens.max().item()),
150 + seq_lens=kv_seqlens,
151 + block_table=self.block_table[:self.num_decodes],
152 + slot_mapping=self.slot_mapping[:self.num_decode_tokens],
153 + tree_attn_bias=self.tree_attn_bias,
154 + )
155 + return self._cached_decode_metadata
156 +
157 +
158 +class TreeAttentionMetadataBuilder(
159 + AttentionMetadataBuilder[TreeAttentionMetadata]):
160 +
161 + def __init__(
162 + self,
163 + kv_cache_spec: AttentionSpec,
164 + layer_names: list[str],
165 + vllm_config: VllmConfig,
166 + device: torch.device,
167 + ):
168 + self.kv_cache_spec = kv_cache_spec
169 + self.block_size = kv_cache_spec.block_size
170 +
171 + spec_config = vllm_config.speculative_config
172 + spec_token_tree = (spec := spec_config) and spec.speculative_token_tree
173 + tree_choices: list[tuple[int,
174 + ...]] = (ast.literal_eval(spec_token_tree)
175 + if spec_token_tree is not None else
176 + [(0, )])
177 + # Construct the tree attention bias.
178 + depth_counts = _get_depth_counts(tree_choices)
179 + self.tree_attn_bias = _prepare_tree_attn_bias(
180 + tree_choices,
181 + depth_counts,
182 + dtype=torch.float32,
183 + device=device,
184 + )
185 +
186 + def reorder_batch(self, input_batch: "InputBatch",
187 + scheduler_output: "SchedulerOutput") -> bool:
188 + return reorder_batch_to_split_decodes_and_prefills(
189 + input_batch,
190 + scheduler_output,
191 + decode_threshold=self.tree_attn_bias.shape[0])
192 +
193 + def build(
194 + self,
195 + common_prefix_len: int,
196 + common_attn_metadata: CommonAttentionMetadata,
197 + fast_build: bool = False,
198 + ) -> TreeAttentionMetadata:
199 + decode_threshold = self.tree_attn_bias.shape[0]
200 + num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (
201 + split_decodes_and_prefills(common_attn_metadata,
202 + decode_threshold=decode_threshold))
203 +
204 + num_actual_tokens = common_attn_metadata.num_actual_tokens
205 + q_start_loc = common_attn_metadata.query_start_loc
206 + max_query_len = common_attn_metadata.max_query_len
207 + kv_seqlens = common_attn_metadata.seq_lens
208 + max_seq_len = int(common_attn_metadata.seq_lens_cpu.max())
209 + block_table = common_attn_metadata.block_table_tensor
210 + slot_mapping = common_attn_metadata.slot_mapping
211 +
212 + return TreeAttentionMetadata(
213 + num_actual_tokens=num_actual_tokens,
214 + num_prefill_tokens=num_prefill_tokens,
215 + num_decode_tokens=num_decode_tokens,
216 + num_prefills=num_prefills,
217 + num_decodes=num_decodes,
218 + max_query_len=max_query_len,
219 + query_start_loc=q_start_loc,
220 + max_seq_len=max_seq_len,
221 + seq_lens=kv_seqlens,
222 + block_table=block_table,
223 + slot_mapping=slot_mapping,
224 + tree_attn_bias=self.tree_attn_bias,
225 + )
226 +
227 + def build_for_drafting(
228 + self,
229 + common_attn_metadata: CommonAttentionMetadata,
230 + draft_index: int,
231 + ) -> TreeAttentionMetadata:
232 + # Cache the original tree attention bias.
233 + orig_tree_attn_bias = self.tree_attn_bias
234 +
235 + if draft_index == 0:
236 + # Use prefill for drafting at the root level.
237 + self.tree_attn_bias = torch.empty(0)
238 + else:
239 + # Slice the tree attention bias for drafting.
240 + query_len = common_attn_metadata.max_query_len
241 + start, end = draft_index, draft_index + query_len
242 + self.tree_attn_bias = self.tree_attn_bias[start:end,
243 + start:end].contiguous()
244 +
245 + # Build attention bias.
246 + attn_metadata = self.build(0, common_attn_metadata, fast_build=True)
247 +
248 + # Reset the tree attention bias to the original value.
249 + self.tree_attn_bias = orig_tree_attn_bias
250 + return attn_metadata
251 +
252 +
253 +def _get_depth_counts(sorted_tree_choices: list[tuple[int, ...]]) -> list[int]:
254 + # Count the number of choices at each depth of the tree.
255 + depth_counts = []
256 + prev_depth = 0
257 + for path in sorted_tree_choices:
258 + depth = len(path)
259 + if depth != prev_depth:
260 + depth_counts.append(0)
261 + depth_counts[depth - 1] += 1
262 + prev_depth = depth
263 + return depth_counts
264 +
265 +
266 +def _prepare_tree_attn_bias(
267 + sorted_tree_choices: list[tuple[int, ...]],
268 + depth_counts: list[int],
269 + dtype: Optional[torch.dtype],
270 + device: Optional[torch.device],
271 +) -> torch.Tensor:
272 + # +1 comes from the additional root node.
273 + tree_len = len(sorted_tree_choices) + 1
274 + tree_attn_mask = torch.full((tree_len, tree_len),
275 + -torch.inf,
276 + device=device,
277 + dtype=dtype)
278 +
279 + # Set diagonal to all zeros. Each token should
280 + # attend to itself.
281 + mask_val = 0
282 + for i in range(tree_len):
283 + tree_attn_mask[i, i] = mask_val
284 +
285 + # Set root to all zeros. All tokens attend to it.
286 + tree_attn_mask[:, 0] = mask_val
287 +
288 + # Set all ancestors to zeros.
289 + start = 0
290 + for i in range(len(depth_counts)):
291 + for j in range(depth_counts[i]):
292 + cur_tree_choice = sorted_tree_choices[start + j]
293 + # Retrieve ancestor position.
294 + if len(cur_tree_choice) == 1:
295 + continue
296 + ancestor_idx = []
297 + for c in range(len(cur_tree_choice) - 1):
298 + ancestor_idx.append(
299 + sorted_tree_choices.index(cur_tree_choice[:c + 1]) + 1)
300 + tree_attn_mask[j + start + 1, ancestor_idx] = mask_val
301 + start += depth_counts[i]
302 + return tree_attn_mask
303 +
304 +
305 +class TreeAttentionImpl(AttentionImpl):
306 +
307 + def __init__(
308 + self,
309 + num_heads: int,
310 + head_size: int,
311 + scale: float,
312 + num_kv_heads: int,
313 + alibi_slopes: Optional[list[float]],
314 + sliding_window: Optional[int],
315 + kv_cache_dtype: str,
316 + blocksparse_params: Optional[dict[str, Any]] = None,
317 + logits_soft_cap: Optional[float] = None,
318 + attn_type: AttentionType = AttentionType.DECODER,
319 + kv_sharing_target_layer_name: Optional[str] = None,
320 + use_irope: bool = False,
321 + ) -> None:
322 + if blocksparse_params is not None:
323 + raise ValueError(
324 + "TreeAttention does not support block-sparse attention.")
325 + self.num_heads = num_heads
326 + self.head_size = head_size
327 + self.scale = float(scale)
328 + self.num_kv_heads = num_kv_heads
329 + self.num_queries_per_kv = self.num_heads // self.num_kv_heads
330 + self.kv_cache_dtype = kv_cache_dtype
331 + self.kv_sharing_target_layer_name = kv_sharing_target_layer_name
332 + if alibi_slopes is not None:
333 + alibi_slopes = torch.tensor(alibi_slopes, dtype=torch.float32)
334 + self.alibi_slopes = alibi_slopes
335 + if logits_soft_cap is None:
336 + # Setting logits_soft_cap to 0 means no soft cap.
337 + logits_soft_cap = 0
338 + self.logits_soft_cap = logits_soft_cap
339 + if sliding_window is None:
340 + self.sliding_window = (-1, -1)
341 + else:
342 + self.sliding_window = (sliding_window - 1, 0)
343 +
344 + TreeAttentionBackend.validate_head_size(head_size)
345 +
346 + if attn_type != AttentionType.DECODER:
347 + raise NotImplementedError("Encoder self-attention and "
348 + "encoder/decoder cross-attention "
349 + "are not implemented for "
350 + "TreeAttentionImpl.")
351 +
352 + def forward(
353 + self,
354 + layer: torch.nn.Module,
355 + query: torch.Tensor,
356 + key: torch.Tensor,
357 + value: torch.Tensor,
358 + kv_cache: torch.Tensor,
359 + attn_metadata: TreeAttentionMetadata,
360 + output: Optional[torch.Tensor] = None,
361 + output_scale: Optional[torch.Tensor] = None,
362 + ) -> torch.Tensor:
363 + """Forward pass with TreeAttention.
364 +
365 + Args:
366 + query: shape = [num_tokens, num_heads, head_size]
367 + key: shape = [num_tokens, num_kv_heads, head_size]
368 + value: shape = [num_tokens, num_kv_heads, head_size]
369 + kv_cache = [2, num_blocks, block_size, num_kv_heads, head_size]
370 + attn_metadata: Metadata for attention.
371 + Returns:
372 + shape = [num_tokens, num_heads * head_size]
373 + """
374 + assert output is not None, "Output tensor must be provided."
375 +
376 + if output_scale is not None:
377 + raise NotImplementedError(
378 + "fused output quantization is not yet supported"
379 + " for TreeAttentionImpl")
380 +
381 + if attn_metadata is None:
382 + # Profiling run.
383 + return output
384 +
385 + # Cache the input KVs.
386 + key_cache, value_cache = kv_cache.unbind(0)
387 + if self.kv_sharing_target_layer_name is None:
388 + # Reshape the input keys and values and store them in the cache.
389 + # Skip this if sharing KV cache with an earlier attention layer.
390 + # NOTE(woosuk): Here, key and value are padded while slot_mapping is
391 + # not padded. However, we don't need to do key[:num_actual_tokens]
392 + # and value[:num_actual_tokens] because the reshape_and_cache_flash
393 + # op uses the slot_mapping's shape to determine the number of
394 + # actual tokens.
395 + ops.reshape_and_cache_flash(
396 + key,
397 + value,
398 + key_cache,
399 + value_cache,
400 + attn_metadata.slot_mapping,
401 + self.kv_cache_dtype,
402 + layer._k_scale,
403 + layer._v_scale,
404 + )
405 +
406 + num_actual_tokens = attn_metadata.num_actual_tokens
407 + num_decode_tokens = attn_metadata.num_decode_tokens
408 + descale_shape = (attn_metadata.query_start_loc.shape[0] - 1,
409 + key.shape[1])
410 + if prefill_meta := attn_metadata.prefill_metadata:
411 + unified_attention(
412 + q=query[num_decode_tokens:num_actual_tokens],
413 + k=key_cache,
414 + v=value_cache,
415 + out=output[num_decode_tokens:num_actual_tokens],
416 + cu_seqlens_q=prefill_meta.query_start_loc,
417 + max_seqlen_q=prefill_meta.max_query_len,
418 + seqused_k=prefill_meta.seq_lens,
419 + max_seqlen_k=prefill_meta.max_seq_len,
420 + softmax_scale=self.scale,
421 + causal=True,
422 + alibi_slopes=self.alibi_slopes,
423 + window_size=self.sliding_window,
424 + block_table=prefill_meta.block_table,
425 + softcap=self.logits_soft_cap,
426 + q_descale=None, # Not supported
427 + k_descale=layer._k_scale.expand(descale_shape),
428 + v_descale=layer._v_scale.expand(descale_shape),
429 + )
430 +
431 + if decode_meta := attn_metadata.decode_metadata:
432 + unified_attention(
433 + q=query[:num_decode_tokens],
434 + k=key_cache,
435 + v=value_cache,
436 + out=output[:num_decode_tokens],
437 + cu_seqlens_q=decode_meta.query_start_loc,
438 + max_seqlen_q=decode_meta.max_query_len,
439 + seqused_k=decode_meta.seq_lens,
440 + max_seqlen_k=decode_meta.max_seq_len,
441 + softmax_scale=self.scale,
442 + causal=True,
443 + alibi_slopes=self.alibi_slopes,
444 + qq_bias=decode_meta.tree_attn_bias,
445 + window_size=self.sliding_window,
446 + block_table=decode_meta.block_table,
447 + softcap=self.logits_soft_cap,
448 + q_descale=None, # Not supported
449 + k_descale=layer._k_scale.expand(descale_shape),
450 + v_descale=layer._v_scale.expand(descale_shape),
451 + )
452 + return output
vllm/v1/attention/backends/utils.py
+20
@@ -214,6 +214,26 @@ class AttentionMetadataBuilder(abc.ABC, Generic[M]):
214 return self.build(common_prefix_len=0,
215 common_attn_metadata=common_attn_metadata)
216
217 + def build_for_drafting(
218 + self,
219 + common_attn_metadata: CommonAttentionMetadata,
220 + draft_index: int,
221 + ) -> M:
222 + """
223 + Build attention metadata for draft model. Uses build by default.
224 +
225 + Args:
226 + common_attn_metadata: The common attention metadata.
227 + draft_index: The index of the current draft operation.
228 + When speculating a chain of tokens, this index refers to the
229 + draft attempt for the i-th token.
230 + For tree-based attention, this index instead refers to the
231 + draft attempt for the i-th level in the tree of tokens.
232 + """
233 + return self.build(common_prefix_len=0,
234 + common_attn_metadata=common_attn_metadata,
235 + fast_build=True)
236 +
237 def use_cascade_attention(
238 self,
239 common_prefix_len: int,
vllm/v1/spec_decode/eagle.py
+251 -18
@@ -1,5 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import ast
4 +from dataclasses import replace
5 from typing import Optional
6
7 import numpy as np
@@ -17,6 +19,8 @@ from vllm.model_executor.models import supports_multimodal
19 from vllm.model_executor.models.llama_eagle3 import Eagle3LlamaForCausalLM
20 from vllm.utils import is_pin_memory_available
21 from vllm.v1.attention.backends.flash_attn import FlashAttentionMetadata
22 +from vllm.v1.attention.backends.tree_attn import (TreeAttentionMetadata,
23 + TreeAttentionMetadataBuilder)
24 from vllm.v1.attention.backends.utils import CommonAttentionMetadata
25 from vllm.v1.kv_cache_interface import KVCacheConfig
26 from vllm.v1.sample.metadata import SamplingMetadata
@@ -74,18 +78,52 @@ class EagleProposer:
78 (self.max_num_tokens, self.hidden_size),
79 dtype=self.dtype,
80 device=device)
77 - # We need +1 here because the arange is used to set query_start_loc,
78 - # which has one more element than batch_size.
79 - self.arange = torch.arange(vllm_config.scheduler_config.max_num_seqs +
80 - 1,
81 - device=device,
82 - dtype=torch.int32)
81 +
82 + max_batch_size = vllm_config.scheduler_config.max_num_seqs
83 + self.arange = torch.arange(
84 + # We need +1 here because the arange is used to set query_start_loc,
85 + # which has one more element than batch_size.
86 + max_batch_size + 1,
87 + device=device,
88 + dtype=torch.int32,
89 + )
90
91 self.inputs_embeds = torch.zeros(
92 (self.max_num_tokens, self.hidden_size),
93 dtype=self.dtype,
94 device=device)
95
96 + # Parse the speculative token tree.
97 + spec_token_tree = self.speculative_config.speculative_token_tree
98 + self.tree_choices: list[tuple[int,
99 + ...]] = ast.literal_eval(spec_token_tree)
100 + tree_depth = len(self.tree_choices[-1])
101 + # Precompute per-level properties of the tree.
102 + num_drafts_per_level = [0] * tree_depth
103 + for node in self.tree_choices:
104 + num_drafts_per_level[len(node) - 1] += 1
105 + self.cu_drafts_per_level = [num_drafts_per_level[0]]
106 + self.child_drafts_per_level = [num_drafts_per_level[0]]
107 + for level in range(1, tree_depth):
108 + self.cu_drafts_per_level.append(self.cu_drafts_per_level[-1] +
109 + num_drafts_per_level[level])
110 + self.child_drafts_per_level.append(num_drafts_per_level[level] //
111 + num_drafts_per_level[level - 1])
112 + # Find the first level where the tree branches off into one or more
113 + # children.
114 + self.first_branching_level = None
115 + for level in range(tree_depth):
116 + if self.cu_drafts_per_level[level] > level + 1:
117 + self.first_branching_level = level
118 + break
119 + # Precompute draft position offsets in flattened tree.
120 + self.tree_draft_pos_offsets = torch.arange(
121 + 1,
122 + len(self.tree_choices) + 1,
123 + device=device,
124 + dtype=torch.int32,
125 + ).repeat(max_batch_size, 1)
126 +
127 def propose(
128 self,
129 # [num_tokens]
@@ -120,11 +158,9 @@ class EagleProposer:
158 assert self.runner is not None
159
160 # FIXME: need to consider multiple kv_cache_groups
123 - attn_metadata = self.runner.attn_metadata_builders[0].build(
124 - common_prefix_len=0,
125 - common_attn_metadata=common_attn_metadata,
126 - fast_build=True,
127 - )
161 + attn_metadata = self.runner.attn_metadata_builders[
162 + 0].build_for_drafting(common_attn_metadata=common_attn_metadata,
163 + draft_index=0)
164
165 # At this moment, we assume all eagle layers belong to the same KV
166 # cache group, thus using the same attention metadata.
@@ -167,6 +203,22 @@ class EagleProposer:
203 last_hidden_states, hidden_states = ret_hidden_states
204 sample_hidden_states = last_hidden_states[last_token_indices]
205 logits = self.model.compute_logits(sample_hidden_states, None)
206 + positions = target_positions[last_token_indices]
207 + hidden_states = hidden_states[last_token_indices]
208 + if self.first_branching_level == 0:
209 + # Branching has occurred at the root level. Draft using tree
210 + # attention.
211 + draft_token_ids_list = self.propose_tree(
212 + tree_root_level=0,
213 + batch_size=batch_size,
214 + logits=logits,
215 + positions=positions,
216 + hidden_states=hidden_states,
217 + common_attn_metadata=common_attn_metadata,
218 + )
219 + # [batch_size, num_tree_tokens]
220 + return torch.cat(draft_token_ids_list, dim=1)
221 +
222 draft_token_ids = logits.argmax(dim=-1)
223
224 # Early exit if there is only one draft token to be generated.
@@ -178,16 +230,15 @@ class EagleProposer:
230 # one layer. Adapt this code to support multiple layers once
231 # there's a multi-layer MTP module.
232
181 - # Currently FlashAttention is the only backend that supports
182 - # multi-token eagle spec decode. This is because the code below
233 + # Currently, only FlashAttention and TreeAttention support multi-token
234 + # eagle spec decode. This is because the code below
235 # makes assumptions about attn_metadata attributes available.
184 - assert isinstance(attn_metadata, FlashAttentionMetadata)
236 + assert isinstance(attn_metadata,
237 + (FlashAttentionMetadata, TreeAttentionMetadata))
238
239 # Generate the remaining draft tokens.
240 draft_token_ids_list = [draft_token_ids]
241
189 - positions = target_positions[last_token_indices]
190 - hidden_states = hidden_states[last_token_indices]
242 if self.use_cuda_graph and \
243 batch_size <= self.cudagraph_batch_sizes[-1]:
244 input_batch_size = self.vllm_config.pad_for_cudagraph(batch_size)
@@ -196,7 +247,7 @@ class EagleProposer:
247 attn_metadata.num_actual_tokens = batch_size
248 attn_metadata.max_query_len = 1
249 attn_metadata.query_start_loc = self.arange[:batch_size + 1]
199 - for _ in range(self.num_speculative_tokens - 1):
250 + for token_index in range(self.num_speculative_tokens - 1):
251 # Update the inputs.
252 # cast to int32 is crucial when eagle model is compiled.
253 # tensor.argmax() returns int64 by default.
@@ -265,7 +316,20 @@ class EagleProposer:
316 logits = self.model.compute_logits(last_hidden_states[:batch_size],
317 None)
318
268 - # TODO(wenlong): get more than one token for tree attention
319 + if self.first_branching_level == token_index + 1:
320 + # Branching has occurred. The remaining tokens are drafted
321 + # using tree attention.
322 + draft_token_ids_list += self.propose_tree(
323 + tree_root_level=token_index + 1,
324 + batch_size=batch_size,
325 + logits=logits,
326 + positions=positions,
327 + hidden_states=hidden_states,
328 + common_attn_metadata=common_attn_metadata,
329 + )
330 + # [batch_size, num_tree_tokens]
331 + return torch.cat(draft_token_ids_list, dim=1)
332 +
333 draft_token_ids = logits.argmax(dim=-1)
334 draft_token_ids_list.append(draft_token_ids)
335
@@ -273,6 +337,175 @@ class EagleProposer:
337 draft_token_ids = torch.stack(draft_token_ids_list, dim=1)
338 return draft_token_ids
339
340 + def propose_tree(
341 + self,
342 + tree_root_level: int,
343 + batch_size: int,
344 + # [num_tokens, vocab_size]
345 + logits: torch.Tensor,
346 + # [num_tokens]
347 + positions: torch.Tensor,
348 + # [num_tokens, hidden_size]
349 + hidden_states: torch.Tensor,
350 + common_attn_metadata: CommonAttentionMetadata,
351 + ) -> list[torch.Tensor]:
352 + tree_attn_metadata_builder = self.runner.attn_metadata_builders[0]
353 + assert isinstance(tree_attn_metadata_builder,
354 + TreeAttentionMetadataBuilder)
355 +
356 + total_num_drafts = self.cu_drafts_per_level[tree_root_level]
357 + level_num_drafts = total_num_drafts
358 + # Sample a draft token for each child at the tree root level.
359 + num_children = self.child_drafts_per_level[tree_root_level]
360 + if num_children == 1:
361 + draft_token_ids = logits.argmax(dim=-1).view(batch_size, -1)
362 + else:
363 + draft_token_ids = torch.topk(logits, num_children,
364 + dim=-1).indices.view(batch_size, -1)
365 + draft_token_ids_list = [draft_token_ids]
366 + draft_hidden_states = hidden_states.view(batch_size, 1, -1)
367 +
368 + # Initialize empty tensors for concatenation with the level outputs.
369 + tree_input_ids = torch.empty(0,
370 + device=self.input_ids.device,
371 + dtype=self.input_ids.dtype)
372 + tree_positions = torch.empty(0,
373 + device=self.positions.device,
374 + dtype=self.positions.dtype)
375 + tree_hidden_states = torch.empty(0,
376 + device=self.hidden_states.device,
377 + dtype=self.hidden_states.dtype)
378 + # Precompute the draft token positions.
379 + flattened_draft_positions = (
380 + positions.view(batch_size, -1) +
381 + self.tree_draft_pos_offsets[:batch_size, :])
382 + tree_depth = len(self.cu_drafts_per_level)
383 + for level in range(tree_root_level, tree_depth - 1):
384 + # Get draft positions for RoPE.
385 + draft_positions = positions + (level + 1)
386 + exceeds_max_model_len = (positions +
387 + total_num_drafts) >= self.max_model_len
388 + # Mask out the position ids that exceed the max model length.
389 + # Otherwise, we may get out-of-range error in RoPE.
390 + clamped_draft_positions = torch.where(
391 + exceeds_max_model_len,
392 + 0,
393 + draft_positions,
394 + )
395 + if level_num_drafts > 1:
396 + # Repeat the positions for each draft at this level.
397 + draft_positions = clamped_draft_positions.repeat_interleave(
398 + level_num_drafts).reshape(batch_size, -1)
399 +
400 + if num_children > 1:
401 + # Repeat draft hidden states for each child.
402 + draft_hidden_states = draft_hidden_states.repeat_interleave(
403 + num_children, dim=1)
404 +
405 + # Concatenate the draft tokens, positions, and hidden states.
406 + tree_input_ids = torch.cat([tree_input_ids, draft_token_ids],
407 + dim=1)
408 + tree_positions = torch.cat([tree_positions, draft_positions],
409 + dim=1)
410 + tree_hidden_states = torch.cat(
411 + [tree_hidden_states, draft_hidden_states], dim=1)
412 +
413 + # Build new attention metadata for the next level of drafts.
414 + # This is necessary to support tree attention.
415 + query_len = total_num_drafts - tree_root_level
416 + common_attn_metadata = replace(
417 + common_attn_metadata,
418 + query_start_loc=query_len * self.arange[:batch_size + 1],
419 + seq_lens=common_attn_metadata.seq_lens + level_num_drafts,
420 + num_actual_tokens=batch_size * query_len,
421 + max_query_len=query_len,
422 + )
423 + attn_metadata = tree_attn_metadata_builder.build_for_drafting(
424 + common_attn_metadata=common_attn_metadata,
425 + draft_index=tree_root_level + 1,
426 + )
427 +
428 + # Apply new attention metadata to all layers.
429 + per_layer_attn_metadata = {}
430 + for layer_name in self.attn_layer_names:
431 + per_layer_attn_metadata[layer_name] = attn_metadata
432 +
433 + # Consider max model length.
434 + attn_metadata.max_seq_len = min(attn_metadata.max_seq_len,
435 + self.max_model_len)
436 + # For the requests that exceed the max model length, we set the
437 + # sequence length to 1 to minimize their overheads in attention.
438 + attn_metadata.seq_lens.masked_fill_(exceeds_max_model_len, 1)
439 +
440 + # Compute the slot mapping.
441 + query_positions = flattened_draft_positions[:, level:level +
442 + query_len]
443 + block_numbers = query_positions // self.block_size
444 + block_ids = attn_metadata.block_table.gather(dim=1,
445 + index=block_numbers)
446 + slot_mapping = (block_ids * self.block_size +
447 + query_positions % self.block_size)
448 + # Mask out the slot mappings that exceed the max model length.
449 + # Otherwise, the KV cache will be inadvertently updated with the
450 + # padding tokens.
451 + slot_mapping[exceeds_max_model_len] = PADDING_SLOT_ID
452 + attn_metadata.slot_mapping = slot_mapping.view(-1)
453 +
454 + # Copy inputs to buffer for cudagraph.
455 + num_tokens = attn_metadata.num_actual_tokens
456 + input_ids = tree_input_ids.view(-1)
457 + self.input_ids[:num_tokens] = input_ids
458 + self.positions[:num_tokens] = tree_positions.view(-1)
459 + self.hidden_states[:num_tokens] = tree_hidden_states.view(
460 + num_tokens, -1)
461 +
462 + if self.use_cuda_graph and \
463 + num_tokens <= self.cudagraph_batch_sizes[-1]:
464 + num_input_tokens = self.vllm_config.pad_for_cudagraph(
465 + num_tokens)
466 + else:
467 + num_input_tokens = num_tokens
468 + # Run the model.
469 + with set_forward_context(per_layer_attn_metadata,
470 + self.vllm_config,
471 + num_tokens=num_input_tokens):
472 + last_hidden_states, hidden_states = self.model(
473 + input_ids=self.input_ids[:num_input_tokens],
474 + positions=self.positions[:num_input_tokens],
475 + hidden_states=self.hidden_states[:num_input_tokens],
476 + inputs_embeds=None,
477 + )
478 +
479 + # Get the output hidden states for the draft tokens.
480 + draft_hidden_states = hidden_states[:num_tokens].view(
481 + batch_size, query_len, -1)[:, -level_num_drafts:]
482 + draft_last_hidden_states = last_hidden_states[:num_tokens].view(
483 + batch_size, query_len, -1)[:, -level_num_drafts:]
484 +
485 + # Get the output logits for the draft tokens.
486 + logits = self.model.compute_logits(
487 + draft_last_hidden_states.reshape(batch_size * level_num_drafts,
488 + -1),
489 + None,
490 + )
491 +
492 + # Sample a draft token for each child at the next tree level.
493 + num_children = self.child_drafts_per_level[level + 1]
494 + if num_children == 1:
495 + draft_token_ids = logits.argmax(dim=-1).view(batch_size, -1)
496 + else:
497 + draft_token_ids = torch.topk(logits, num_children,
498 + dim=-1).indices.view(
499 + batch_size, -1)
500 + draft_token_ids_list.append(draft_token_ids)
501 +
502 + # Update the # drafts counters for the next tree level.
503 + level_num_drafts = self.cu_drafts_per_level[level +
504 + 1] - total_num_drafts
505 + total_num_drafts = self.cu_drafts_per_level[level + 1]
506 +
507 + return draft_token_ids_list
508 +
509 def prepare_inputs(
510 self,
511 common_attn_metadata: CommonAttentionMetadata,