diff --git a/.buildkite/intel_jobs/test-intel.yaml b/.buildkite/intel_jobs/test-intel.yaml index c59be699502..336f71fdbcd 100644 --- a/.buildkite/intel_jobs/test-intel.yaml +++ b/.buildkite/intel_jobs/test-intel.yaml @@ -61,5 +61,5 @@ steps: pytest -v -s v1/worker --ignore=v1/worker/test_gpu_model_runner.py --ignore=v1/worker/test_worker_memory_snapshot.py && pytest -v -s v1/structured_output && pytest -v -s v1/test_serial_utils.py && - 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 && + 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 && 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' diff --git a/docs/design/attention_backends.md b/docs/design/attention_backends.md index e711694b3b9..3f6feecf5f7 100644 --- a/docs/design/attention_backends.md +++ b/docs/design/attention_backends.md @@ -179,7 +179,6 @@ Priority is **1 = highest** (tried first). | `ROCM_AITER_FA` | | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32 | 64, 128, 256 | ❌ | ✅ | ❌ | ❌ | Decoder | N/A | | `ROCM_AITER_UNIFIED_ATTN` | | fp16, bf16 | `auto` | %16 | Any | ✅ | ❌ | ✅ | ❌ | All | N/A | | `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 | -| `TREE_ATTN` | | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | 32, 64, 96, 128, 160, 192, 224, 256 | ❌ | ❌ | ❌ | ❌ | Decoder | Any | | `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 | | `TURBOQUANT` | | fp16, bf16 | `turboquant_k8v4`, `turboquant_4bit_nc`, `turboquant_k3v4_nc`, `turboquant_3bit_nc` | 16, 32, 64, 128 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | Any | diff --git a/tests/utils.py b/tests/utils.py index 0ded7e1d696..cff601374b0 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -1865,7 +1865,7 @@ def has_module_attribute(module_name, attribute_name): def get_attn_backend_list_based_on_platform() -> list[str]: if current_platform.is_cuda(): - return ["FLASH_ATTN", "TRITON_ATTN", "TREE_ATTN"] + return ["FLASH_ATTN", "TRITON_ATTN"] elif current_platform.is_rocm(): attn_backend_list = ["TRITON_ATTN"] try: diff --git a/tests/v1/attention/test_attention_backends.py b/tests/v1/attention/test_attention_backends.py index 41218c41f4f..46d25c017f2 100644 --- a/tests/v1/attention/test_attention_backends.py +++ b/tests/v1/attention/test_attention_backends.py @@ -36,7 +36,6 @@ BACKENDS_TO_TEST = [ AttentionBackendEnum.FLASHINFER, AttentionBackendEnum.FLEX_ATTENTION, AttentionBackendEnum.TRITON_ATTN, - AttentionBackendEnum.TREE_ATTN, "FLEX_ATTENTION_SLOW", ] diff --git a/tests/v1/e2e/spec_decode/test_spec_decode.py b/tests/v1/e2e/spec_decode/test_spec_decode.py index 2ab2245b790..e14f7681bf4 100644 --- a/tests/v1/e2e/spec_decode/test_spec_decode.py +++ b/tests/v1/e2e/spec_decode/test_spec_decode.py @@ -386,11 +386,6 @@ def _run_eagle_correctness( Compare the outputs of an original LLM and a speculative LLM which should be the same when using eagle speculative decoding. """ - if attn_backend == "TREE_ATTN": - pytest.skip( - "TREE_ATTN is flaky in the test disable for now until it can be " - "resolved (see https://github.com/vllm-project/vllm/issues/22922)" - ) if model_impl == "transformers": import transformers from packaging.version import Version diff --git a/tests/v1/spec_decode/test_eagle.py b/tests/v1/spec_decode/test_eagle.py index 462ddfdfe50..32c152bf754 100644 --- a/tests/v1/spec_decode/test_eagle.py +++ b/tests/v1/spec_decode/test_eagle.py @@ -49,7 +49,6 @@ def _create_proposer( method: str, num_speculative_tokens: int, attention_backend: str | None = None, - speculative_token_tree: list[tuple[int, ...]] | None = None, parallel_drafting: bool = False, ) -> EagleProposer: # Method-dependent setup @@ -75,18 +74,12 @@ def _create_proposer( trust_remote_code=(method == "dflash"), ) - spec_token_tree_str = None - if speculative_token_tree is not None: - assert num_speculative_tokens == len(speculative_token_tree) - spec_token_tree_str = str(speculative_token_tree) - speculative_config = SpeculativeConfig( target_model_config=model_config, target_parallel_config=ParallelConfig(), model=draft_model_dir, method=method, num_speculative_tokens=num_speculative_tokens, - speculative_token_tree=spec_token_tree_str, parallel_drafting=parallel_drafting, ) if parallel_drafting: @@ -844,12 +837,6 @@ def test_propose(method, attn_backend, num_speculative_tokens, monkeypatch): "multi-token eagle spec decode on current platform" ) - if attn_backend == "TREE_ATTN": - pytest.skip( - "TREE_ATTN is tested separately in test_propose_tree" - "because it requires special input mocking." - ) - if attn_backend == "ROCM_AITER_FA" and current_platform.is_rocm(): monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1") @@ -955,10 +942,6 @@ def test_propose(method, attn_backend, num_speculative_tokens, monkeypatch): attn_metadata_builder_cls, _ = try_get_attention_backend( AttentionBackendEnum.TRITON_ATTN ) - elif attn_backend == "TREE_ATTN": - attn_metadata_builder_cls, _ = try_get_attention_backend( - AttentionBackendEnum.TREE_ATTN - ) elif attn_backend == "ROCM_AITER_FA": attn_metadata_builder_cls, _ = try_get_attention_backend( AttentionBackendEnum.ROCM_AITER_FA @@ -1014,149 +997,6 @@ def test_propose(method, attn_backend, num_speculative_tokens, monkeypatch): assert torch.equal(result, expected_tokens) -@pytest.mark.parametrize( - "spec_token_tree", - [ - [(0,)], # A single token - [(0,), (0, 0), (0, 0, 0)], # Chain - [(0,), (1,), (2,)], # Parallel - [(0,), (1,), (2,), (0, 0), (0, 1), (1, 0), (1, 1), (2, 0), (2, 1)], # Tree - ], -) -def test_propose_tree(spec_token_tree): - # Get GPU device. - device = torch.device(DEVICE_TYPE) - - # Setup test parameters. - batch_size = 2 - seq_len_1 = 5 - seq_len_2 = 3 - total_tokens = seq_len_1 + seq_len_2 - vocab_size = 100 - seq_lens = [seq_len_1, seq_len_2] - num_speculative_tokens = len(spec_token_tree) - - # Create proposer first so we can use its actual hidden_size. - proposer = _create_proposer( - "eagle", - num_speculative_tokens, - speculative_token_tree=spec_token_tree, - ) - # Get the hidden_size from the proposer to ensure consistency. - hidden_size = proposer.hidden_size - - # Helper to create deterministic logits that will produce specific tokens - def create_deterministic_logits(token_ids, k: int): - logits = torch.full((batch_size, vocab_size), -100.0, device=device) - for i, token_id in enumerate(token_ids): - # Assign decreasing values to the k, consecutive, tokens. - for j in range(k): - logits[i, token_id + j] = 100.0 - j - return logits - - # Mock a model that returns deterministic logits. - base_token_ids = torch.tensor([42, 60], dtype=torch.int64, device=device) - - # Skip loading the model and replace it with a mock that returns - # deterministic outputs. - model_mock = mock.MagicMock() - - # Mock the model forward calls. - forward_returns = [ - ( - torch.zeros(total_tokens, hidden_size, device=device), - torch.zeros(total_tokens, hidden_size, device=device), - ) - ] - for cu_num_drafts in proposer.cu_drafts_per_level: - h_logits = torch.zeros(batch_size * cu_num_drafts, hidden_size, device=device) - h_states = torch.zeros(batch_size * cu_num_drafts, hidden_size, device=device) - forward_returns.append((h_logits, h_states)) - model_mock.side_effect = forward_returns - - # Mock the compute_logits calls. - cu_num_drafts_tensor = torch.tensor( - [0] + proposer.cu_drafts_per_level, dtype=torch.int32, device=device - ) - logits_returns = [] - for level, num_children in enumerate(proposer.child_drafts_per_level): - token_ids = base_token_ids + cu_num_drafts_tensor[level] - level_num_drafts = cu_num_drafts_tensor[level + 1] - cu_num_drafts_tensor[level] - level_logits = [] - for i in range(level_num_drafts // num_children): - level_logits.append( - create_deterministic_logits(token_ids + i * num_children, num_children) - ) - logits_returns.append(torch.stack(level_logits, dim=1)) - model_mock.compute_logits.side_effect = logits_returns - - # Assign the mock to the proposer - proposer.model = model_mock - - # Assign draft attn_layer_names since load_model is not invoked - proposer._draft_attn_layer_names = {"layer.0"} - - # Get the tree attention metadata builder. - attn_metadata_builder_cls, _ = try_get_attention_backend( - AttentionBackendEnum.TREE_ATTN - ) - attn_metadata_builder = attn_metadata_builder_cls( - kv_cache_spec=create_standard_kv_cache_spec(proposer.vllm_config), - layer_names=proposer._draft_attn_layer_names, - vllm_config=proposer.vllm_config, - device=device, - ) - - # Mock runner and draft_attn_groups for attention metadata building. - proposer.runner = mock.MagicMock() - mock_attn_group = mock.MagicMock() - mock_attn_group.get_metadata_builder.return_value = attn_metadata_builder - mock_attn_group.layer_names = list(proposer._draft_attn_layer_names) - mock_attn_group.kv_cache_spec = attn_metadata_builder.kv_cache_spec - proposer.draft_attn_groups = [mock_attn_group] - - # Setup inputs for the proposer. - target_token_ids = torch.randint(0, vocab_size, (total_tokens,), device=device) - target_positions = torch.cat( - [torch.arange(seq_len_1, device=device), torch.arange(seq_len_2, device=device)] - ) - target_hidden_states = torch.randn(total_tokens, hidden_size, device=device) - next_token_ids = torch.randint( - 0, vocab_size, (batch_size,), dtype=torch.int32, device=device - ) - batch_spec = BatchSpec( - seq_lens=seq_lens, - query_lens=seq_lens, - ) - common_attn_metadata = create_common_attn_metadata( - batch_spec, - block_size=BLOCK_SIZE, - device=device, - ) - sampling_metadata = mock.MagicMock() - - # Propose draft tokens. - result = proposer.propose( - target_token_ids=target_token_ids, - target_positions=target_positions, - target_hidden_states=target_hidden_states, - next_token_ids=next_token_ids, - token_indices_to_sample=None, - common_attn_metadata=common_attn_metadata, - sampling_metadata=sampling_metadata, - ) - assert result.shape == (batch_size, num_speculative_tokens) - - # The tokens are expected to be consecutive integers starting - # from the base token IDs. - expected_tokens = base_token_ids[:, None] + torch.arange( - num_speculative_tokens, dtype=torch.int64, device=device - ) - - # Verify that the draft tokens match our expectations. - assert torch.equal(result, expected_tokens) - - def test_set_inputs_first_pass_dflash(): """ Test for DFlash set_inputs_first_pass. diff --git a/tests/v1/spec_decode/test_tree_attention.py b/tests/v1/spec_decode/test_tree_attention.py deleted file mode 100644 index 3c126c49f8c..00000000000 --- a/tests/v1/spec_decode/test_tree_attention.py +++ /dev/null @@ -1,506 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -import math - -import pytest -import torch - -from tests.v1.attention.utils import ( - create_standard_kv_cache_spec, - create_vllm_config, - try_backend_includes_kv_cache_update, - try_get_attention_backend, -) -from vllm.config import ParallelConfig, SpeculativeConfig -from vllm.platforms import current_platform -from vllm.utils.torch_utils import set_random_seed -from vllm.v1.attention.backend import CommonAttentionMetadata -from vllm.v1.attention.backends.fa_utils import is_flash_attn_varlen_func_available -from vllm.v1.attention.backends.registry import AttentionBackendEnum - -DEVICE_TYPE = current_platform.device_type - -if not is_flash_attn_varlen_func_available(): - pytest.skip( - "This test requires flash_attn_varlen_func, but it's not available.", - allow_module_level=True, - ) - -# --------------------------------------------------------------------------- # -# KV cache layout adaptation -# --------------------------------------------------------------------------- # -# Two KV cache layouts exist across backends: -# -# Flash layout: (2, num_blocks, block_size, num_kv_heads, head_size) -# - dim 0 separates key (index 0) and value (index 1) -# - Used by: FLASH_ATTN, TREE_ATTN, ROCM_AITER_FA, ROCM_ATTN -# -# Block layout: (num_blocks, 2, block_size, num_kv_heads, head_size) -# - dim 1 separates key (index 0) and value (index 1) -# - Used by: TRITON_ATTN -# -# The test creates KV caches in flash layout (the canonical format used by -# tree attention). When a reference backend needs block layout we transpose -# dims 0 and 1. -# -# Note: ROCM_ATTN uses flash layout for storage but its forward path calls -# PagedAttention.split_kv_cache which reinterprets the raw memory as paged -# layout (num_blocks, num_kv_heads, head_size//x, block_size, x). This is -# a view-level incompatibility, not a transpose - see the TODO in -# _get_available_reference_backends for details. -# -# TODO: Replace this mapping with a `KV_CACHE_LAYOUT` class attribute on each -# AttentionImpl so the layout is self-documented by the backend itself, e.g.: -# class TritonAttentionImpl(AttentionImpl): -# KV_CACHE_LAYOUT = "block" -# --------------------------------------------------------------------------- # - -_BLOCK_KV_LAYOUT_BACKENDS = frozenset( - { - AttentionBackendEnum.TRITON_ATTN, - } -) - -# Backends whose do_kv_cache_update requires engine-level state (e.g. -# ForwardContext) that is not available in this test harness, but whose -# KV cache is flash layout and can be written with reshape_and_cache_flash. -# When a backend is listed here, forward_attention() bypasses -# do_kv_cache_update and writes directly to the cache. -_NEEDS_DIRECT_CACHE_UPDATE = frozenset( - { - AttentionBackendEnum.ROCM_AITER_FA, - } -) - -# Backends with known test-harness incompatibilities - see the TODOs -# inside _get_available_reference_backends for details. -_INCOMPATIBLE_REFERENCE_BACKENDS = frozenset( - { - AttentionBackendEnum.ROCM_AITER_FA, - AttentionBackendEnum.ROCM_ATTN, - } -) - - -def _adapt_kv_cache_for_backend( - kv_cache: torch.Tensor, - backend: AttentionBackendEnum, -) -> torch.Tensor: - """Convert kv_cache from flash layout ``(2, num_blocks, ...)`` to block - layout ``(num_blocks, 2, ...)`` if the backend requires it. Returns the - original tensor unchanged when no conversion is needed.""" - if backend in _BLOCK_KV_LAYOUT_BACKENDS: - return kv_cache.transpose(0, 1).contiguous() - return kv_cache - - -def _get_platform_default_backend() -> AttentionBackendEnum: - """Ask the platform what backend it would auto-select at runtime.""" - from vllm.v1.attention.selector import AttentionSelectorConfig - - config = AttentionSelectorConfig( - block_size=32, - kv_cache_dtype="auto", - use_mla=False, - use_sparse=False, - head_size=128, - dtype=torch.bfloat16, - ) - backend_path = current_platform.get_attn_backend_cls( - selected_backend=None, - attn_selector_config=config, - ) - for backend in AttentionBackendEnum: - try: - if backend.get_path() == backend_path: - return backend - except ValueError: - continue - raise RuntimeError( - f"Platform returned backend path '{backend_path}' " - f"that doesn't match any AttentionBackendEnum member." - ) - - -def _get_available_reference_backends() -> list[AttentionBackendEnum]: - """Collect all reference backends the current platform can run. - - On CUDA this is just FLASH_ATTN. On ROCm this includes the platform - default plus every backend the hardware supports, so the test validates - tree attention against all of them. - """ - if current_platform.is_rocm(): - backends: list[AttentionBackendEnum] = [] - - # 1. Whatever the platform would auto-select at runtime. - default_backend = _get_platform_default_backend() - if default_backend not in _INCOMPATIBLE_REFERENCE_BACKENDS: - backends.append(default_backend) - - # 2. TRITON_ATTN - always available on ROCm. - if AttentionBackendEnum.TRITON_ATTN not in backends: - backends.append(AttentionBackendEnum.TRITON_ATTN) - - # TODO: Enable ROCM_ATTN. Its forward path uses - # PagedAttention.split_kv_cache which reinterprets the raw - # cache memory as paged layout: - # key: (num_blocks, num_kv_heads, head_size//x, block_size, x) - # value: (num_blocks, num_kv_heads, head_size, block_size) - # Tree attention writes prefix data in NHD flash layout, so the - # same bytes produce completely different values when read in - # paged format. Supporting ROCM_ATTN would require writing - # prefix data via PagedAttention.write_to_paged_cache into a - # separate paged-format KV cache. - - # TODO: Enable ROCM_AITER_FA. Its metadata builder reads head - # counts from the model config at construction time and - # allocates extend_workspace with those dimensions. The test - # uses independent head count parameters (num_heads=2/4, - # num_kv_heads=2) that don't match the model config - # (Llama-3-8B: 32 q heads, 8 kv heads), causing a head count - # mismatch in flash_attn_varlen_func during extend_forward. - # Fixing this requires either matching test head counts to the - # model config or decoupling the builder from model config - # head geometry. The direct cache update path - # (_NEEDS_DIRECT_CACHE_UPDATE) is already in place for when - # this is resolved. - - return backends - - # CUDA: flash attention. - return [AttentionBackendEnum.FLASH_ATTN] - - -class MockAttentionLayer(torch.nn.Module): - _q_scale = torch.tensor(1.0, dtype=torch.float32, device=DEVICE_TYPE) - _k_scale = torch.tensor(1.0, dtype=torch.float32, device=DEVICE_TYPE) - _v_scale = torch.tensor(1.0, dtype=torch.float32, device=DEVICE_TYPE) - layer_name = "mock_layer" - - def __init__(self): - super().__init__() - - def forward(self, x): - return x - - -def forward_attention( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - kv_cache: torch.Tensor, - block_table: torch.Tensor, - slot_mapping: torch.Tensor, - seqlen_k: int, - backend: AttentionBackendEnum, - spec_token_tree: str | None = None, - num_spec_tokens: int = 0, -) -> torch.Tensor: - """Run a single attention forward pass through the given backend. - - ``kv_cache`` is expected in **flash layout** - ``(2, num_blocks, block_size, num_kv_heads, head_size)``. - It is automatically converted when the target backend needs a - different layout. - """ - batch_size, q_len, num_heads, dim_per_head = q.shape - num_kv_heads = k.shape[-2] - # Initialize the query and KV sequence lengths. - query_start_loc = q_len * torch.arange( - batch_size + 1, device=q.device, dtype=torch.int32 - ) - query_lens = torch.diff(query_start_loc) - seq_lens = torch.full( - (batch_size,), - seqlen_k, - device=q.device, - dtype=torch.int32, - ) - context_lens = seq_lens - query_lens - max_seq_len = int(seq_lens.max()) - max_query_len = q_len - num_actual_tokens = query_start_loc[-1] - - softmax_scale = q.shape[-1] ** (-0.5) - layer = MockAttentionLayer() - - # Build common metadata. - model_name = "meta-llama/Meta-Llama-3-8B" - builder_cls, impl_cls = try_get_attention_backend(backend) - vllm_config = create_vllm_config(model_name=model_name, max_model_len=max(seq_lens)) - if spec_token_tree is not None: - # Create speculative config if token tree is specified. - vllm_config.speculative_config = SpeculativeConfig( - target_model_config=vllm_config.model_config, - target_parallel_config=ParallelConfig(), - model=model_name, - method="eagle", - num_speculative_tokens=num_spec_tokens, - speculative_token_tree=spec_token_tree, - ) - kv_cache_spec = create_standard_kv_cache_spec(vllm_config) - builder = builder_cls(kv_cache_spec, [], vllm_config, q.device) - seq_lens_cpu = seq_lens.cpu() - common_attn_metadata = CommonAttentionMetadata( - query_start_loc=query_start_loc, - query_start_loc_cpu=query_start_loc.cpu(), - seq_lens=seq_lens, - seq_lens_cpu_upper_bound=seq_lens_cpu, - _seq_lens_cpu=seq_lens_cpu, - _num_computed_tokens_cpu=context_lens.cpu(), - num_reqs=batch_size, - num_actual_tokens=num_actual_tokens, - max_query_len=max_query_len, - max_seq_len=max_seq_len, - block_table_tensor=block_table, - slot_mapping=slot_mapping, - ) - - # Build attention metadata. - attn_metadata = builder.build( - common_prefix_len=0, - common_attn_metadata=common_attn_metadata, - ) - - # Initialize the backend implementation. - instance = impl_cls( - num_heads=num_heads, - head_size=dim_per_head, - scale=softmax_scale, - num_kv_heads=num_kv_heads, - alibi_slopes=None, - sliding_window=None, - kv_cache_dtype="auto", - ) - - # Adapt KV cache layout for this backend. - adapted_kv_cache = _adapt_kv_cache_for_backend(kv_cache, backend) - - # Run forward pass and return output. - query = q.view(-1, num_heads, dim_per_head) - key = k.view(-1, num_kv_heads, dim_per_head) - value = v.view(-1, num_kv_heads, dim_per_head) - output = torch.empty_like(query) - if not try_backend_includes_kv_cache_update(backend): - if backend in _NEEDS_DIRECT_CACHE_UPDATE: - # This backend's do_kv_cache_update requires engine-level - # ForwardContext that isn't available in this test harness. - # Write directly using reshape_and_cache_flash since the - # KV cache layout is identical (flash layout, unbind on dim 0). - key_cache, value_cache = adapted_kv_cache.unbind(0) - torch.ops._C_cache_ops.reshape_and_cache_flash( - key, - value, - key_cache, - value_cache, - attn_metadata.slot_mapping, - "auto", - layer._k_scale, - layer._v_scale, - ) - else: - instance.do_kv_cache_update( - layer=layer, - key=key, - value=value, - kv_cache=adapted_kv_cache, - slot_mapping=attn_metadata.slot_mapping, - ) - return instance.forward( - layer=layer, - query=query, - key=key, - value=value, - kv_cache=adapted_kv_cache.clone(), - attn_metadata=attn_metadata, - output=output, - ) - - -@pytest.mark.parametrize( - "reference_backend", - _get_available_reference_backends(), - ids=lambda b: b.name, -) -def test_tree_attn_correctness( - reference_backend: AttentionBackendEnum, -) -> None: - set_random_seed(42) - - device = "cuda" - tree_attn_masks = { - # Chain. - "[(0,), (0, 0), (0, 0, 0)]": torch.tensor( - [ - [1, 0, 0, 0], - [1, 1, 0, 0], - [1, 1, 1, 0], - [1, 1, 1, 1], - ], - device=device, - dtype=torch.int32, - ), - # Tree. - "[(0,), (1,), (0, 0), (0, 1), (1, 0), (1, 1)]": torch.tensor( - [ - [1, 0, 0, 0, 0, 0, 0], - [1, 1, 0, 0, 0, 0, 0], - [1, 0, 1, 0, 0, 0, 0], - [1, 1, 0, 1, 0, 0, 0], - [1, 1, 0, 0, 1, 0, 0], - [1, 0, 1, 0, 0, 1, 0], - [1, 0, 1, 0, 0, 0, 1], - ], - device=device, - dtype=torch.int32, - ), - } - - dim_per_head = 128 - num_kv_heads = 2 - block_size = 32 - max_sequence_length = 8192 - randomize_blocks = True - for batch_size in [1, 16, 32]: - for num_heads in [2, 4]: - for sequence_position in [16, 1024, 2048]: - for spec_token_tree, tree_attn_mask in tree_attn_masks.items(): - # Assert that the number of heads is divisible - # by the number of KV heads. - assert num_heads % num_kv_heads == 0 - - # Initialize q, k, and v. - tree_size_q = tree_attn_mask.shape[0] - seqlen_k = sequence_position + tree_size_q - q = torch.randn( - (batch_size, tree_size_q, num_heads, dim_per_head), - device=device, - dtype=torch.bfloat16, - ) - k = torch.randn( - (batch_size, tree_size_q, num_kv_heads, dim_per_head), - device=device, - dtype=torch.bfloat16, - ) - v = torch.randn( - (batch_size, tree_size_q, num_kv_heads, dim_per_head), - device=device, - dtype=torch.bfloat16, - ) - - # KV cache in flash layout - the canonical format for - # tree attention. forward_attention() handles conversion - # when needed. - assert max_sequence_length % block_size == 0 - max_blocks_per_batch = max_sequence_length // block_size - kv_cache = torch.randn( - ( - 2, - batch_size * max_blocks_per_batch, - block_size, - num_kv_heads, - dim_per_head, - ), - device=q.device, - dtype=torch.bfloat16, - ) - num_alloc_blocks_per_batch = math.ceil(seqlen_k / block_size) - block_table = torch.zeros( - (batch_size, max_blocks_per_batch), - device=q.device, - dtype=torch.int32, - ) - block_ids = torch.arange( - 0, - batch_size * num_alloc_blocks_per_batch, - device=q.device, - dtype=torch.int32, - ) - if randomize_blocks: - # Randomize the block ids. - block_ids = block_ids[torch.randperm(block_ids.numel())] - block_table[:, :num_alloc_blocks_per_batch] = block_ids.view( - -1, num_alloc_blocks_per_batch - ) - - # Set up the slot mapping for the input KVs. - tree_positions = sequence_position + torch.arange( - 0, - tree_size_q, - device=q.device, - dtype=torch.int64, - ).repeat(batch_size, 1) - tree_slot_mapping = _gen_slot_mapping( - tree_positions, block_table, block_size - ) - - # Compute attention for the tree. - tree_attn_output = forward_attention( - q=q, - k=k, - v=v, - kv_cache=kv_cache, - block_table=block_table, - slot_mapping=tree_slot_mapping, - seqlen_k=seqlen_k, - backend=AttentionBackendEnum.TREE_ATTN, - spec_token_tree=spec_token_tree, - num_spec_tokens=tree_size_q - 1, - ).view(batch_size, -1, num_heads, dim_per_head) - - # Verify each branch against the reference backend. - for q_index in range(tree_size_q): - # Get the q, k, and v for the branch. - branch_mask = tree_attn_mask[q_index, :] - branch_indices = torch.nonzero(branch_mask, as_tuple=True)[0] - q_len = branch_indices.shape[0] - q_branch = q[:, branch_indices] - k_branch = k[:, branch_indices] - v_branch = v[:, branch_indices] - - # Setup slot mapping for the branch. - branch_positions = sequence_position + torch.arange( - 0, - q_len, - device=q.device, - dtype=torch.int64, - ).repeat(batch_size, 1) - branch_slot_mapping = _gen_slot_mapping( - branch_positions, block_table, block_size - ) - - # Reference attention for this branch. - ref_output = forward_attention( - q=q_branch, - k=k_branch, - v=v_branch, - kv_cache=kv_cache, - block_table=block_table, - slot_mapping=branch_slot_mapping, - seqlen_k=sequence_position + q_len, - backend=reference_backend, - ).view(batch_size, -1, num_heads, dim_per_head) - - # Compare the outputs. - assert torch.allclose( - tree_attn_output[:, branch_indices], - ref_output, - atol=7.81e-3, - ), ( - f"outputs are not close for " - f"reference_backend: {reference_backend.name}, " - f"batch_size: {batch_size}, " - f"num_heads: {num_heads}, " - f"sequence_position: {sequence_position}, " - f"tree_attn_mask: {tree_attn_mask}, " - f"q_index: {q_index}." - ) - - -def _gen_slot_mapping( - positions: torch.Tensor, block_table: torch.Tensor, block_size: int -): - block_indices = positions // block_size - blocks = block_table.gather(dim=1, index=block_indices) - return (blocks * block_size + positions % block_size).view(-1) diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index 0f8adc6e4ac..4589820c9fd 100644 --- a/vllm/config/speculative.py +++ b/vllm/config/speculative.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import ast import copy from typing import TYPE_CHECKING, Any, Literal, get_args @@ -145,9 +144,6 @@ class SpeculativeConfig: provided. Defaults to 1.""" # Alternative drafting strategies - speculative_token_tree: str | None = None - """Specifies the tree structure for speculative token generation. - """ parallel_drafting: bool = False """Enable parallel drafting, where all speculative tokens are generated in parallel rather than sequentially. This can improve performance but @@ -739,23 +735,10 @@ class SpeculativeConfig: f" must be divisible by {n_predict=}" ) - if self.speculative_token_tree is None: - if self.num_speculative_tokens is None: - raise ValueError( - "A speculative model was provided, but neither " - "`speculative_token_tree` nor `num_speculative_tokens` " - "was provided" - ) - - # Generate chain of tokens. - self.speculative_token_tree = str( - [(i + 1) * (0,) for i in range(self.num_speculative_tokens)] - ) - else: - # Sort the token tree breadth-first. - tree_choices = ast.literal_eval(self.speculative_token_tree) - self.speculative_token_tree = str( - sorted(tree_choices, key=lambda t: (len(t), t)) + if self.num_speculative_tokens is None: + raise ValueError( + "A speculative model was provided, but " + "`num_speculative_tokens` was not provided" ) self.draft_tensor_parallel_size = ( diff --git a/vllm/v1/attention/backends/registry.py b/vllm/v1/attention/backends/registry.py index f31edfafc38..552b7f3373c 100644 --- a/vllm/v1/attention/backends/registry.py +++ b/vllm/v1/attention/backends/registry.py @@ -76,7 +76,6 @@ class AttentionBackendEnum(Enum, metaclass=_AttentionBackendEnumMeta): FLASH_ATTN_MLA = "vllm.v1.attention.backends.mla.flashattn_mla.FlashAttnMLABackend" NO_ATTENTION = "vllm.v1.attention.backends.no_attention.NoAttentionBackend" FLEX_ATTENTION = "vllm.v1.attention.backends.flex_attention.FlexAttentionBackend" - TREE_ATTN = "vllm.v1.attention.backends.tree_attn.TreeAttentionBackend" ROCM_AITER_UNIFIED_ATTN = ( "vllm.v1.attention.backends.rocm_aiter_unified_attn." "RocmAiterUnifiedAttentionBackend" diff --git a/vllm/v1/attention/backends/tree_attn.py b/vllm/v1/attention/backends/tree_attn.py deleted file mode 100644 index af9c91d11ee..00000000000 --- a/vllm/v1/attention/backends/tree_attn.py +++ /dev/null @@ -1,488 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Attention layer with TreeAttention.""" - -import ast -from dataclasses import dataclass -from typing import ClassVar - -import torch - -from vllm import _custom_ops as ops -from vllm.config import VllmConfig -from vllm.config.cache import CacheDType -from vllm.logger import init_logger -from vllm.v1.attention.backend import ( - AttentionBackend, - AttentionImpl, - AttentionMetadataBuilder, - AttentionType, - CommonAttentionMetadata, - MultipleOf, -) -from vllm.v1.attention.backends.utils import ( - split_decodes_and_prefills, -) -from vllm.v1.attention.ops.triton_unified_attention import unified_attention -from vllm.v1.kv_cache_interface import AttentionSpec - -logger = init_logger(__name__) - - -class TreeAttentionBackend(AttentionBackend): - supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16] - supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [ - "auto", - "float16", - "bfloat16", - ] - forward_includes_kv_cache_update: bool = False - - @staticmethod - def get_supported_kernel_block_sizes() -> list[int | MultipleOf]: - return [MultipleOf(16)] - - @classmethod - def get_supported_head_sizes(cls) -> list[int]: - return [32, 64, 96, 128, 160, 192, 224, 256] - - @staticmethod - def get_name() -> str: - return "TREE_ATTN" - - @staticmethod - def get_impl_cls() -> type["TreeAttentionImpl"]: - return TreeAttentionImpl - - @staticmethod - def get_kv_cache_shape( - num_blocks: int, - block_size: int, - num_kv_heads: int, - head_size: int, - cache_dtype_str: str = "auto", - ) -> tuple[int, ...]: - if block_size % 16 != 0: - raise ValueError("Block size must be a multiple of 16.") - return (2, num_blocks, block_size, num_kv_heads, head_size) - - @staticmethod - def get_builder_cls() -> type["TreeAttentionMetadataBuilder"]: - return TreeAttentionMetadataBuilder - - @staticmethod - def use_cascade_attention(*args, **kwargs) -> bool: - return False - - -@dataclass -class TreeAttentionMetadata: - num_actual_tokens: int # Number of tokens excluding padding. - max_query_len: int - query_start_loc: torch.Tensor - max_seq_len: int - seq_lens: torch.Tensor - block_table: torch.Tensor - slot_mapping: torch.Tensor - - num_prefill_tokens: int = 0 - num_decode_tokens: int = 0 - num_prefills: int = 0 - num_decodes: int = 0 - - # Precomputed (on CPU in the builder) max_query_len and max_seq_len for - # the prefill-only and decode-only sub-batches. Used by the properties - # below to avoid a GPU->CPU sync via `.max().item()` on every forward. - max_query_len_prefill: int = 0 - max_seq_len_prefill: int = 0 - max_query_len_decode: int = 0 - max_seq_len_decode: int = 0 - - tree_attn_bias: torch.Tensor | None = None - - # Cached Prefill/decode metadata. - _cached_prefill_metadata: "TreeAttentionMetadata | None" = None - _cached_decode_metadata: "TreeAttentionMetadata | None" = None - - @property - def prefill_metadata(self) -> "TreeAttentionMetadata | None": - if self.num_prefills == 0: - return None - - if self._cached_prefill_metadata is not None: - # Recover cached prefill-phase attention - # metadata structure - return self._cached_prefill_metadata - - q_start_loc = self.query_start_loc[self.num_decodes :] - kv_seqlens = self.seq_lens[self.num_decodes :] - # Construct & cache prefill-phase attention metadata structure - self._cached_prefill_metadata = TreeAttentionMetadata( - num_actual_tokens=self.num_prefill_tokens, - max_query_len=self.max_query_len_prefill, - query_start_loc=q_start_loc - q_start_loc[0], - max_seq_len=self.max_seq_len_prefill, - seq_lens=kv_seqlens, - block_table=self.block_table[self.num_decodes :], - slot_mapping=self.slot_mapping[self.num_decode_tokens :], - ) - return self._cached_prefill_metadata - - @property - def decode_metadata(self) -> "TreeAttentionMetadata | None": - if self.num_decode_tokens == 0: - return None - - if self._cached_decode_metadata is not None: - # Recover cached decode-phase attention - # metadata structure - return self._cached_decode_metadata - - q_start_loc = self.query_start_loc[: self.num_decodes + 1] - kv_seqlens = self.seq_lens[: self.num_decodes] - # Construct & cache decode-phase attention metadata structure - self._cached_decode_metadata = TreeAttentionMetadata( - num_actual_tokens=self.num_decode_tokens, - max_query_len=self.max_query_len_decode, - query_start_loc=q_start_loc, - max_seq_len=self.max_seq_len_decode, - seq_lens=kv_seqlens, - block_table=self.block_table[: self.num_decodes], - slot_mapping=self.slot_mapping[: self.num_decode_tokens], - tree_attn_bias=self.tree_attn_bias, - ) - return self._cached_decode_metadata - - -class TreeAttentionMetadataBuilder(AttentionMetadataBuilder[TreeAttentionMetadata]): - def __init__( - self, - kv_cache_spec: AttentionSpec, - layer_names: list[str], - vllm_config: VllmConfig, - device: torch.device, - ): - super().__init__(kv_cache_spec, layer_names, vllm_config, device) - - self.block_size = kv_cache_spec.block_size - - spec_config = vllm_config.speculative_config - spec_token_tree: str | None = None - if spec := spec_config: - spec_token_tree = spec.speculative_token_tree - tree_choices: list[tuple[int, ...]] = ( - ast.literal_eval(spec_token_tree) if spec_token_tree is not None else [(0,)] - ) - # Construct the tree attention bias. - depth_counts = _get_depth_counts(tree_choices) - self.tree_attn_bias = _prepare_tree_attn_bias( - tree_choices, - depth_counts, - dtype=torch.float32, - device=device, - ) - - self.reorder_batch_threshold = self.tree_attn_bias.shape[0] - - def build( - self, - common_prefix_len: int, - common_attn_metadata: CommonAttentionMetadata, - fast_build: bool = False, - ) -> TreeAttentionMetadata: - decode_threshold = self.tree_attn_bias.shape[0] - num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = ( - split_decodes_and_prefills( - common_attn_metadata, decode_threshold=decode_threshold - ) - ) - - num_actual_tokens = common_attn_metadata.num_actual_tokens - q_start_loc = common_attn_metadata.query_start_loc - max_query_len = common_attn_metadata.max_query_len - kv_seqlens = common_attn_metadata.seq_lens - max_seq_len = common_attn_metadata.max_seq_len - block_table = common_attn_metadata.block_table_tensor - slot_mapping = common_attn_metadata.slot_mapping - - # Precompute prefill/decode sub-batch max_query_len / max_seq_len on - # CPU so the prefill_metadata / decode_metadata properties don't need - # a GPU->CPU sync via `.max().item()` on every forward. - # Prefer `seq_lens_cpu_upper_bound` over the (deprecated) - # `seq_lens_cpu` property: the upper bound is precise for prefill - # rows and optimistic-but-safe for decode rows (workspace sizing - # from `max()` is fine with an over-estimate), and avoids the - # `seq_lens.to("cpu")` sync the property would fall through to in - # async-spec-decode mode. The draft-attention path (eagle - # speculator) doesn't populate it; fall back to the batch-wide - # `max_seq_len` as a safe upper bound for both sub-batches. - q_start_loc_cpu = common_attn_metadata.query_start_loc_cpu - seq_lens_cpu = common_attn_metadata.seq_lens_cpu_upper_bound - if num_prefills > 0: - q_seqlens_p = torch.diff(q_start_loc_cpu[num_decodes:]) - max_query_len_prefill = int(q_seqlens_p.max()) - max_seq_len_prefill = ( - int(seq_lens_cpu[num_decodes:].max()) - if seq_lens_cpu is not None - else max_seq_len - ) - else: - max_query_len_prefill = 0 - max_seq_len_prefill = 0 - if num_decodes > 0: - q_seqlens_d = torch.diff(q_start_loc_cpu[: num_decodes + 1]) - max_query_len_decode = int(q_seqlens_d.max()) - max_seq_len_decode = ( - int(seq_lens_cpu[:num_decodes].max()) - if seq_lens_cpu is not None - else max_seq_len - ) - else: - max_query_len_decode = 0 - max_seq_len_decode = 0 - - return TreeAttentionMetadata( - num_actual_tokens=num_actual_tokens, - num_prefill_tokens=num_prefill_tokens, - num_decode_tokens=num_decode_tokens, - num_prefills=num_prefills, - num_decodes=num_decodes, - max_query_len=max_query_len, - query_start_loc=q_start_loc, - max_seq_len=max_seq_len, - seq_lens=kv_seqlens, - block_table=block_table, - slot_mapping=slot_mapping, - max_query_len_prefill=max_query_len_prefill, - max_seq_len_prefill=max_seq_len_prefill, - max_query_len_decode=max_query_len_decode, - max_seq_len_decode=max_seq_len_decode, - tree_attn_bias=self.tree_attn_bias, - ) - - def build_for_drafting( - self, - common_attn_metadata: CommonAttentionMetadata, - draft_index: int, - ) -> TreeAttentionMetadata: - # Cache the original tree attention bias. - orig_tree_attn_bias = self.tree_attn_bias - - if draft_index == 0: - # Use prefill for drafting at the root level. - self.tree_attn_bias = torch.empty(0) - else: - # Slice the tree attention bias for drafting. Exclude - # the root level. - start, end = 1, 1 + common_attn_metadata.max_query_len - self.tree_attn_bias = self.tree_attn_bias[start:end, start:end].contiguous() - - # Build attention bias. - attn_metadata = self.build(0, common_attn_metadata, fast_build=True) - - # Reset the tree attention bias to the original value. - self.tree_attn_bias = orig_tree_attn_bias - return attn_metadata - - -def _get_depth_counts(sorted_tree_choices: list[tuple[int, ...]]) -> list[int]: - # Count the number of choices at each depth of the tree. - depth_counts = [] - prev_depth = 0 - for path in sorted_tree_choices: - depth = len(path) - if depth != prev_depth: - depth_counts.append(0) - depth_counts[depth - 1] += 1 - prev_depth = depth - return depth_counts - - -def _prepare_tree_attn_bias( - sorted_tree_choices: list[tuple[int, ...]], - depth_counts: list[int], - dtype: torch.dtype | None, - device: torch.device | None, -) -> torch.Tensor: - # +1 comes from the additional root node. - tree_len = len(sorted_tree_choices) + 1 - tree_attn_mask = torch.full( - (tree_len, tree_len), -torch.inf, device=device, dtype=dtype - ) - - # Set diagonal to all zeros. Each token should - # attend to itself. - mask_val = 0 - for i in range(tree_len): - tree_attn_mask[i, i] = mask_val - - # Set root to all zeros. All tokens attend to it. - tree_attn_mask[:, 0] = mask_val - - # Set all ancestors to zeros. - start = 0 - for i in range(len(depth_counts)): - for j in range(depth_counts[i]): - cur_tree_choice = sorted_tree_choices[start + j] - # Retrieve ancestor position. - if len(cur_tree_choice) == 1: - continue - ancestor_idx = [] - for c in range(len(cur_tree_choice) - 1): - ancestor_idx.append( - sorted_tree_choices.index(cur_tree_choice[: c + 1]) + 1 - ) - tree_attn_mask[j + start + 1, ancestor_idx] = mask_val - start += depth_counts[i] - return tree_attn_mask - - -class TreeAttentionImpl(AttentionImpl): - def __init__( - self, - num_heads: int, - head_size: int, - scale: float, - num_kv_heads: int, - alibi_slopes: list[float] | None, - sliding_window: int | None, - kv_cache_dtype: str, - logits_soft_cap: float | None = None, - attn_type: AttentionType = AttentionType.DECODER, - kv_sharing_target_layer_name: str | None = None, - ) -> None: - self.num_heads = num_heads - self.head_size = head_size - self.scale = float(scale) - self.num_kv_heads = num_kv_heads - self.num_queries_per_kv = self.num_heads // self.num_kv_heads - self.kv_cache_dtype = kv_cache_dtype - self.kv_sharing_target_layer_name = kv_sharing_target_layer_name - if alibi_slopes is not None: - alibi_slopes = torch.tensor(alibi_slopes, dtype=torch.float32) - self.alibi_slopes = alibi_slopes - if logits_soft_cap is None: - # Setting logits_soft_cap to 0 means no soft cap. - logits_soft_cap = 0 - self.logits_soft_cap = logits_soft_cap - if sliding_window is None: - self.sliding_window = (-1, -1) - else: - self.sliding_window = (sliding_window - 1, 0) - - if attn_type != AttentionType.DECODER: - raise NotImplementedError( - "Encoder self-attention and " - "encoder/decoder cross-attention " - "are not implemented for " - "TreeAttentionImpl." - ) - - def do_kv_cache_update( - self, - layer: torch.nn.Module, - key: torch.Tensor, - value: torch.Tensor, - kv_cache: torch.Tensor, - slot_mapping: torch.Tensor, - ) -> None: - key_cache, value_cache = kv_cache.unbind(0) - - # Reshape the input keys and values and store them in the cache. - # NOTE(woosuk): Here, key and value are padded while slot_mapping is - # not padded. However, we don't need to do key[:num_actual_tokens] - # and value[:num_actual_tokens] because the reshape_and_cache_flash - # op uses the slot_mapping's shape to determine the number of - # actual tokens. - ops.reshape_and_cache_flash( - key, - value, - key_cache, - value_cache, - slot_mapping, - self.kv_cache_dtype, - layer._k_scale, - layer._v_scale, - ) - - def forward( - self, - layer: torch.nn.Module, - query: torch.Tensor, - key: torch.Tensor, - value: torch.Tensor, - kv_cache: torch.Tensor, - attn_metadata: TreeAttentionMetadata, - output: torch.Tensor, - output_scale: torch.Tensor | None = None, - output_block_scale: torch.Tensor | None = None, - ) -> torch.Tensor: - """Forward pass with TreeAttention. - - Args: - query: shape = [num_tokens, num_heads, head_size] - key: shape = [num_tokens, num_kv_heads, head_size] - value: shape = [num_tokens, num_kv_heads, head_size] - kv_cache: shape = - [2, num_blocks, block_size, num_kv_heads, head_size] - attn_metadata: Metadata for attention. - Returns: - shape = [num_tokens, num_heads * head_size] - """ - if output_scale is not None or output_block_scale is not None: - raise NotImplementedError( - "fused output quantization is not yet supported for TreeAttentionImpl" - ) - - if attn_metadata is None: - # Profiling run. - return output.fill_(0) - - key_cache, value_cache = kv_cache.unbind(0) - - num_actual_tokens = attn_metadata.num_actual_tokens - num_decode_tokens = attn_metadata.num_decode_tokens - descale_shape = (attn_metadata.query_start_loc.shape[0] - 1, key.shape[1]) - if prefill_meta := attn_metadata.prefill_metadata: - unified_attention( - q=query[num_decode_tokens:num_actual_tokens], - k=key_cache, - v=value_cache, - out=output[num_decode_tokens:num_actual_tokens], - cu_seqlens_q=prefill_meta.query_start_loc, - max_seqlen_q=prefill_meta.max_query_len, - seqused_k=prefill_meta.seq_lens, - max_seqlen_k=prefill_meta.max_seq_len, - softmax_scale=self.scale, - causal=True, - alibi_slopes=self.alibi_slopes, - window_size=self.sliding_window, - block_table=prefill_meta.block_table, - softcap=self.logits_soft_cap, - q_descale=None, # Not supported - k_descale=layer._k_scale.expand(descale_shape), - v_descale=layer._v_scale.expand(descale_shape), - ) - - if decode_meta := attn_metadata.decode_metadata: - unified_attention( - q=query[:num_decode_tokens], - k=key_cache, - v=value_cache, - out=output[:num_decode_tokens], - cu_seqlens_q=decode_meta.query_start_loc, - max_seqlen_q=decode_meta.max_query_len, - seqused_k=decode_meta.seq_lens, - max_seqlen_k=decode_meta.max_seq_len, - softmax_scale=self.scale, - causal=True, - alibi_slopes=self.alibi_slopes, - qq_bias=decode_meta.tree_attn_bias, - window_size=self.sliding_window, - block_table=decode_meta.block_table, - softcap=self.logits_soft_cap, - q_descale=None, # Not supported - k_descale=layer._k_scale.expand(descale_shape), - v_descale=layer._v_scale.expand(descale_shape), - ) - return output diff --git a/vllm/v1/spec_decode/llm_base_proposer.py b/vllm/v1/spec_decode/llm_base_proposer.py index 8ee349a1cc0..08ed9798ed6 100644 --- a/vllm/v1/spec_decode/llm_base_proposer.py +++ b/vllm/v1/spec_decode/llm_base_proposer.py @@ -1,6 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import ast from importlib.util import find_spec from typing import Any, cast @@ -29,10 +28,6 @@ from vllm.platforms import current_platform from vllm.utils.platform_utils import is_pin_memory_available from vllm.v1.attention.backend import CommonAttentionMetadata from vllm.v1.attention.backends.registry import AttentionBackendEnum -from vllm.v1.attention.backends.tree_attn import ( - TreeAttentionMetadata, - TreeAttentionMetadataBuilder, -) from vllm.v1.attention.backends.triton_attn import TritonAttentionMetadata from vllm.v1.cudagraph_dispatcher import CudagraphDispatcher from vllm.v1.kv_cache_interface import KVCacheConfig, UniformTypeKVCacheSpecs @@ -284,29 +279,6 @@ class SpecDecodeBaseProposer: self.allowed_attn_types = tuple(rocm_types) - # Parse the speculative token tree. - spec_token_tree = self.speculative_config.speculative_token_tree - assert spec_token_tree is not None - self.tree_choices: list[tuple[int, ...]] = ast.literal_eval(spec_token_tree) - tree_depth = len(self.tree_choices[-1]) - # Precompute per-level properties of the tree. - num_drafts_per_level = [0] * tree_depth - for node in self.tree_choices: - num_drafts_per_level[len(node) - 1] += 1 - self.cu_drafts_per_level = [num_drafts_per_level[0]] - self.child_drafts_per_level = [num_drafts_per_level[0]] - for level in range(1, tree_depth): - self.cu_drafts_per_level.append( - self.cu_drafts_per_level[-1] + num_drafts_per_level[level] - ) - self.child_drafts_per_level.append( - num_drafts_per_level[level] // num_drafts_per_level[level - 1] - ) - # Precompute draft position offsets in flattened tree. - self.tree_draft_pos_offsets = torch.arange( - 1, len(self.tree_choices) + 1, device=device, dtype=torch.int32 - ).repeat(self.max_batch_size, 1) - def _raise_if_padded_drafter_batch_disabled(self): if self.speculative_config.disable_padded_drafter_batch: raise NotImplementedError( @@ -512,20 +484,6 @@ class SpecDecodeBaseProposer: # (which read via _get_positions) use the correct values. self.positions[:batch_size] = positions - if any(isinstance(md, TreeAttentionMetadata) for md in per_group_attn_metadata): - # Draft using tree attention - requires full logits for top-k - logits = self.model.compute_logits(sample_hidden_states) - draft_token_ids_list = self.propose_tree( - batch_size=batch_size, - logits=logits, - positions=positions, - hidden_states=hidden_states, - common_attn_metadata=common_attn_metadata, - slot_mappings=slot_mappings, - ) - # [batch_size, num_tree_tokens] - return torch.cat(draft_token_ids_list, dim=1) - draft_token_ids = self._greedy_sample(sample_hidden_states) if self.allowed_attn_types is not None: @@ -1014,178 +972,6 @@ class SpecDecodeBaseProposer: num_rejected_tokens_gpu, ) - def propose_tree( - self, - batch_size: int, - # [num_tokens, vocab_size] - logits: torch.Tensor, - # [num_tokens] - positions: torch.Tensor, - # [num_tokens, hidden_size] - hidden_states: torch.Tensor, - common_attn_metadata: CommonAttentionMetadata, - slot_mappings: dict[str, torch.Tensor] - | list[dict[str, torch.Tensor]] - | None = None, - ) -> list[torch.Tensor]: - tree_attn_metadata_builder = self.draft_attn_groups[0].get_metadata_builder() - assert isinstance(tree_attn_metadata_builder, TreeAttentionMetadataBuilder) - - total_num_drafts = self.cu_drafts_per_level[0] - level_num_drafts = total_num_drafts - # Sample a draft token for each child at the tree root level. - num_children = self.child_drafts_per_level[0] - if num_children == 1: - draft_token_ids = logits.argmax(dim=-1).view(batch_size, -1) - else: - draft_token_ids = torch.topk(logits, num_children, dim=-1).indices.view( - batch_size, -1 - ) - draft_token_ids_list = [draft_token_ids] - draft_hidden_states = hidden_states.view(batch_size, 1, -1) - - # Initialize empty tensors for concatenation with the level outputs. - tree_input_ids = torch.empty( - 0, device=self.input_ids.device, dtype=self.input_ids.dtype - ) - tree_positions = torch.empty( - 0, device=self.positions.device, dtype=self.positions.dtype - ) - tree_hidden_states = torch.empty( - 0, device=self.hidden_states.device, dtype=self.hidden_states.dtype - ) - # Precompute the draft token positions. - flattened_draft_positions = ( - positions.view(batch_size, -1) + self.tree_draft_pos_offsets[:batch_size, :] - ) - tree_depth = len(self.cu_drafts_per_level) - for level in range(tree_depth - 1): - # Get draft positions for RoPE. - draft_positions = positions + (level + 1) - exceeds_max_model_len = (positions + total_num_drafts) >= self.max_model_len - # Mask out the position ids that exceed the max model length. - # Otherwise, we may get out-of-range error in RoPE. - draft_positions = torch.where( - exceeds_max_model_len, - 0, - draft_positions, - ).view(batch_size, -1) - - if level_num_drafts > 1: - # Repeat the positions for each draft at this level. - draft_positions = draft_positions.repeat_interleave( - level_num_drafts, dim=1 - ) - - if num_children > 1: - # Repeat draft hidden states for each child. - draft_hidden_states = draft_hidden_states.repeat_interleave( - num_children, dim=1 - ) - - # Concatenate the draft tokens, positions, and hidden states. - tree_input_ids = torch.cat([tree_input_ids, draft_token_ids], dim=1) - tree_positions = torch.cat([tree_positions, draft_positions], dim=1) - tree_hidden_states = torch.cat( - [tree_hidden_states, draft_hidden_states], dim=1 - ) - - # Build new attention metadata for the next level of drafts. - # This is necessary to support tree attention. - query_len = total_num_drafts - common_attn_metadata = replace( - common_attn_metadata, - query_start_loc=query_len * self.arange[: batch_size + 1], - seq_lens=common_attn_metadata.seq_lens + level_num_drafts, - num_actual_tokens=batch_size * query_len, - max_query_len=query_len, - ) - attn_metadata = tree_attn_metadata_builder.build_for_drafting( - common_attn_metadata=common_attn_metadata, draft_index=level + 1 - ) - - # Apply new attention metadata to all draft layers. - per_layer_attn_metadata = {} - for attn_group in self.draft_attn_groups: - for layer_name in attn_group.layer_names: - per_layer_attn_metadata[layer_name] = attn_metadata - - # Consider max model length. - attn_metadata.max_seq_len = min( - attn_metadata.max_seq_len, self.max_model_len - ) - # For the requests that exceed the max model length, we set the - # sequence length to 1 to minimize their overheads in attention. - attn_metadata.seq_lens.masked_fill_(exceeds_max_model_len, 1) - - # Compute the slot mapping. - block_size = tree_attn_metadata_builder.kv_cache_spec.block_size - query_positions = flattened_draft_positions[:, level : level + query_len] - block_numbers = query_positions // block_size - block_ids = attn_metadata.block_table.gather(dim=1, index=block_numbers) - slot_mapping = block_ids * block_size + query_positions % block_size - # Mask out the slot mappings that exceed the max model length. - # Otherwise, the KV cache will be inadvertently updated with the - # padding tokens. - slot_mapping[exceeds_max_model_len] = PADDING_SLOT_ID - attn_metadata.slot_mapping = slot_mapping.view(-1) - - # Copy inputs to buffer for cudagraph. - num_tokens = attn_metadata.num_actual_tokens - input_ids = tree_input_ids.view(-1) - self.input_ids[:num_tokens] = input_ids - self.positions[:num_tokens] = tree_positions.view(-1) - self.hidden_states[:num_tokens] = tree_hidden_states.view(num_tokens, -1) - - cudagraph_runtime_mode, batch_desc = self.cudagraph_dispatcher.dispatch( - num_tokens - ) - num_input_tokens = batch_desc.num_tokens - # Run the model. - with set_forward_context( - per_layer_attn_metadata, - self.vllm_config, - num_tokens=num_input_tokens, - cudagraph_runtime_mode=cudagraph_runtime_mode, - slot_mapping=self._get_slot_mapping( - num_input_tokens, attn_metadata.slot_mapping - ), - ): - last_hidden_states, hidden_states = self.model( - input_ids=self.input_ids[:num_input_tokens], - positions=self.positions[:num_input_tokens], - hidden_states=self.hidden_states[:num_input_tokens], - inputs_embeds=None, - ) - - # Get the output hidden states for the draft tokens. - draft_hidden_states = hidden_states[:num_tokens].view( - batch_size, query_len, -1 - )[:, -level_num_drafts:] - draft_last_hidden_states = last_hidden_states[:num_tokens].view( - batch_size, query_len, -1 - )[:, -level_num_drafts:] - - # Get the output logits for the draft tokens. - logits = self.model.compute_logits( - draft_last_hidden_states.reshape(batch_size * level_num_drafts, -1) - ) - - # Sample a draft token for each child at the next tree level. - num_children = self.child_drafts_per_level[level + 1] - if num_children == 1: - draft_token_ids = logits.argmax(dim=-1).view(batch_size, -1) - else: - draft_token_ids = torch.topk(logits, num_children, dim=-1).indices.view( - batch_size, -1 - ) - draft_token_ids_list.append(draft_token_ids) - - # Update the # drafts counters for the next tree level. - level_num_drafts = self.cu_drafts_per_level[level + 1] - total_num_drafts - total_num_drafts = self.cu_drafts_per_level[level + 1] - return draft_token_ids_list - def prepare_inputs( self, common_attn_metadata: CommonAttentionMetadata,