@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / b1728c1e66

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 4, 02:30 UTC

[Attention][Cleanup] Remove tree attention (#42121)

Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>

Lucas Wilkinson committed May 8, 2026 at 21:36 UTC b1728c1e660d35f0443605c779cf0643e7737cec

11 files changed +6 -1399

.buildkite/intel_jobs/test-intel.yaml

+1 -1

@@ -61,5 +61,5 @@ steps:

61 pytest -v -s v1/worker --ignore=v1/worker/test_gpu_model_runner.py --ignore=v1/worker/test_worker_memory_snapshot.py &&

62 pytest -v -s v1/structured_output &&

63 pytest -v -s v1/test_serial_utils.py &&

64 - pytest -v -s v1/spec_decode --ignore=v1/spec_decode/test_max_len.py --ignore=v1/spec_decode/test_tree_attention.py --ignore=v1/spec_decode/test_speculators_eagle3.py --ignore=v1/spec_decode/test_acceptance_length.py &&

64 + pytest -v -s v1/spec_decode --ignore=v1/spec_decode/test_max_len.py --ignore=v1/spec_decode/test_speculators_eagle3.py --ignore=v1/spec_decode/test_acceptance_length.py &&

65 pytest -v -s v1/kv_connector/unit --ignore=v1/kv_connector/unit/test_multi_connector.py --ignore=v1/kv_connector/unit/test_example_connector.py --ignore=v1/kv_connector/unit/test_lmcache_integration.py --ignore=v1/kv_connector/unit/test_hf3fs_client.py --ignore=v1/kv_connector/unit/test_hf3fs_connector.py --ignore=v1/kv_connector/unit/test_hf3fs_metadata_server.py'

docs/design/attention_backends.md

-1

@@ -179,7 +179,6 @@ Priority is **1 = highest** (tried first).

179 | `ROCM_AITER_FA` | | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32 | 64, 128, 256 | ❌ | ✅ | ❌ | ❌ | Decoder | N/A |

180 | `ROCM_AITER_UNIFIED_ATTN` | | fp16, bf16 | `auto` | %16 | Any | ✅ | ❌ | ✅ | ❌ | All | N/A |

181 | `ROCM_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | 32, 64, 80, 96, 128, 160, 192, 224, 256 | ❌ | ✅ | ✅ | ❌ | Decoder, Encoder, Encoder Only | N/A |

182 -| `TREE_ATTN` | | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | 32, 64, 96, 128, 160, 192, 224, 256 | ❌ | ❌ | ❌ | ❌ | Decoder | Any |

182 | `TRITON_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `int8_per_token_head`, `fp8_per_token_head` | %16 | Any | ✅ | ❌ | ✅ | ❌ | All | Any |

183 | `TURBOQUANT` | | fp16, bf16 | `turboquant_k8v4`, `turboquant_4bit_nc`, `turboquant_k3v4_nc`, `turboquant_3bit_nc` | 16, 32, 64, 128 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | Any |

184

tests/utils.py

+1 -1

@@ -1865,7 +1865,7 @@ def has_module_attribute(module_name, attribute_name):

1865

1866 def get_attn_backend_list_based_on_platform() -> list[str]:

1867 if current_platform.is_cuda():

1868 - return ["FLASH_ATTN", "TRITON_ATTN", "TREE_ATTN"]

1868 + return ["FLASH_ATTN", "TRITON_ATTN"]

1869 elif current_platform.is_rocm():

1870 attn_backend_list = ["TRITON_ATTN"]

1871 try:

tests/v1/attention/test_attention_backends.py

-1

@@ -36,7 +36,6 @@ BACKENDS_TO_TEST = [

36 AttentionBackendEnum.FLASHINFER,

37 AttentionBackendEnum.FLEX_ATTENTION,

38 AttentionBackendEnum.TRITON_ATTN,

39 - AttentionBackendEnum.TREE_ATTN,

39 "FLEX_ATTENTION_SLOW",

40 ]

tests/v1/e2e/spec_decode/test_spec_decode.py

-5

@@ -386,11 +386,6 @@ def _run_eagle_correctness(

386 Compare the outputs of an original LLM and a speculative LLM

387 which should be the same when using eagle speculative decoding.

388 """

389 - if attn_backend == "TREE_ATTN":

390 - pytest.skip(

391 - "TREE_ATTN is flaky in the test disable for now until it can be "

392 - "resolved (see https://github.com/vllm-project/vllm/issues/22922)"

393 - )

389 if model_impl == "transformers":

390 import transformers

391 from packaging.version import Version

tests/v1/spec_decode/test_eagle.py

-160

@@ -49,7 +49,6 @@ def _create_proposer(

49 method: str,

50 num_speculative_tokens: int,

51 attention_backend: str | None = None,

52 - speculative_token_tree: list[tuple[int, ...]] | None = None,

52 parallel_drafting: bool = False,

53 ) -> EagleProposer:

54 # Method-dependent setup

@@ -75,18 +74,12 @@ def _create_proposer(

74 trust_remote_code=(method == "dflash"),

75 )

78 - spec_token_tree_str = None

79 - if speculative_token_tree is not None:

80 - assert num_speculative_tokens == len(speculative_token_tree)

81 - spec_token_tree_str = str(speculative_token_tree)

82 -

77 speculative_config = SpeculativeConfig(

78 target_model_config=model_config,

79 target_parallel_config=ParallelConfig(),

80 model=draft_model_dir,

81 method=method,

82 num_speculative_tokens=num_speculative_tokens,

89 - speculative_token_tree=spec_token_tree_str,

83 parallel_drafting=parallel_drafting,

84 )

85 if parallel_drafting:

@@ -844,12 +837,6 @@ def test_propose(method, attn_backend, num_speculative_tokens, monkeypatch):

837 "multi-token eagle spec decode on current platform"

838 )

839

847 - if attn_backend == "TREE_ATTN":

848 - pytest.skip(

849 - "TREE_ATTN is tested separately in test_propose_tree"

850 - "because it requires special input mocking."

851 - )

852 -

840 if attn_backend == "ROCM_AITER_FA" and current_platform.is_rocm():

841 monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1")

842

@@ -955,10 +942,6 @@ def test_propose(method, attn_backend, num_speculative_tokens, monkeypatch):

942 attn_metadata_builder_cls, _ = try_get_attention_backend(

943 AttentionBackendEnum.TRITON_ATTN

944 )

958 - elif attn_backend == "TREE_ATTN":

959 - attn_metadata_builder_cls, _ = try_get_attention_backend(

960 - AttentionBackendEnum.TREE_ATTN

961 - )

945 elif attn_backend == "ROCM_AITER_FA":

946 attn_metadata_builder_cls, _ = try_get_attention_backend(

947 AttentionBackendEnum.ROCM_AITER_FA

@@ -1014,149 +997,6 @@ def test_propose(method, attn_backend, num_speculative_tokens, monkeypatch):

997 assert torch.equal(result, expected_tokens)

998

999

1017 -@pytest.mark.parametrize(

1018 - "spec_token_tree",

1019 - [

1020 - [(0,)], # A single token

1021 - [(0,), (0, 0), (0, 0, 0)], # Chain

1022 - [(0,), (1,), (2,)], # Parallel

1023 - [(0,), (1,), (2,), (0, 0), (0, 1), (1, 0), (1, 1), (2, 0), (2, 1)], # Tree

1024 - ],

1025 -)

1026 -def test_propose_tree(spec_token_tree):

1027 - # Get GPU device.

1028 - device = torch.device(DEVICE_TYPE)

1029 -

1030 - # Setup test parameters.

1031 - batch_size = 2

1032 - seq_len_1 = 5

1033 - seq_len_2 = 3

1034 - total_tokens = seq_len_1 + seq_len_2

1035 - vocab_size = 100

1036 - seq_lens = [seq_len_1, seq_len_2]

1037 - num_speculative_tokens = len(spec_token_tree)

1038 -

1039 - # Create proposer first so we can use its actual hidden_size.

1040 - proposer = _create_proposer(

1041 - "eagle",

1042 - num_speculative_tokens,

1043 - speculative_token_tree=spec_token_tree,

1044 - )

1045 - # Get the hidden_size from the proposer to ensure consistency.

1046 - hidden_size = proposer.hidden_size

1047 -

1048 - # Helper to create deterministic logits that will produce specific tokens

1049 - def create_deterministic_logits(token_ids, k: int):

1050 - logits = torch.full((batch_size, vocab_size), -100.0, device=device)

1051 - for i, token_id in enumerate(token_ids):

1052 - # Assign decreasing values to the k, consecutive, tokens.

1053 - for j in range(k):

1054 - logits[i, token_id + j] = 100.0 - j

1055 - return logits

1056 -

1057 - # Mock a model that returns deterministic logits.

1058 - base_token_ids = torch.tensor([42, 60], dtype=torch.int64, device=device)

1059 -

1060 - # Skip loading the model and replace it with a mock that returns

1061 - # deterministic outputs.

1062 - model_mock = mock.MagicMock()

1063 -

1064 - # Mock the model forward calls.

1065 - forward_returns = [

1066 - (

1067 - torch.zeros(total_tokens, hidden_size, device=device),

1068 - torch.zeros(total_tokens, hidden_size, device=device),

1069 - )

1070 - ]

1071 - for cu_num_drafts in proposer.cu_drafts_per_level:

1072 - h_logits = torch.zeros(batch_size * cu_num_drafts, hidden_size, device=device)

1073 - h_states = torch.zeros(batch_size * cu_num_drafts, hidden_size, device=device)

1074 - forward_returns.append((h_logits, h_states))

1075 - model_mock.side_effect = forward_returns

1076 -

1077 - # Mock the compute_logits calls.

1078 - cu_num_drafts_tensor = torch.tensor(

1079 - [0] + proposer.cu_drafts_per_level, dtype=torch.int32, device=device

1080 - )

1081 - logits_returns = []

1082 - for level, num_children in enumerate(proposer.child_drafts_per_level):

1083 - token_ids = base_token_ids + cu_num_drafts_tensor[level]

1084 - level_num_drafts = cu_num_drafts_tensor[level + 1] - cu_num_drafts_tensor[level]

1085 - level_logits = []

1086 - for i in range(level_num_drafts // num_children):

1087 - level_logits.append(

1088 - create_deterministic_logits(token_ids + i * num_children, num_children)

1089 - )

1090 - logits_returns.append(torch.stack(level_logits, dim=1))

1091 - model_mock.compute_logits.side_effect = logits_returns

1092 -

1093 - # Assign the mock to the proposer

1094 - proposer.model = model_mock

1095 -

1096 - # Assign draft attn_layer_names since load_model is not invoked

1097 - proposer._draft_attn_layer_names = {"layer.0"}

1098 -

1099 - # Get the tree attention metadata builder.

1100 - attn_metadata_builder_cls, _ = try_get_attention_backend(

1101 - AttentionBackendEnum.TREE_ATTN

1102 - )

1103 - attn_metadata_builder = attn_metadata_builder_cls(

1104 - kv_cache_spec=create_standard_kv_cache_spec(proposer.vllm_config),

1105 - layer_names=proposer._draft_attn_layer_names,

1106 - vllm_config=proposer.vllm_config,

1107 - device=device,

1108 - )

1109 -

1110 - # Mock runner and draft_attn_groups for attention metadata building.

1111 - proposer.runner = mock.MagicMock()

1112 - mock_attn_group = mock.MagicMock()

1113 - mock_attn_group.get_metadata_builder.return_value = attn_metadata_builder

1114 - mock_attn_group.layer_names = list(proposer._draft_attn_layer_names)

1115 - mock_attn_group.kv_cache_spec = attn_metadata_builder.kv_cache_spec

1116 - proposer.draft_attn_groups = [mock_attn_group]

1117 -

1118 - # Setup inputs for the proposer.

1119 - target_token_ids = torch.randint(0, vocab_size, (total_tokens,), device=device)

1120 - target_positions = torch.cat(

1121 - [torch.arange(seq_len_1, device=device), torch.arange(seq_len_2, device=device)]

1122 - )

1123 - target_hidden_states = torch.randn(total_tokens, hidden_size, device=device)

1124 - next_token_ids = torch.randint(

1125 - 0, vocab_size, (batch_size,), dtype=torch.int32, device=device

1126 - )

1127 - batch_spec = BatchSpec(

1128 - seq_lens=seq_lens,

1129 - query_lens=seq_lens,

1130 - )

1131 - common_attn_metadata = create_common_attn_metadata(

1132 - batch_spec,

1133 - block_size=BLOCK_SIZE,

1134 - device=device,

1135 - )

1136 - sampling_metadata = mock.MagicMock()

1137 -

1138 - # Propose draft tokens.

1139 - result = proposer.propose(

1140 - target_token_ids=target_token_ids,

1141 - target_positions=target_positions,

1142 - target_hidden_states=target_hidden_states,

1143 - next_token_ids=next_token_ids,

1144 - token_indices_to_sample=None,

1145 - common_attn_metadata=common_attn_metadata,

1146 - sampling_metadata=sampling_metadata,

1147 - )

1148 - assert result.shape == (batch_size, num_speculative_tokens)

1149 -

1150 - # The tokens are expected to be consecutive integers starting

1151 - # from the base token IDs.

1152 - expected_tokens = base_token_ids[:, None] + torch.arange(

1153 - num_speculative_tokens, dtype=torch.int64, device=device

1154 - )

1155 -

1156 - # Verify that the draft tokens match our expectations.

1157 - assert torch.equal(result, expected_tokens)

1158 -

1159 -

1000 def test_set_inputs_first_pass_dflash():

1001 """

1002 Test for DFlash set_inputs_first_pass.

tests/v1/spec_decode/test_tree_attention.py deleted

-506

@@ -1,506 +0,0 @@

1 -# SPDX-License-Identifier: Apache-2.0

2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 -

4 -import math

5 -

6 -import pytest

7 -import torch

8 -

9 -from tests.v1.attention.utils import (

10 - create_standard_kv_cache_spec,

11 - create_vllm_config,

12 - try_backend_includes_kv_cache_update,

13 - try_get_attention_backend,

14 -)

15 -from vllm.config import ParallelConfig, SpeculativeConfig

16 -from vllm.platforms import current_platform

17 -from vllm.utils.torch_utils import set_random_seed

18 -from vllm.v1.attention.backend import CommonAttentionMetadata

19 -from vllm.v1.attention.backends.fa_utils import is_flash_attn_varlen_func_available

20 -from vllm.v1.attention.backends.registry import AttentionBackendEnum

21 -

22 -DEVICE_TYPE = current_platform.device_type

23 -

24 -if not is_flash_attn_varlen_func_available():

25 - pytest.skip(

26 - "This test requires flash_attn_varlen_func, but it's not available.",

27 - allow_module_level=True,

28 - )

29 -

30 -# --------------------------------------------------------------------------- #

31 -# KV cache layout adaptation

32 -# --------------------------------------------------------------------------- #

33 -# Two KV cache layouts exist across backends:

34 -#

35 -# Flash layout: (2, num_blocks, block_size, num_kv_heads, head_size)

36 -# - dim 0 separates key (index 0) and value (index 1)

37 -# - Used by: FLASH_ATTN, TREE_ATTN, ROCM_AITER_FA, ROCM_ATTN

38 -#

39 -# Block layout: (num_blocks, 2, block_size, num_kv_heads, head_size)

40 -# - dim 1 separates key (index 0) and value (index 1)

41 -# - Used by: TRITON_ATTN

42 -#

43 -# The test creates KV caches in flash layout (the canonical format used by

44 -# tree attention). When a reference backend needs block layout we transpose

45 -# dims 0 and 1.

46 -#

47 -# Note: ROCM_ATTN uses flash layout for storage but its forward path calls

48 -# PagedAttention.split_kv_cache which reinterprets the raw memory as paged

49 -# layout (num_blocks, num_kv_heads, head_size//x, block_size, x). This is

50 -# a view-level incompatibility, not a transpose - see the TODO in

51 -# _get_available_reference_backends for details.

52 -#

53 -# TODO: Replace this mapping with a `KV_CACHE_LAYOUT` class attribute on each

54 -# AttentionImpl so the layout is self-documented by the backend itself, e.g.:

55 -# class TritonAttentionImpl(AttentionImpl):

56 -# KV_CACHE_LAYOUT = "block"

57 -# --------------------------------------------------------------------------- #

58 -

59 -_BLOCK_KV_LAYOUT_BACKENDS = frozenset(

60 - {

61 - AttentionBackendEnum.TRITON_ATTN,

62 - }

63 -)

64 -

65 -# Backends whose do_kv_cache_update requires engine-level state (e.g.

66 -# ForwardContext) that is not available in this test harness, but whose

67 -# KV cache is flash layout and can be written with reshape_and_cache_flash.

68 -# When a backend is listed here, forward_attention() bypasses

69 -# do_kv_cache_update and writes directly to the cache.

70 -_NEEDS_DIRECT_CACHE_UPDATE = frozenset(

71 - {

72 - AttentionBackendEnum.ROCM_AITER_FA,

73 - }

74 -)

75 -

76 -# Backends with known test-harness incompatibilities - see the TODOs

77 -# inside _get_available_reference_backends for details.

78 -_INCOMPATIBLE_REFERENCE_BACKENDS = frozenset(

79 - {

80 - AttentionBackendEnum.ROCM_AITER_FA,

81 - AttentionBackendEnum.ROCM_ATTN,

82 - }

83 -)

84 -

85 -

86 -def _adapt_kv_cache_for_backend(

87 - kv_cache: torch.Tensor,

88 - backend: AttentionBackendEnum,

89 -) -> torch.Tensor:

90 - """Convert kv_cache from flash layout ``(2, num_blocks, ...)`` to block

91 - layout ``(num_blocks, 2, ...)`` if the backend requires it. Returns the

92 - original tensor unchanged when no conversion is needed."""

93 - if backend in _BLOCK_KV_LAYOUT_BACKENDS:

94 - return kv_cache.transpose(0, 1).contiguous()

95 - return kv_cache

96 -

97 -

98 -def _get_platform_default_backend() -> AttentionBackendEnum:

99 - """Ask the platform what backend it would auto-select at runtime."""

100 - from vllm.v1.attention.selector import AttentionSelectorConfig

101 -

102 - config = AttentionSelectorConfig(

103 - block_size=32,

104 - kv_cache_dtype="auto",

105 - use_mla=False,

106 - use_sparse=False,

107 - head_size=128,

108 - dtype=torch.bfloat16,

109 - )

110 - backend_path = current_platform.get_attn_backend_cls(

111 - selected_backend=None,

112 - attn_selector_config=config,

113 - )

114 - for backend in AttentionBackendEnum:

115 - try:

116 - if backend.get_path() == backend_path:

117 - return backend

118 - except ValueError:

119 - continue

120 - raise RuntimeError(

121 - f"Platform returned backend path '{backend_path}' "

122 - f"that doesn't match any AttentionBackendEnum member."

123 - )

124 -

125 -

126 -def _get_available_reference_backends() -> list[AttentionBackendEnum]:

127 - """Collect all reference backends the current platform can run.

128 -

129 - On CUDA this is just FLASH_ATTN. On ROCm this includes the platform

130 - default plus every backend the hardware supports, so the test validates

131 - tree attention against all of them.

132 - """

133 - if current_platform.is_rocm():

134 - backends: list[AttentionBackendEnum] = []

135 -

136 - # 1. Whatever the platform would auto-select at runtime.

137 - default_backend = _get_platform_default_backend()

138 - if default_backend not in _INCOMPATIBLE_REFERENCE_BACKENDS:

139 - backends.append(default_backend)

140 -

141 - # 2. TRITON_ATTN - always available on ROCm.

142 - if AttentionBackendEnum.TRITON_ATTN not in backends:

143 - backends.append(AttentionBackendEnum.TRITON_ATTN)

144 -

145 - # TODO: Enable ROCM_ATTN. Its forward path uses

146 - # PagedAttention.split_kv_cache which reinterprets the raw

147 - # cache memory as paged layout:

148 - # key: (num_blocks, num_kv_heads, head_size//x, block_size, x)

149 - # value: (num_blocks, num_kv_heads, head_size, block_size)

150 - # Tree attention writes prefix data in NHD flash layout, so the

151 - # same bytes produce completely different values when read in

152 - # paged format. Supporting ROCM_ATTN would require writing

153 - # prefix data via PagedAttention.write_to_paged_cache into a

154 - # separate paged-format KV cache.

155 -

156 - # TODO: Enable ROCM_AITER_FA. Its metadata builder reads head

157 - # counts from the model config at construction time and

158 - # allocates extend_workspace with those dimensions. The test

159 - # uses independent head count parameters (num_heads=2/4,

160 - # num_kv_heads=2) that don't match the model config

161 - # (Llama-3-8B: 32 q heads, 8 kv heads), causing a head count

162 - # mismatch in flash_attn_varlen_func during extend_forward.

163 - # Fixing this requires either matching test head counts to the

164 - # model config or decoupling the builder from model config

165 - # head geometry. The direct cache update path

166 - # (_NEEDS_DIRECT_CACHE_UPDATE) is already in place for when

167 - # this is resolved.

168 -

169 - return backends

170 -

171 - # CUDA: flash attention.

172 - return [AttentionBackendEnum.FLASH_ATTN]

173 -

174 -

175 -class MockAttentionLayer(torch.nn.Module):

176 - _q_scale = torch.tensor(1.0, dtype=torch.float32, device=DEVICE_TYPE)

177 - _k_scale = torch.tensor(1.0, dtype=torch.float32, device=DEVICE_TYPE)

178 - _v_scale = torch.tensor(1.0, dtype=torch.float32, device=DEVICE_TYPE)

179 - layer_name = "mock_layer"

180 -

181 - def __init__(self):

182 - super().__init__()

183 -

184 - def forward(self, x):

185 - return x

186 -

187 -

188 -def forward_attention(

189 - q: torch.Tensor,

190 - k: torch.Tensor,

191 - v: torch.Tensor,

192 - kv_cache: torch.Tensor,

193 - block_table: torch.Tensor,

194 - slot_mapping: torch.Tensor,

195 - seqlen_k: int,

196 - backend: AttentionBackendEnum,

197 - spec_token_tree: str | None = None,

198 - num_spec_tokens: int = 0,

199 -) -> torch.Tensor:

200 - """Run a single attention forward pass through the given backend.

201 -

202 - ``kv_cache`` is expected in **flash layout**

203 - ``(2, num_blocks, block_size, num_kv_heads, head_size)``.

204 - It is automatically converted when the target backend needs a

205 - different layout.

206 - """

207 - batch_size, q_len, num_heads, dim_per_head = q.shape

208 - num_kv_heads = k.shape[-2]

209 - # Initialize the query and KV sequence lengths.

210 - query_start_loc = q_len * torch.arange(

211 - batch_size + 1, device=q.device, dtype=torch.int32

212 - )

213 - query_lens = torch.diff(query_start_loc)

214 - seq_lens = torch.full(

215 - (batch_size,),

216 - seqlen_k,

217 - device=q.device,

218 - dtype=torch.int32,

219 - )

220 - context_lens = seq_lens - query_lens

221 - max_seq_len = int(seq_lens.max())

222 - max_query_len = q_len

223 - num_actual_tokens = query_start_loc[-1]

224 -

225 - softmax_scale = q.shape[-1] ** (-0.5)

226 - layer = MockAttentionLayer()

227 -

228 - # Build common metadata.

229 - model_name = "meta-llama/Meta-Llama-3-8B"

230 - builder_cls, impl_cls = try_get_attention_backend(backend)

231 - vllm_config = create_vllm_config(model_name=model_name, max_model_len=max(seq_lens))

232 - if spec_token_tree is not None:

233 - # Create speculative config if token tree is specified.

234 - vllm_config.speculative_config = SpeculativeConfig(

235 - target_model_config=vllm_config.model_config,

236 - target_parallel_config=ParallelConfig(),

237 - model=model_name,

238 - method="eagle",

239 - num_speculative_tokens=num_spec_tokens,

240 - speculative_token_tree=spec_token_tree,

241 - )

242 - kv_cache_spec = create_standard_kv_cache_spec(vllm_config)

243 - builder = builder_cls(kv_cache_spec, [], vllm_config, q.device)

244 - seq_lens_cpu = seq_lens.cpu()

245 - common_attn_metadata = CommonAttentionMetadata(

246 - query_start_loc=query_start_loc,

247 - query_start_loc_cpu=query_start_loc.cpu(),

248 - seq_lens=seq_lens,

249 - seq_lens_cpu_upper_bound=seq_lens_cpu,

250 - _seq_lens_cpu=seq_lens_cpu,

251 - _num_computed_tokens_cpu=context_lens.cpu(),

252 - num_reqs=batch_size,

253 - num_actual_tokens=num_actual_tokens,

254 - max_query_len=max_query_len,

255 - max_seq_len=max_seq_len,

256 - block_table_tensor=block_table,

257 - slot_mapping=slot_mapping,

258 - )

259 -

260 - # Build attention metadata.

261 - attn_metadata = builder.build(

262 - common_prefix_len=0,

263 - common_attn_metadata=common_attn_metadata,

264 - )

265 -

266 - # Initialize the backend implementation.

267 - instance = impl_cls(

268 - num_heads=num_heads,

269 - head_size=dim_per_head,

270 - scale=softmax_scale,

271 - num_kv_heads=num_kv_heads,

272 - alibi_slopes=None,

273 - sliding_window=None,

274 - kv_cache_dtype="auto",

275 - )

276 -

277 - # Adapt KV cache layout for this backend.

278 - adapted_kv_cache = _adapt_kv_cache_for_backend(kv_cache, backend)

279 -

280 - # Run forward pass and return output.

281 - query = q.view(-1, num_heads, dim_per_head)

282 - key = k.view(-1, num_kv_heads, dim_per_head)

283 - value = v.view(-1, num_kv_heads, dim_per_head)

284 - output = torch.empty_like(query)

285 - if not try_backend_includes_kv_cache_update(backend):

286 - if backend in _NEEDS_DIRECT_CACHE_UPDATE:

287 - # This backend's do_kv_cache_update requires engine-level

288 - # ForwardContext that isn't available in this test harness.

289 - # Write directly using reshape_and_cache_flash since the

290 - # KV cache layout is identical (flash layout, unbind on dim 0).

291 - key_cache, value_cache = adapted_kv_cache.unbind(0)

292 - torch.ops._C_cache_ops.reshape_and_cache_flash(

293 - key,

294 - value,

295 - key_cache,

296 - value_cache,

297 - attn_metadata.slot_mapping,

298 - "auto",

299 - layer._k_scale,

300 - layer._v_scale,

301 - )

302 - else:

303 - instance.do_kv_cache_update(

304 - layer=layer,

305 - key=key,

306 - value=value,

307 - kv_cache=adapted_kv_cache,

308 - slot_mapping=attn_metadata.slot_mapping,

309 - )

310 - return instance.forward(

311 - layer=layer,

312 - query=query,

313 - key=key,

314 - value=value,

315 - kv_cache=adapted_kv_cache.clone(),

316 - attn_metadata=attn_metadata,

317 - output=output,

318 - )

319 -

320 -

321 -@pytest.mark.parametrize(

322 - "reference_backend",

323 - _get_available_reference_backends(),

324 - ids=lambda b: b.name,

325 -)

326 -def test_tree_attn_correctness(

327 - reference_backend: AttentionBackendEnum,

328 -) -> None:

329 - set_random_seed(42)

330 -

331 - device = "cuda"

332 - tree_attn_masks = {

333 - # Chain.

334 - "[(0,), (0, 0), (0, 0, 0)]": torch.tensor(

335 - [

336 - [1, 0, 0, 0],

337 - [1, 1, 0, 0],

338 - [1, 1, 1, 0],

339 - [1, 1, 1, 1],

340 - ],

341 - device=device,

342 - dtype=torch.int32,

343 - ),

344 - # Tree.

345 - "[(0,), (1,), (0, 0), (0, 1), (1, 0), (1, 1)]": torch.tensor(

346 - [

347 - [1, 0, 0, 0, 0, 0, 0],

348 - [1, 1, 0, 0, 0, 0, 0],

349 - [1, 0, 1, 0, 0, 0, 0],

350 - [1, 1, 0, 1, 0, 0, 0],

351 - [1, 1, 0, 0, 1, 0, 0],

352 - [1, 0, 1, 0, 0, 1, 0],

353 - [1, 0, 1, 0, 0, 0, 1],

354 - ],

355 - device=device,

356 - dtype=torch.int32,

357 - ),

358 - }

359 -

360 - dim_per_head = 128

361 - num_kv_heads = 2

362 - block_size = 32

363 - max_sequence_length = 8192

364 - randomize_blocks = True

365 - for batch_size in [1, 16, 32]:

366 - for num_heads in [2, 4]:

367 - for sequence_position in [16, 1024, 2048]:

368 - for spec_token_tree, tree_attn_mask in tree_attn_masks.items():

369 - # Assert that the number of heads is divisible

370 - # by the number of KV heads.

371 - assert num_heads % num_kv_heads == 0

372 -

373 - # Initialize q, k, and v.

374 - tree_size_q = tree_attn_mask.shape[0]

375 - seqlen_k = sequence_position + tree_size_q

376 - q = torch.randn(

377 - (batch_size, tree_size_q, num_heads, dim_per_head),

378 - device=device,

379 - dtype=torch.bfloat16,

380 - )

381 - k = torch.randn(

382 - (batch_size, tree_size_q, num_kv_heads, dim_per_head),

383 - device=device,

384 - dtype=torch.bfloat16,

385 - )

386 - v = torch.randn(

387 - (batch_size, tree_size_q, num_kv_heads, dim_per_head),

388 - device=device,

389 - dtype=torch.bfloat16,

390 - )

391 -

392 - # KV cache in flash layout - the canonical format for

393 - # tree attention. forward_attention() handles conversion

394 - # when needed.

395 - assert max_sequence_length % block_size == 0

396 - max_blocks_per_batch = max_sequence_length // block_size

397 - kv_cache = torch.randn(

398 - (

399 - 2,

400 - batch_size * max_blocks_per_batch,

401 - block_size,

402 - num_kv_heads,

403 - dim_per_head,

404 - ),

405 - device=q.device,

406 - dtype=torch.bfloat16,

407 - )

408 - num_alloc_blocks_per_batch = math.ceil(seqlen_k / block_size)

409 - block_table = torch.zeros(

410 - (batch_size, max_blocks_per_batch),

411 - device=q.device,

412 - dtype=torch.int32,

413 - )

414 - block_ids = torch.arange(

415 - 0,

416 - batch_size * num_alloc_blocks_per_batch,

417 - device=q.device,

418 - dtype=torch.int32,

419 - )

420 - if randomize_blocks:

421 - # Randomize the block ids.

422 - block_ids = block_ids[torch.randperm(block_ids.numel())]

423 - block_table[:, :num_alloc_blocks_per_batch] = block_ids.view(

424 - -1, num_alloc_blocks_per_batch

425 - )

426 -

427 - # Set up the slot mapping for the input KVs.

428 - tree_positions = sequence_position + torch.arange(

429 - 0,

430 - tree_size_q,

431 - device=q.device,

432 - dtype=torch.int64,

433 - ).repeat(batch_size, 1)

434 - tree_slot_mapping = _gen_slot_mapping(

435 - tree_positions, block_table, block_size

436 - )

437 -

438 - # Compute attention for the tree.

439 - tree_attn_output = forward_attention(

440 - q=q,

441 - k=k,

442 - v=v,

443 - kv_cache=kv_cache,

444 - block_table=block_table,

445 - slot_mapping=tree_slot_mapping,

446 - seqlen_k=seqlen_k,

447 - backend=AttentionBackendEnum.TREE_ATTN,

448 - spec_token_tree=spec_token_tree,

449 - num_spec_tokens=tree_size_q - 1,

450 - ).view(batch_size, -1, num_heads, dim_per_head)

451 -

452 - # Verify each branch against the reference backend.

453 - for q_index in range(tree_size_q):

454 - # Get the q, k, and v for the branch.

455 - branch_mask = tree_attn_mask[q_index, :]

456 - branch_indices = torch.nonzero(branch_mask, as_tuple=True)[0]

457 - q_len = branch_indices.shape[0]

458 - q_branch = q[:, branch_indices]

459 - k_branch = k[:, branch_indices]

460 - v_branch = v[:, branch_indices]

461 -

462 - # Setup slot mapping for the branch.

463 - branch_positions = sequence_position + torch.arange(

464 - 0,

465 - q_len,

466 - device=q.device,

467 - dtype=torch.int64,

468 - ).repeat(batch_size, 1)

469 - branch_slot_mapping = _gen_slot_mapping(

470 - branch_positions, block_table, block_size

471 - )

472 -

473 - # Reference attention for this branch.

474 - ref_output = forward_attention(

475 - q=q_branch,

476 - k=k_branch,

477 - v=v_branch,

478 - kv_cache=kv_cache,

479 - block_table=block_table,

480 - slot_mapping=branch_slot_mapping,

481 - seqlen_k=sequence_position + q_len,

482 - backend=reference_backend,

483 - ).view(batch_size, -1, num_heads, dim_per_head)

484 -

485 - # Compare the outputs.

486 - assert torch.allclose(

487 - tree_attn_output[:, branch_indices],

488 - ref_output,

489 - atol=7.81e-3,

490 - ), (

491 - f"outputs are not close for "

492 - f"reference_backend: {reference_backend.name}, "

493 - f"batch_size: {batch_size}, "

494 - f"num_heads: {num_heads}, "

495 - f"sequence_position: {sequence_position}, "

496 - f"tree_attn_mask: {tree_attn_mask}, "

497 - f"q_index: {q_index}."

498 - )

499 -

500 -

501 -def _gen_slot_mapping(

502 - positions: torch.Tensor, block_table: torch.Tensor, block_size: int

503 -):

504 - block_indices = positions // block_size

505 - blocks = block_table.gather(dim=1, index=block_indices)

506 - return (blocks * block_size + positions % block_size).view(-1)

vllm/config/speculative.py

+4 -21

@@ -1,7 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 -import ast

4 import copy

5 from typing import TYPE_CHECKING, Any, Literal, get_args

@@ -145,9 +144,6 @@ class SpeculativeConfig:

144 provided. Defaults to 1."""

145

146 # Alternative drafting strategies

148 - speculative_token_tree: str | None = None

149 - """Specifies the tree structure for speculative token generation.

150 - """

147 parallel_drafting: bool = False

148 """Enable parallel drafting, where all speculative tokens are generated

149 in parallel rather than sequentially. This can improve performance but

@@ -739,23 +735,10 @@ class SpeculativeConfig:

735 f" must be divisible by {n_predict=}"

736 )

737

742 - if self.speculative_token_tree is None:

743 - if self.num_speculative_tokens is None:

744 - raise ValueError(

745 - "A speculative model was provided, but neither "

746 - "`speculative_token_tree` nor `num_speculative_tokens` "

747 - "was provided"

748 - )

749 -

750 - # Generate chain of tokens.

751 - self.speculative_token_tree = str(

752 - [(i + 1) * (0,) for i in range(self.num_speculative_tokens)]

753 - )

754 - else:

755 - # Sort the token tree breadth-first.

756 - tree_choices = ast.literal_eval(self.speculative_token_tree)

757 - self.speculative_token_tree = str(

758 - sorted(tree_choices, key=lambda t: (len(t), t))

738 + if self.num_speculative_tokens is None:

739 + raise ValueError(

740 + "A speculative model was provided, but "

741 + "`num_speculative_tokens` was not provided"

742 )

743

744 self.draft_tensor_parallel_size = (

vllm/v1/attention/backends/registry.py

-1

@@ -76,7 +76,6 @@ class AttentionBackendEnum(Enum, metaclass=_AttentionBackendEnumMeta):

76 FLASH_ATTN_MLA = "vllm.v1.attention.backends.mla.flashattn_mla.FlashAttnMLABackend"

77 NO_ATTENTION = "vllm.v1.attention.backends.no_attention.NoAttentionBackend"

78 FLEX_ATTENTION = "vllm.v1.attention.backends.flex_attention.FlexAttentionBackend"

79 - TREE_ATTN = "vllm.v1.attention.backends.tree_attn.TreeAttentionBackend"

79 ROCM_AITER_UNIFIED_ATTN = (

80 "vllm.v1.attention.backends.rocm_aiter_unified_attn."

81 "RocmAiterUnifiedAttentionBackend"

vllm/v1/attention/backends/tree_attn.py deleted

-488

@@ -1,488 +0,0 @@

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 ClassVar

8 -

9 -import torch

10 -

11 -from vllm import _custom_ops as ops

12 -from vllm.config import VllmConfig

13 -from vllm.config.cache import CacheDType

14 -from vllm.logger import init_logger

15 -from vllm.v1.attention.backend import (

16 - AttentionBackend,

17 - AttentionImpl,

18 - AttentionMetadataBuilder,

19 - AttentionType,

20 - CommonAttentionMetadata,

21 - MultipleOf,

22 -)

23 -from vllm.v1.attention.backends.utils import (

24 - split_decodes_and_prefills,

25 -)

26 -from vllm.v1.attention.ops.triton_unified_attention import unified_attention

27 -from vllm.v1.kv_cache_interface import AttentionSpec

28 -

29 -logger = init_logger(__name__)

30 -

31 -

32 -class TreeAttentionBackend(AttentionBackend):

33 - supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]

34 - supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [

35 - "auto",

36 - "float16",

37 - "bfloat16",

38 - ]

39 - forward_includes_kv_cache_update: bool = False

40 -

41 - @staticmethod

42 - def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:

43 - return [MultipleOf(16)]

44 -

45 - @classmethod

46 - def get_supported_head_sizes(cls) -> list[int]:

47 - return [32, 64, 96, 128, 160, 192, 224, 256]

48 -

49 - @staticmethod

50 - def get_name() -> str:

51 - return "TREE_ATTN"

52 -

53 - @staticmethod

54 - def get_impl_cls() -> type["TreeAttentionImpl"]:

55 - return TreeAttentionImpl

56 -

57 - @staticmethod

58 - def get_kv_cache_shape(

59 - num_blocks: int,

60 - block_size: int,

61 - num_kv_heads: int,

62 - head_size: int,

63 - cache_dtype_str: str = "auto",

64 - ) -> tuple[int, ...]:

65 - if block_size % 16 != 0:

66 - raise ValueError("Block size must be a multiple of 16.")

67 - return (2, num_blocks, block_size, num_kv_heads, head_size)

68 -

69 - @staticmethod

70 - def get_builder_cls() -> type["TreeAttentionMetadataBuilder"]:

71 - return TreeAttentionMetadataBuilder

72 -

73 - @staticmethod

74 - def use_cascade_attention(*args, **kwargs) -> bool:

75 - return False

76 -

77 -

78 -@dataclass

79 -class TreeAttentionMetadata:

80 - num_actual_tokens: int # Number of tokens excluding padding.

81 - max_query_len: int

82 - query_start_loc: torch.Tensor

83 - max_seq_len: int

84 - seq_lens: torch.Tensor

85 - block_table: torch.Tensor

86 - slot_mapping: torch.Tensor

87 -

88 - num_prefill_tokens: int = 0

89 - num_decode_tokens: int = 0

90 - num_prefills: int = 0

91 - num_decodes: int = 0

92 -

93 - # Precomputed (on CPU in the builder) max_query_len and max_seq_len for

94 - # the prefill-only and decode-only sub-batches. Used by the properties

95 - # below to avoid a GPU->CPU sync via `.max().item()` on every forward.

96 - max_query_len_prefill: int = 0

97 - max_seq_len_prefill: int = 0

98 - max_query_len_decode: int = 0

99 - max_seq_len_decode: int = 0

100 -

101 - tree_attn_bias: torch.Tensor | None = None

102 -

103 - # Cached Prefill/decode metadata.

104 - _cached_prefill_metadata: "TreeAttentionMetadata | None" = None

105 - _cached_decode_metadata: "TreeAttentionMetadata | None" = None

106 -

107 - @property

108 - def prefill_metadata(self) -> "TreeAttentionMetadata | None":

109 - if self.num_prefills == 0:

110 - return None

111 -

112 - if self._cached_prefill_metadata is not None:

113 - # Recover cached prefill-phase attention

114 - # metadata structure

115 - return self._cached_prefill_metadata

116 -

117 - q_start_loc = self.query_start_loc[self.num_decodes :]

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=self.max_query_len_prefill,

123 - query_start_loc=q_start_loc - q_start_loc[0],

124 - max_seq_len=self.max_seq_len_prefill,

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) -> "TreeAttentionMetadata | None":

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 - kv_seqlens = self.seq_lens[: self.num_decodes]

143 - # Construct & cache decode-phase attention metadata structure

144 - self._cached_decode_metadata = TreeAttentionMetadata(

145 - num_actual_tokens=self.num_decode_tokens,

146 - max_query_len=self.max_query_len_decode,

147 - query_start_loc=q_start_loc,

148 - max_seq_len=self.max_seq_len_decode,

149 - seq_lens=kv_seqlens,

150 - block_table=self.block_table[: self.num_decodes],

151 - slot_mapping=self.slot_mapping[: self.num_decode_tokens],

152 - tree_attn_bias=self.tree_attn_bias,

153 - )

154 - return self._cached_decode_metadata

155 -

156 -

157 -class TreeAttentionMetadataBuilder(AttentionMetadataBuilder[TreeAttentionMetadata]):

158 - def __init__(

159 - self,

160 - kv_cache_spec: AttentionSpec,

161 - layer_names: list[str],

162 - vllm_config: VllmConfig,

163 - device: torch.device,

164 - ):

165 - super().__init__(kv_cache_spec, layer_names, vllm_config, device)

166 -

167 - self.block_size = kv_cache_spec.block_size

168 -

169 - spec_config = vllm_config.speculative_config

170 - spec_token_tree: str | None = None

171 - if spec := spec_config:

172 - spec_token_tree = spec.speculative_token_tree

173 - tree_choices: list[tuple[int, ...]] = (

174 - ast.literal_eval(spec_token_tree) if spec_token_tree is not None else [(0,)]

175 - )

176 - # Construct the tree attention bias.

177 - depth_counts = _get_depth_counts(tree_choices)

178 - self.tree_attn_bias = _prepare_tree_attn_bias(

179 - tree_choices,

180 - depth_counts,

181 - dtype=torch.float32,

182 - device=device,

183 - )

184 -

185 - self.reorder_batch_threshold = self.tree_attn_bias.shape[0]

186 -

187 - def build(

188 - self,

189 - common_prefix_len: int,

190 - common_attn_metadata: CommonAttentionMetadata,

191 - fast_build: bool = False,

192 - ) -> TreeAttentionMetadata:

193 - decode_threshold = self.tree_attn_bias.shape[0]

194 - num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (

195 - split_decodes_and_prefills(

196 - common_attn_metadata, decode_threshold=decode_threshold

197 - )

198 - )

199 -

200 - num_actual_tokens = common_attn_metadata.num_actual_tokens

201 - q_start_loc = common_attn_metadata.query_start_loc

202 - max_query_len = common_attn_metadata.max_query_len

203 - kv_seqlens = common_attn_metadata.seq_lens

204 - max_seq_len = common_attn_metadata.max_seq_len

205 - block_table = common_attn_metadata.block_table_tensor

206 - slot_mapping = common_attn_metadata.slot_mapping

207 -

208 - # Precompute prefill/decode sub-batch max_query_len / max_seq_len on

209 - # CPU so the prefill_metadata / decode_metadata properties don't need

210 - # a GPU->CPU sync via `.max().item()` on every forward.

211 - # Prefer `seq_lens_cpu_upper_bound` over the (deprecated)

212 - # `seq_lens_cpu` property: the upper bound is precise for prefill

213 - # rows and optimistic-but-safe for decode rows (workspace sizing

214 - # from `max()` is fine with an over-estimate), and avoids the

215 - # `seq_lens.to("cpu")` sync the property would fall through to in

216 - # async-spec-decode mode. The draft-attention path (eagle

217 - # speculator) doesn't populate it; fall back to the batch-wide

218 - # `max_seq_len` as a safe upper bound for both sub-batches.

219 - q_start_loc_cpu = common_attn_metadata.query_start_loc_cpu

220 - seq_lens_cpu = common_attn_metadata.seq_lens_cpu_upper_bound

221 - if num_prefills > 0:

222 - q_seqlens_p = torch.diff(q_start_loc_cpu[num_decodes:])

223 - max_query_len_prefill = int(q_seqlens_p.max())

224 - max_seq_len_prefill = (

225 - int(seq_lens_cpu[num_decodes:].max())

226 - if seq_lens_cpu is not None

227 - else max_seq_len

228 - )

229 - else:

230 - max_query_len_prefill = 0

231 - max_seq_len_prefill = 0

232 - if num_decodes > 0:

233 - q_seqlens_d = torch.diff(q_start_loc_cpu[: num_decodes + 1])

234 - max_query_len_decode = int(q_seqlens_d.max())

235 - max_seq_len_decode = (

236 - int(seq_lens_cpu[:num_decodes].max())

237 - if seq_lens_cpu is not None

238 - else max_seq_len

239 - )

240 - else:

241 - max_query_len_decode = 0

242 - max_seq_len_decode = 0

243 -

244 - return TreeAttentionMetadata(

245 - num_actual_tokens=num_actual_tokens,

246 - num_prefill_tokens=num_prefill_tokens,

247 - num_decode_tokens=num_decode_tokens,

248 - num_prefills=num_prefills,

249 - num_decodes=num_decodes,

250 - max_query_len=max_query_len,

251 - query_start_loc=q_start_loc,

252 - max_seq_len=max_seq_len,

253 - seq_lens=kv_seqlens,

254 - block_table=block_table,

255 - slot_mapping=slot_mapping,

256 - max_query_len_prefill=max_query_len_prefill,

257 - max_seq_len_prefill=max_seq_len_prefill,

258 - max_query_len_decode=max_query_len_decode,

259 - max_seq_len_decode=max_seq_len_decode,

260 - tree_attn_bias=self.tree_attn_bias,

261 - )

262 -

263 - def build_for_drafting(

264 - self,

265 - common_attn_metadata: CommonAttentionMetadata,

266 - draft_index: int,

267 - ) -> TreeAttentionMetadata:

268 - # Cache the original tree attention bias.

269 - orig_tree_attn_bias = self.tree_attn_bias

270 -

271 - if draft_index == 0:

272 - # Use prefill for drafting at the root level.

273 - self.tree_attn_bias = torch.empty(0)

274 - else:

275 - # Slice the tree attention bias for drafting. Exclude

276 - # the root level.

277 - start, end = 1, 1 + common_attn_metadata.max_query_len

278 - self.tree_attn_bias = self.tree_attn_bias[start:end, start:end].contiguous()

279 -

280 - # Build attention bias.

281 - attn_metadata = self.build(0, common_attn_metadata, fast_build=True)

282 -

283 - # Reset the tree attention bias to the original value.

284 - self.tree_attn_bias = orig_tree_attn_bias

285 - return attn_metadata

286 -

287 -

288 -def _get_depth_counts(sorted_tree_choices: list[tuple[int, ...]]) -> list[int]:

289 - # Count the number of choices at each depth of the tree.

290 - depth_counts = []

291 - prev_depth = 0

292 - for path in sorted_tree_choices:

293 - depth = len(path)

294 - if depth != prev_depth:

295 - depth_counts.append(0)

296 - depth_counts[depth - 1] += 1

297 - prev_depth = depth

298 - return depth_counts

299 -

300 -

301 -def _prepare_tree_attn_bias(

302 - sorted_tree_choices: list[tuple[int, ...]],

303 - depth_counts: list[int],

304 - dtype: torch.dtype | None,

305 - device: torch.device | None,

306 -) -> torch.Tensor:

307 - # +1 comes from the additional root node.

308 - tree_len = len(sorted_tree_choices) + 1

309 - tree_attn_mask = torch.full(

310 - (tree_len, tree_len), -torch.inf, device=device, dtype=dtype

311 - )

312 -

313 - # Set diagonal to all zeros. Each token should

314 - # attend to itself.

315 - mask_val = 0

316 - for i in range(tree_len):

317 - tree_attn_mask[i, i] = mask_val

318 -

319 - # Set root to all zeros. All tokens attend to it.

320 - tree_attn_mask[:, 0] = mask_val

321 -

322 - # Set all ancestors to zeros.

323 - start = 0

324 - for i in range(len(depth_counts)):

325 - for j in range(depth_counts[i]):

326 - cur_tree_choice = sorted_tree_choices[start + j]

327 - # Retrieve ancestor position.

328 - if len(cur_tree_choice) == 1:

329 - continue

330 - ancestor_idx = []

331 - for c in range(len(cur_tree_choice) - 1):

332 - ancestor_idx.append(

333 - sorted_tree_choices.index(cur_tree_choice[: c + 1]) + 1

334 - )

335 - tree_attn_mask[j + start + 1, ancestor_idx] = mask_val

336 - start += depth_counts[i]

337 - return tree_attn_mask

338 -

339 -

340 -class TreeAttentionImpl(AttentionImpl):

341 - def __init__(

342 - self,

343 - num_heads: int,

344 - head_size: int,

345 - scale: float,

346 - num_kv_heads: int,

347 - alibi_slopes: list[float] | None,

348 - sliding_window: int | None,

349 - kv_cache_dtype: str,

350 - logits_soft_cap: float | None = None,

351 - attn_type: AttentionType = AttentionType.DECODER,

352 - kv_sharing_target_layer_name: str | None = None,

353 - ) -> None:

354 - self.num_heads = num_heads

355 - self.head_size = head_size

356 - self.scale = float(scale)

357 - self.num_kv_heads = num_kv_heads

358 - self.num_queries_per_kv = self.num_heads // self.num_kv_heads

359 - self.kv_cache_dtype = kv_cache_dtype

360 - self.kv_sharing_target_layer_name = kv_sharing_target_layer_name

361 - if alibi_slopes is not None:

362 - alibi_slopes = torch.tensor(alibi_slopes, dtype=torch.float32)

363 - self.alibi_slopes = alibi_slopes

364 - if logits_soft_cap is None:

365 - # Setting logits_soft_cap to 0 means no soft cap.

366 - logits_soft_cap = 0

367 - self.logits_soft_cap = logits_soft_cap

368 - if sliding_window is None:

369 - self.sliding_window = (-1, -1)

370 - else:

371 - self.sliding_window = (sliding_window - 1, 0)

372 -

373 - if attn_type != AttentionType.DECODER:

374 - raise NotImplementedError(

375 - "Encoder self-attention and "

376 - "encoder/decoder cross-attention "

377 - "are not implemented for "

378 - "TreeAttentionImpl."

379 - )

380 -

381 - def do_kv_cache_update(

382 - self,

383 - layer: torch.nn.Module,

384 - key: torch.Tensor,

385 - value: torch.Tensor,

386 - kv_cache: torch.Tensor,

387 - slot_mapping: torch.Tensor,

388 - ) -> None:

389 - key_cache, value_cache = kv_cache.unbind(0)

390 -

391 - # Reshape the input keys and values and store them in the cache.

392 - # NOTE(woosuk): Here, key and value are padded while slot_mapping is

393 - # not padded. However, we don't need to do key[:num_actual_tokens]

394 - # and value[:num_actual_tokens] because the reshape_and_cache_flash

395 - # op uses the slot_mapping's shape to determine the number of

396 - # actual tokens.

397 - ops.reshape_and_cache_flash(

398 - key,

399 - value,

400 - key_cache,

401 - value_cache,

402 - slot_mapping,

403 - self.kv_cache_dtype,

404 - layer._k_scale,

405 - layer._v_scale,

406 - )

407 -

408 - def forward(

409 - self,

410 - layer: torch.nn.Module,

411 - query: torch.Tensor,

412 - key: torch.Tensor,

413 - value: torch.Tensor,

414 - kv_cache: torch.Tensor,

415 - attn_metadata: TreeAttentionMetadata,

416 - output: torch.Tensor,

417 - output_scale: torch.Tensor | None = None,

418 - output_block_scale: torch.Tensor | None = None,

419 - ) -> torch.Tensor:

420 - """Forward pass with TreeAttention.

421 -

422 - Args:

423 - query: shape = [num_tokens, num_heads, head_size]

424 - key: shape = [num_tokens, num_kv_heads, head_size]

425 - value: shape = [num_tokens, num_kv_heads, head_size]

426 - kv_cache: shape =

427 - [2, num_blocks, block_size, num_kv_heads, head_size]

428 - attn_metadata: Metadata for attention.

429 - Returns:

430 - shape = [num_tokens, num_heads * head_size]

431 - """

432 - if output_scale is not None or output_block_scale is not None:

433 - raise NotImplementedError(

434 - "fused output quantization is not yet supported for TreeAttentionImpl"

435 - )

436 -

437 - if attn_metadata is None:

438 - # Profiling run.

439 - return output.fill_(0)

440 -

441 - key_cache, value_cache = kv_cache.unbind(0)

442 -

443 - num_actual_tokens = attn_metadata.num_actual_tokens

444 - num_decode_tokens = attn_metadata.num_decode_tokens

445 - descale_shape = (attn_metadata.query_start_loc.shape[0] - 1, key.shape[1])

446 - if prefill_meta := attn_metadata.prefill_metadata:

447 - unified_attention(

448 - q=query[num_decode_tokens:num_actual_tokens],

449 - k=key_cache,

450 - v=value_cache,

451 - out=output[num_decode_tokens:num_actual_tokens],

452 - cu_seqlens_q=prefill_meta.query_start_loc,

453 - max_seqlen_q=prefill_meta.max_query_len,

454 - seqused_k=prefill_meta.seq_lens,

455 - max_seqlen_k=prefill_meta.max_seq_len,

456 - softmax_scale=self.scale,

457 - causal=True,

458 - alibi_slopes=self.alibi_slopes,

459 - window_size=self.sliding_window,

460 - block_table=prefill_meta.block_table,

461 - softcap=self.logits_soft_cap,

462 - q_descale=None, # Not supported

463 - k_descale=layer._k_scale.expand(descale_shape),

464 - v_descale=layer._v_scale.expand(descale_shape),

465 - )

466 -

467 - if decode_meta := attn_metadata.decode_metadata:

468 - unified_attention(

469 - q=query[:num_decode_tokens],

470 - k=key_cache,

471 - v=value_cache,

472 - out=output[:num_decode_tokens],

473 - cu_seqlens_q=decode_meta.query_start_loc,

474 - max_seqlen_q=decode_meta.max_query_len,

475 - seqused_k=decode_meta.seq_lens,

476 - max_seqlen_k=decode_meta.max_seq_len,

477 - softmax_scale=self.scale,

478 - causal=True,

479 - alibi_slopes=self.alibi_slopes,

480 - qq_bias=decode_meta.tree_attn_bias,

481 - window_size=self.sliding_window,

482 - block_table=decode_meta.block_table,

483 - softcap=self.logits_soft_cap,

484 - q_descale=None, # Not supported

485 - k_descale=layer._k_scale.expand(descale_shape),

486 - v_descale=layer._v_scale.expand(descale_shape),

487 - )

488 - return output

vllm/v1/spec_decode/llm_base_proposer.py

-214

@@ -1,6 +1,5 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 -import ast

3 from importlib.util import find_spec

4 from typing import Any, cast

@@ -29,10 +28,6 @@ from vllm.platforms import current_platform

28 from vllm.utils.platform_utils import is_pin_memory_available

29 from vllm.v1.attention.backend import CommonAttentionMetadata

30 from vllm.v1.attention.backends.registry import AttentionBackendEnum

32 -from vllm.v1.attention.backends.tree_attn import (

33 - TreeAttentionMetadata,

34 - TreeAttentionMetadataBuilder,

35 -)

31 from vllm.v1.attention.backends.triton_attn import TritonAttentionMetadata

32 from vllm.v1.cudagraph_dispatcher import CudagraphDispatcher

33 from vllm.v1.kv_cache_interface import KVCacheConfig, UniformTypeKVCacheSpecs

@@ -284,29 +279,6 @@ class SpecDecodeBaseProposer:

279

280 self.allowed_attn_types = tuple(rocm_types)

281

287 - # Parse the speculative token tree.

288 - spec_token_tree = self.speculative_config.speculative_token_tree

289 - assert spec_token_tree is not None

290 - self.tree_choices: list[tuple[int, ...]] = ast.literal_eval(spec_token_tree)

291 - tree_depth = len(self.tree_choices[-1])

292 - # Precompute per-level properties of the tree.

293 - num_drafts_per_level = [0] * tree_depth

294 - for node in self.tree_choices:

295 - num_drafts_per_level[len(node) - 1] += 1

296 - self.cu_drafts_per_level = [num_drafts_per_level[0]]

297 - self.child_drafts_per_level = [num_drafts_per_level[0]]

298 - for level in range(1, tree_depth):

299 - self.cu_drafts_per_level.append(

300 - self.cu_drafts_per_level[-1] + num_drafts_per_level[level]

301 - )

302 - self.child_drafts_per_level.append(

303 - num_drafts_per_level[level] // num_drafts_per_level[level - 1]

304 - )

305 - # Precompute draft position offsets in flattened tree.

306 - self.tree_draft_pos_offsets = torch.arange(

307 - 1, len(self.tree_choices) + 1, device=device, dtype=torch.int32

308 - ).repeat(self.max_batch_size, 1)

309 -

282 def _raise_if_padded_drafter_batch_disabled(self):

283 if self.speculative_config.disable_padded_drafter_batch:

284 raise NotImplementedError(

@@ -512,20 +484,6 @@ class SpecDecodeBaseProposer:

484 # (which read via _get_positions) use the correct values.

485 self.positions[:batch_size] = positions

486

515 - if any(isinstance(md, TreeAttentionMetadata) for md in per_group_attn_metadata):

516 - # Draft using tree attention - requires full logits for top-k

517 - logits = self.model.compute_logits(sample_hidden_states)

518 - draft_token_ids_list = self.propose_tree(

519 - batch_size=batch_size,

520 - logits=logits,

521 - positions=positions,

522 - hidden_states=hidden_states,

523 - common_attn_metadata=common_attn_metadata,

524 - slot_mappings=slot_mappings,

525 - )

526 - # [batch_size, num_tree_tokens]

527 - return torch.cat(draft_token_ids_list, dim=1)

528 -

487 draft_token_ids = self._greedy_sample(sample_hidden_states)

488

489 if self.allowed_attn_types is not None:

@@ -1014,178 +972,6 @@ class SpecDecodeBaseProposer:

972 num_rejected_tokens_gpu,

973 )

974

1017 - def propose_tree(

1018 - self,

1019 - batch_size: int,

1020 - # [num_tokens, vocab_size]

1021 - logits: torch.Tensor,

1022 - # [num_tokens]

1023 - positions: torch.Tensor,

1024 - # [num_tokens, hidden_size]

1025 - hidden_states: torch.Tensor,

1026 - common_attn_metadata: CommonAttentionMetadata,

1027 - slot_mappings: dict[str, torch.Tensor]

1028 - | list[dict[str, torch.Tensor]]

1029 - | None = None,

1030 - ) -> list[torch.Tensor]:

1031 - tree_attn_metadata_builder = self.draft_attn_groups[0].get_metadata_builder()

1032 - assert isinstance(tree_attn_metadata_builder, TreeAttentionMetadataBuilder)

1033 -

1034 - total_num_drafts = self.cu_drafts_per_level[0]

1035 - level_num_drafts = total_num_drafts

1036 - # Sample a draft token for each child at the tree root level.

1037 - num_children = self.child_drafts_per_level[0]

1038 - if num_children == 1:

1039 - draft_token_ids = logits.argmax(dim=-1).view(batch_size, -1)

1040 - else:

1041 - draft_token_ids = torch.topk(logits, num_children, dim=-1).indices.view(

1042 - batch_size, -1

1043 - )

1044 - draft_token_ids_list = [draft_token_ids]

1045 - draft_hidden_states = hidden_states.view(batch_size, 1, -1)

1046 -

1047 - # Initialize empty tensors for concatenation with the level outputs.

1048 - tree_input_ids = torch.empty(

1049 - 0, device=self.input_ids.device, dtype=self.input_ids.dtype

1050 - )

1051 - tree_positions = torch.empty(

1052 - 0, device=self.positions.device, dtype=self.positions.dtype

1053 - )

1054 - tree_hidden_states = torch.empty(

1055 - 0, device=self.hidden_states.device, dtype=self.hidden_states.dtype

1056 - )

1057 - # Precompute the draft token positions.

1058 - flattened_draft_positions = (

1059 - positions.view(batch_size, -1) + self.tree_draft_pos_offsets[:batch_size, :]

1060 - )

1061 - tree_depth = len(self.cu_drafts_per_level)

1062 - for level in range(tree_depth - 1):

1063 - # Get draft positions for RoPE.

1064 - draft_positions = positions + (level + 1)

1065 - exceeds_max_model_len = (positions + total_num_drafts) >= self.max_model_len

1066 - # Mask out the position ids that exceed the max model length.

1067 - # Otherwise, we may get out-of-range error in RoPE.

1068 - draft_positions = torch.where(

1069 - exceeds_max_model_len,

1070 - 0,

1071 - draft_positions,

1072 - ).view(batch_size, -1)

1073 -

1074 - if level_num_drafts > 1:

1075 - # Repeat the positions for each draft at this level.

1076 - draft_positions = draft_positions.repeat_interleave(

1077 - level_num_drafts, dim=1

1078 - )

1079 -

1080 - if num_children > 1:

1081 - # Repeat draft hidden states for each child.

1082 - draft_hidden_states = draft_hidden_states.repeat_interleave(

1083 - num_children, dim=1

1084 - )

1085 -

1086 - # Concatenate the draft tokens, positions, and hidden states.

1087 - tree_input_ids = torch.cat([tree_input_ids, draft_token_ids], dim=1)

1088 - tree_positions = torch.cat([tree_positions, draft_positions], dim=1)

1089 - tree_hidden_states = torch.cat(

1090 - [tree_hidden_states, draft_hidden_states], dim=1

1091 - )

1092 -

1093 - # Build new attention metadata for the next level of drafts.

1094 - # This is necessary to support tree attention.

1095 - query_len = total_num_drafts

1096 - common_attn_metadata = replace(

1097 - common_attn_metadata,

1098 - query_start_loc=query_len * self.arange[: batch_size + 1],

1099 - seq_lens=common_attn_metadata.seq_lens + level_num_drafts,

1100 - num_actual_tokens=batch_size * query_len,

1101 - max_query_len=query_len,

1102 - )

1103 - attn_metadata = tree_attn_metadata_builder.build_for_drafting(

1104 - common_attn_metadata=common_attn_metadata, draft_index=level + 1

1105 - )

1106 -

1107 - # Apply new attention metadata to all draft layers.

1108 - per_layer_attn_metadata = {}

1109 - for attn_group in self.draft_attn_groups:

1110 - for layer_name in attn_group.layer_names:

1111 - per_layer_attn_metadata[layer_name] = attn_metadata

1112 -

1113 - # Consider max model length.

1114 - attn_metadata.max_seq_len = min(

1115 - attn_metadata.max_seq_len, self.max_model_len

1116 - )

1117 - # For the requests that exceed the max model length, we set the

1118 - # sequence length to 1 to minimize their overheads in attention.

1119 - attn_metadata.seq_lens.masked_fill_(exceeds_max_model_len, 1)

1120 -

1121 - # Compute the slot mapping.

1122 - block_size = tree_attn_metadata_builder.kv_cache_spec.block_size

1123 - query_positions = flattened_draft_positions[:, level : level + query_len]

1124 - block_numbers = query_positions // block_size

1125 - block_ids = attn_metadata.block_table.gather(dim=1, index=block_numbers)

1126 - slot_mapping = block_ids * block_size + query_positions % block_size

1127 - # Mask out the slot mappings that exceed the max model length.

1128 - # Otherwise, the KV cache will be inadvertently updated with the

1129 - # padding tokens.

1130 - slot_mapping[exceeds_max_model_len] = PADDING_SLOT_ID

1131 - attn_metadata.slot_mapping = slot_mapping.view(-1)

1132 -

1133 - # Copy inputs to buffer for cudagraph.

1134 - num_tokens = attn_metadata.num_actual_tokens

1135 - input_ids = tree_input_ids.view(-1)

1136 - self.input_ids[:num_tokens] = input_ids

1137 - self.positions[:num_tokens] = tree_positions.view(-1)

1138 - self.hidden_states[:num_tokens] = tree_hidden_states.view(num_tokens, -1)

1139 -

1140 - cudagraph_runtime_mode, batch_desc = self.cudagraph_dispatcher.dispatch(

1141 - num_tokens

1142 - )

1143 - num_input_tokens = batch_desc.num_tokens

1144 - # Run the model.

1145 - with set_forward_context(

1146 - per_layer_attn_metadata,

1147 - self.vllm_config,

1148 - num_tokens=num_input_tokens,

1149 - cudagraph_runtime_mode=cudagraph_runtime_mode,

1150 - slot_mapping=self._get_slot_mapping(

1151 - num_input_tokens, attn_metadata.slot_mapping

1152 - ),

1153 - ):

1154 - last_hidden_states, hidden_states = self.model(

1155 - input_ids=self.input_ids[:num_input_tokens],

1156 - positions=self.positions[:num_input_tokens],

1157 - hidden_states=self.hidden_states[:num_input_tokens],

1158 - inputs_embeds=None,

1159 - )

1160 -

1161 - # Get the output hidden states for the draft tokens.

1162 - draft_hidden_states = hidden_states[:num_tokens].view(

1163 - batch_size, query_len, -1

1164 - )[:, -level_num_drafts:]

1165 - draft_last_hidden_states = last_hidden_states[:num_tokens].view(

1166 - batch_size, query_len, -1

1167 - )[:, -level_num_drafts:]

1168 -

1169 - # Get the output logits for the draft tokens.

1170 - logits = self.model.compute_logits(

1171 - draft_last_hidden_states.reshape(batch_size * level_num_drafts, -1)

1172 - )

1173 -

1174 - # Sample a draft token for each child at the next tree level.

1175 - num_children = self.child_drafts_per_level[level + 1]

1176 - if num_children == 1:

1177 - draft_token_ids = logits.argmax(dim=-1).view(batch_size, -1)

1178 - else:

1179 - draft_token_ids = torch.topk(logits, num_children, dim=-1).indices.view(

1180 - batch_size, -1

1181 - )

1182 - draft_token_ids_list.append(draft_token_ids)

1183 -

1184 - # Update the # drafts counters for the next tree level.

1185 - level_num_drafts = self.cu_drafts_per_level[level + 1] - total_num_drafts

1186 - total_num_drafts = self.cu_drafts_per_level[level + 1]

1187 - return draft_token_ids_list

1188 -

975 def prepare_inputs(

976 self,

977 common_attn_metadata: CommonAttentionMetadata,