mirror of
https://github.com/vllm-project/vllm.git
synced 2026-08-06 05:48:08 +00:00
[Attention][Cleanup] Remove tree attention (#42121)
Signed-off-by: Lucas Wilkinson <[email protected]>
This commit is contained in:
Lucas Wilkinson
@@ -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'
|
||||
|
||||
@@ -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 |
|
||||
|
||||
|
||||
+1
-1
@@ -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:
|
||||
|
||||
@@ -36,7 +36,6 @@ BACKENDS_TO_TEST = [
|
||||
AttentionBackendEnum.FLASHINFER,
|
||||
AttentionBackendEnum.FLEX_ATTENTION,
|
||||
AttentionBackendEnum.TRITON_ATTN,
|
||||
AttentionBackendEnum.TREE_ATTN,
|
||||
"FLEX_ATTENTION_SLOW",
|
||||
]
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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)
|
||||
@@ -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 = (
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -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
|
||||
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user