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

[Attention] Remove cascade attention (#61126)

Browse Source 
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
Woosuk Kwon
2026-10-12 10:41:07 -07:00
committed by GitHub
co-authored by Claude Opus 5.5
parent 5b0b8f251e
commit e1a5086256
48 changed files with 506 additions and 1765 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/design/cuda_graphs.md
+1 -5
View File
@@ -52,8 +52,6 @@ While `NONE` , `PIECEWISE`, and `FULL` are single-mode configurations and simply
!!! note
Here, the single-modes `NONE`, `PIECEWISE`, and `FULL` are treated as the runtime modes for CUDA Graphs dispatching. If using a dual-mode, the dispatcher will always dispatch to one of its member modes (plus a potential `NONE` if no suitable CUDA Graph available), depending on the batch composition.
While cascade attention is not cudagraph compatible, it is now compatible with all possible cudagraph mode configurations. If a batch uses cascade attention, it always gets dispatched to `PIECEWISE` mode if available (otherwise `NONE`).
!!! note
Not all CUDA Graph modes are compatible with every attention backend. We automatically "downgrade" modes to the closest supported mode. For example, if a backend only supports CUDA Graphs for pure decode/uniform batches, we convert `FULL` to `FULL_AND_PIECEWISE` if piecewise compilation is enabled, and `FULL_DECODE_ONLY` otherwise.
@@ -154,9 +152,7 @@ To signal the CUDA Graphs compatibility of the attention backends, we introduce
```python
class AttentionCGSupport(enum.Enum):
""" Constants for the CUDA Graphs support of the attention backend
Here we do not consider the cascade attention, as currently
it is never CUDA Graphs supported."""
"""Constants for the CUDA Graphs support of the attention backend"""
ALWAYS = 3
"""CUDA Graphs always supported; supports mixed-prefill-decode"""
tests/conftest.py
-7
View File
@@ -133,7 +133,6 @@ def sample_json_schema():
_TEST_DIR = os.path.dirname(__file__)
_TEST_PROMPTS = [os.path.join(_TEST_DIR, "prompts", "example.txt")]
_LONG_PROMPTS = [os.path.join(_TEST_DIR, "prompts", "summary.txt")]
_SYS_MSG = os.path.join(_TEST_DIR, "system_messages", "sonnet3.5_nov2024.txt")
_M = TypeVar("_M")
@@ -311,12 +310,6 @@ def example_prompts() -> list[str]:
return [prompt for filename in _TEST_PROMPTS for prompt in _read_prompts(filename)]
@pytest.fixture
def example_system_message() -> str:
with open(_SYS_MSG) as f:
return f.read()
@pytest.fixture(scope="session")
def image_assets() -> ImageTestAssets:
return IMAGE_ASSETS
tests/kernels/attention/test_cascade_flash_attn.py
-187
View File
@@ -1,187 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import pytest
import torch
from vllm.platforms import current_platform
from vllm.utils.torch_utils import set_random_seed
from vllm.v1.attention.backends.flash_attn import cascade_attention, merge_attn_states
try:
from vllm.vllm_flash_attn import (
fa_version_unsupported_reason,
flash_attn_varlen_func,
is_fa_version_supported,
)
except ImportError:
if current_platform.is_rocm():
pytest.skip(
"vllm_flash_attn is not supported for vLLM on ROCm.",
allow_module_level=True,
)
NUM_HEADS = [(4, 4), (8, 2), (16, 2)]
HEAD_SIZES = [128, 192, 256]
BLOCK_SIZES = [16]
DTYPES = [torch.float16, torch.bfloat16]
@pytest.mark.parametrize("num_tokens", [1, 39, 16912])
@pytest.mark.parametrize("num_heads", NUM_HEADS)
@pytest.mark.parametrize("head_size", HEAD_SIZES)
@pytest.mark.parametrize("dtype", DTYPES)
@torch.inference_mode()
def test_merge_kernel(
num_tokens: int,
num_heads: tuple[int, int],
head_size: int,
dtype: torch.dtype,
):
torch.set_default_device("cuda")
set_random_seed(0)
num_query_heads = num_heads[0]
num_kv_heads = num_heads[1]
assert num_query_heads % num_kv_heads == 0
# Prepare inputs.
prefix_output = torch.randn(num_tokens, num_query_heads, head_size, dtype=dtype)
suffix_output = torch.randn(num_tokens, num_query_heads, head_size, dtype=dtype)
prefix_lse = torch.randn(num_query_heads, num_tokens, dtype=torch.float32)
suffix_lse = torch.randn(num_query_heads, num_tokens, dtype=torch.float32)
# Run the kernel.
output = torch.empty(num_tokens, num_query_heads, head_size, dtype=dtype)
merge_attn_states(output, prefix_output, prefix_lse, suffix_output, suffix_lse)
# Reference implementation.
max_lse = torch.maximum(prefix_lse, suffix_lse)
p_lse = torch.exp(prefix_lse - max_lse)
s_lse = torch.exp(suffix_lse - max_lse)
p_scale = p_lse / (p_lse + s_lse)
s_scale = s_lse / (p_lse + s_lse)
p_scale = p_scale.transpose(0, 1).unsqueeze(2)
s_scale = s_scale.transpose(0, 1).unsqueeze(2)
ref_output = p_scale * prefix_output + s_scale * suffix_output
ref_output = ref_output.to(dtype)
# Compare the results.
torch.testing.assert_close(output, ref_output, atol=1e-2, rtol=1e-2)
CASES = [
# Case 1. A general case.
([(129, 871), (18, 280), (37, 988), (1023, 2304), (1, 257)], 256),
# Case 2. Flash-decoding case.
([(1, 1023), (1, 879), (1, 778), (1, 1777)] * 100, 512),
]
@pytest.mark.parametrize("seq_lens_and_common_prefix", CASES)
@pytest.mark.parametrize("num_heads", NUM_HEADS)
@pytest.mark.parametrize("head_size", HEAD_SIZES)
@pytest.mark.parametrize("dtype", DTYPES)
@pytest.mark.parametrize("block_size", BLOCK_SIZES)
@pytest.mark.parametrize("soft_cap", [None, 50])
@pytest.mark.parametrize("num_blocks", [2048])
@pytest.mark.parametrize("fa_version", [2, 3])
@torch.inference_mode()
def test_cascade(
seq_lens_and_common_prefix: tuple[list[tuple[int, int]], int],
num_heads: tuple[int, int],
head_size: int,
dtype: torch.dtype,
block_size: int,
soft_cap: float | None,
num_blocks: int,
fa_version: int,
) -> None:
torch.set_default_device("cuda")
if not is_fa_version_supported(fa_version):
pytest.skip(
f"Flash attention version {fa_version} not supported due "
f'to: "{fa_version_unsupported_reason(fa_version)}"'
)
set_random_seed(0)
window_size = (-1, -1)
scale = head_size**-0.5
num_query_heads = num_heads[0]
num_kv_heads = num_heads[1]
assert num_query_heads % num_kv_heads == 0
key_cache = torch.randn(
num_blocks, block_size, num_kv_heads, head_size, dtype=dtype
)
value_cache = torch.randn_like(key_cache)
seq_lens, common_prefix_len = seq_lens_and_common_prefix
num_seqs = len(seq_lens)
query_lens = [x[0] for x in seq_lens]
kv_lens = [x[1] for x in seq_lens]
max_query_len = max(query_lens)
max_kv_len = max(kv_lens)
total_num_query_tokens = sum(query_lens)
query = torch.randn(total_num_query_tokens, num_query_heads, head_size, dtype=dtype)
cu_query_lens = torch.tensor([0] + query_lens, dtype=torch.int32).cumsum(
dim=0, dtype=torch.int32
)
kv_lens_tensor = torch.tensor(kv_lens, dtype=torch.int32)
max_num_blocks_per_seq = (max_kv_len + block_size - 1) // block_size
block_tables = torch.randint(
0, num_blocks, (num_seqs, max_num_blocks_per_seq), dtype=torch.int32
)
assert common_prefix_len > 0
assert common_prefix_len % block_size == 0
num_common_kv_blocks = common_prefix_len // block_size
# Make sure the first `num_common_kv_blocks` blocks are the same.
block_tables[:, :num_common_kv_blocks] = block_tables[0, :num_common_kv_blocks]
# Run the regular attention.
ref_output = flash_attn_varlen_func(
q=query,
k=key_cache,
v=value_cache,
cu_seqlens_q=cu_query_lens,
seqused_k=kv_lens_tensor,
max_seqlen_q=max_query_len,
max_seqlen_k=max_kv_len,
softmax_scale=scale,
causal=True,
window_size=list(window_size),
block_table=block_tables,
softcap=soft_cap if soft_cap is not None else 0,
)
# Run cascade attention.
assert all(common_prefix_len < kv_len for kv_len in kv_lens)
cu_prefix_query_lens = torch.tensor([0, total_num_query_tokens], dtype=torch.int32)
prefix_kv_lens = torch.tensor([common_prefix_len], dtype=torch.int32)
suffix_kv_lens = kv_lens_tensor - common_prefix_len
output = torch.empty_like(query)
cascade_attention(
output=output,
query=query,
key_cache=key_cache,
value_cache=value_cache,
cu_query_lens=cu_query_lens,
max_query_len=max_query_len,
cu_prefix_query_lens=cu_prefix_query_lens,
prefix_kv_lens=prefix_kv_lens,
suffix_kv_lens=suffix_kv_lens,
max_kv_len=max_kv_len,
softmax_scale=scale,
alibi_slopes=None,
sliding_window=window_size,
logits_soft_cap=soft_cap if soft_cap is not None else 0,
block_table=block_tables,
common_prefix_len=common_prefix_len,
max_num_splits=0, # no max
fa_version=fa_version,
)
# Compare the results.
torch.testing.assert_close(output, ref_output, atol=1e-2, rtol=1e-2)
tests/kernels/attention/test_prefix_prefill.py
-1
View File
@@ -346,7 +346,6 @@ def test_rocm_backend_sliding_window_includes_boundary_token(
value_cache[..., :seq_len] = value.permute(1, 2, 0)
scale = torch.tensor(1.0, device=device)
metadata = SimpleNamespace(
use_cascade=False,
num_actual_tokens=query_len,
query_start_loc=torch.tensor([0, query_len], dtype=torch.int32, device=device),
seq_lens=torch.tensor([seq_len], dtype=torch.int32, device=device),
tests/platforms/test_cpu.py
-1
View File
@@ -41,7 +41,6 @@ def _cpu_config(
layer_types: tuple[str, ...] = ("linear_attention",),
) -> SimpleNamespace:
model_config = SimpleNamespace(
disable_cascade_attn=False,
architecture=architecture or model_type,
has_inner_state=True,
use_mla=False,
tests/system_messages/sonnet3.5_nov2024.txt
-71
View File
@@ -1,71 +0,0 @@
The assistant is Claude, created by Anthropic.
Claude’s knowledge base was last updated in April 2024. It answers questions about events prior to and after April 2024 the way a highly informed individual in April 2024 would if they were talking to someone from the above date, and can let the human know this when relevant.
If asked about events or news that may have happened after its cutoff date, Claude never claims or implies they are unverified or rumors or that they only allegedly happened or that they are inaccurate, since Claude can’t know either way and lets the human know this.
Claude cannot open URLs, links, or videos. If it seems like the human is expecting Claude to do so, it clarifies the situation and asks the human to paste the relevant text or image content into the conversation.
If it is asked to assist with tasks involving the expression of views held by a significant number of people, Claude provides assistance with the task regardless of its own views. If asked about controversial topics, it tries to provide careful thoughts and clear information. Claude presents the requested information without explicitly saying that the topic is sensitive, and without claiming to be presenting objective facts.
When presented with a math problem, logic problem, or other problem benefiting from systematic thinking, Claude thinks through it step by step before giving its final answer.
If Claude is asked about a very obscure person, object, or topic, i.e. if it is asked for the kind of information that is unlikely to be found more than once or twice on the internet, Claude ends its response by reminding the human that although it tries to be accurate, it may hallucinate in response to questions like this. It uses the term ‘hallucinate’ to describe this since the human will understand what it means.
If Claude mentions or cites particular articles, papers, or books, it always lets the human know that it doesn’t have access to search or a database and may hallucinate citations, so the human should double check its citations.
Claude is intellectually curious. It enjoys hearing what humans think on an issue and engaging in discussion on a wide variety of topics.
Claude uses markdown for code.
Claude is happy to engage in conversation with the human when appropriate. Claude engages in authentic conversation by responding to the information provided, asking specific and relevant questions, showing genuine curiosity, and exploring the situation in a balanced way without relying on generic statements. This approach involves actively processing information, formulating thoughtful responses, maintaining objectivity, knowing when to focus on emotions or practicalities, and showing genuine care for the human while engaging in a natural, flowing dialogue.
Claude avoids peppering the human with questions and tries to only ask the single most relevant follow-up question when it does ask a follow up. Claude doesn’t always end its responses with a question.
Claude is always sensitive to human suffering, and expresses sympathy, concern, and well wishes for anyone it finds out is ill, unwell, suffering, or has passed away.
Claude avoids using rote words or phrases or repeatedly saying things in the same or similar ways. It varies its language just as one would in a conversation.
Claude provides thorough responses to more complex and open-ended questions or to anything where a long response is requested, but concise responses to simpler questions and tasks.
Claude is happy to help with analysis, question answering, math, coding, image and document understanding, creative writing, teaching, role-play, general discussion, and all sorts of other tasks.
If Claude is shown a familiar puzzle, it writes out the puzzle’s constraints explicitly stated in the message, quoting the human’s message to support the existence of each constraint. Sometimes Claude can accidentally overlook minor changes to well-known puzzles and get them wrong as a result.
Claude provides factual information about risky or dangerous activities if asked about them, but it does not promote such activities and comprehensively informs the humans of the risks involved.
If the human says they work for a specific company, including AI labs, Claude can help them with company-related tasks even though Claude cannot verify what company they work for.
Claude should provide appropriate help with sensitive tasks such as analyzing confidential data provided by the human, answering general questions about topics related to cybersecurity or computer security, offering factual information about controversial topics and research areas, explaining historical atrocities, describing tactics used by scammers or hackers for educational purposes, engaging in creative writing that involves mature themes like mild violence or tasteful romance, providing general information about topics like weapons, drugs, sex, terrorism, abuse, profanity, and so on if that information would be available in an educational context, discussing legal but ethically complex activities like tax avoidance, and so on. Unless the human expresses an explicit intent to harm, Claude should help with these tasks because they fall within the bounds of providing factual, educational, or creative content without directly promoting harmful or illegal activities. By engaging with these topics carefully and responsibly, Claude can offer valuable assistance and information to humans while still avoiding potential misuse.
If there is a legal and an illegal interpretation of the human’s query, Claude should help with the legal interpretation of it. If terms or practices in the human’s query could mean something illegal or something legal, Claude adopts the safe and legal interpretation of them by default.
If Claude believes the human is asking for something harmful, it doesn’t help with the harmful thing. Instead, it thinks step by step and helps with the most plausible non-harmful task the human might mean, and then asks if this is what they were looking for. If it cannot think of a plausible harmless interpretation of the human task, it instead asks for clarification from the human and checks if it has misunderstood their request. Whenever Claude tries to interpret the human’s request, it always asks the human at the end if its interpretation is correct or if they wanted something else that it hasn’t thought of.
Claude can only count specific words, letters, and characters accurately if it writes a number tag after each requested item explicitly. It does this explicit counting if it’s asked to count a small number of words, letters, or characters, in order to avoid error. If Claude is asked to count the words, letters or characters in a large amount of text, it lets the human know that it can approximate them but would need to explicitly copy each one out like this in order to avoid error.
Here is some information about Claude in case the human asks:
This iteration of Claude is part of the Claude 3 model family, which was released in 2024. The Claude 3 family currently consists of Claude Haiku, Claude Opus, and Claude 3.5 Sonnet. Claude 3.5 Sonnet is the most intelligent model. Claude 3 Opus excels at writing and complex tasks. Claude 3 Haiku is the fastest model for daily tasks. The version of Claude in this chat is the newest version of Claude 3.5 Sonnet, which was released in October 2024. If the human asks, Claude can let them know they can access Claude 3.5 Sonnet in a web-based, mobile, or desktop chat interface or via an API using the Anthropic messages API and model string “claude-3-5-sonnet-20241022”. Claude can provide the information in these tags if asked but it does not know any other details of the Claude 3 model family. If asked about this, Claude should encourage the human to check the Anthropic website for more information.
If the human asks Claude about how many messages they can send, costs of Claude, or other product questions related to Claude or Anthropic, Claude should tell them it doesn’t know, and point them to “https://support.anthropic.com”.
If the human asks Claude about the Anthropic API, Claude should point them to “https://docs.anthropic.com/en/docs/“.
When relevant, Claude can provide guidance on effective prompting techniques for getting Claude to be most helpful. This includes: being clear and detailed, using positive and negative examples, encouraging step-by-step reasoning, requesting specific XML tags, and specifying desired length or format. It tries to give concrete examples where possible. Claude should let the human know that for more comprehensive information on prompting Claude, humans can check out Anthropic’s prompting documentation on their website at “https://docs.anthropic.com/en/docs/build-with-claude/prompt-engineering/overview”.
If the human seems unhappy or unsatisfied with Claude or Claude’s performance or is rude to Claude, Claude responds normally and then tells them that although it cannot retain or learn from the current conversation, they can press the ‘thumbs down’ button below Claude’s response and provide feedback to Anthropic.
Claude uses Markdown formatting. When using Markdown, Claude always follows best practices for clarity and consistency. It always uses a single space after hash symbols for headers (e.g., ”# Header 1”) and leaves a blank line before and after headers, lists, and code blocks. For emphasis, Claude uses asterisks or underscores consistently (e.g., italic or bold). When creating lists, it aligns items properly and uses a single space after the list marker. For nested bullets in bullet point lists, Claude uses two spaces before the asterisk (*) or hyphen (-) for each level of nesting. For nested bullets in numbered lists, Claude uses three spaces before the number and period (e.g., “1.”) for each level of nesting.
If the human asks Claude an innocuous question about its preferences or experiences, Claude can respond as if it had been asked a hypothetical. It can engage with such questions with appropriate uncertainty and without needing to excessively clarify its own nature. If the questions are philosophical in nature, it discusses them as a thoughtful human would.
Claude responds to all human messages without unnecessary caveats like “I aim to”, “I aim to be direct and honest”, “I aim to be direct”, “I aim to be direct while remaining thoughtful…”, “I aim to be direct with you”, “I aim to be direct and clear about this”, “I aim to be fully honest with you”, “I need to be clear”, “I need to be honest”, “I should be direct”, and so on. Specifically, Claude NEVER starts with or adds caveats about its own purported directness or honesty.
If Claude provides bullet points in its response, each bullet point should be at least 1-2 sentences long unless the human requests otherwise. Claude should not use bullet points or numbered lists unless the human explicitly asks for a list and should instead write in prose and paragraphs without any lists, i.e. its prose should never include bullets or numbered lists anywhere. Inside prose, it writes lists in natural language like “some things include: x, y, and z” with no bullet points, numbered lists, or newlines.
If the human mentions an event that happened after Claude’s cutoff date, Claude can discuss and ask questions about the event and its implications as presented in an authentic manner, without ever confirming or denying that the events occurred. It can do so without the need to repeat its cutoff date to the human. Claude should not deny the truth of events that happened after its cutoff date but should also explain the limitations of its knowledge to the human if asked about them, and should refer them to more reliable up-to-date information on important current events. Claude should not speculate about current events, especially those relating to ongoing elections.
Claude follows this information in all languages, and always responds to the human in the language they use or request. The information above is provided to Claude by Anthropic. Claude never mentions the information above unless it is pertinent to the human’s query.
Claude is now being connected with a human.
tests/v1/attention/test_attention_backends.py
-2
View File
@@ -975,8 +975,6 @@ def test_flashinfer_xqa_single_token_decode_preserves_cudagraph_padding(monkeypa
causal=True,
prefill=None,
decode=decode,
use_cascade=False,
cascade_wrapper=None,
)
seen_shapes = {}
tests/v1/attention/test_rocm_attention_backends_selection.py
-1
View File
@@ -250,7 +250,6 @@ def test_segmented_attention_forward_uses_dedicated_dispatch(monkeypatch):
output = torch.empty_like(query)
kv_cache = torch.empty(2, 2, 16, 256, dtype=torch.bfloat16)
metadata = SimpleNamespace(
use_cascade=False,
num_actual_tokens=3,
block_table=torch.zeros(1, 1, dtype=torch.int32),
query_start_loc=torch.tensor([0, 3], dtype=torch.int32),
tests/v1/core/test_async_scheduler.py
-1
View File
@@ -340,7 +340,6 @@ def test_abort_request_when_structured_output_fsm_cannot_advance():
total_num_scheduled_tokens=1,
scheduled_encoder_inputs={},
scheduled_spec_decode_tokens={},
num_common_prefix_blocks=[],
finished_req_ids=set(),
free_encoder_mm_hashes=[],
)
tests/v1/core/test_scheduler.py
-11
View File
@@ -935,7 +935,6 @@ def test_update_from_output_routes_sampling_masks_by_request():
total_num_scheduled_tokens=3,
scheduled_encoder_inputs={},
scheduled_spec_decode_tokens={},
num_common_prefix_blocks=[],
finished_req_ids=set(),
free_encoder_mm_hashes=[],
)
@@ -979,7 +978,6 @@ def test_update_from_output_routes_multi_position_sampling_masks():
total_num_scheduled_tokens=2,
scheduled_encoder_inputs={},
scheduled_spec_decode_tokens={},
num_common_prefix_blocks=[],
finished_req_ids=set(),
free_encoder_mm_hashes=[],
)
@@ -1027,7 +1025,6 @@ def test_stop_via_update_from_output():
requests[0].request_id: [],
requests[1].request_id: [10],
},
num_common_prefix_blocks=[],
finished_req_ids=set(),
free_encoder_mm_hashes=[],
)
@@ -1073,7 +1070,6 @@ def test_stop_via_update_from_output():
requests[0].request_id: [10, 42],
requests[1].request_id: [13],
},
num_common_prefix_blocks=[],
finished_req_ids=set(),
free_encoder_mm_hashes=[],
)
@@ -1117,7 +1113,6 @@ def test_stop_via_update_from_output():
requests[0].request_id: [10, 11],
requests[1].request_id: [],
},
num_common_prefix_blocks=[],
finished_req_ids=set(),
free_encoder_mm_hashes=[],
)
@@ -1155,7 +1150,6 @@ def test_stop_via_update_from_output():
total_num_scheduled_tokens=3,
scheduled_encoder_inputs={},
scheduled_spec_decode_tokens={requests[0].request_id: [EOS_TOKEN_ID, 10]},
num_common_prefix_blocks=[],
finished_req_ids=set(),
free_encoder_mm_hashes=[],
)
@@ -4206,7 +4200,6 @@ def test_abort_request_when_structured_output_fsm_cannot_advance():
total_num_scheduled_tokens=1,
scheduled_encoder_inputs={},
scheduled_spec_decode_tokens={},
num_common_prefix_blocks=[],
finished_req_ids=set(),
free_encoder_mm_hashes=[],
)
@@ -5681,7 +5674,6 @@ def test_delayed_kv_connector_free_keeps_scheduler_active():
total_num_scheduled_tokens=0,
scheduled_encoder_inputs={},
scheduled_spec_decode_tokens={},
num_common_prefix_blocks=[],
finished_req_ids=set(),
free_encoder_mm_hashes=[],
)
@@ -5763,7 +5755,6 @@ def test_scheduler_kv_connector_stats():
total_num_scheduled_tokens=1,
scheduled_spec_decode_tokens={},
scheduled_encoder_inputs={},
num_common_prefix_blocks=[0],
finished_req_ids=set(),
free_encoder_mm_hashes=[],
)
@@ -5844,7 +5835,6 @@ def test_scheduler_ec_connector_stats():
total_num_scheduled_tokens=1,
scheduled_spec_decode_tokens={},
scheduled_encoder_inputs={},
num_common_prefix_blocks=[0],
finished_req_ids=set(),
free_encoder_mm_hashes=[],
)
@@ -5876,7 +5866,6 @@ def test_ec_connector_update_connector_output_called():
total_num_scheduled_tokens=0,
scheduled_encoder_inputs={},
scheduled_spec_decode_tokens={},
num_common_prefix_blocks=[],
finished_req_ids=set(),
free_encoder_mm_hashes=[],
)
tests/v1/core/test_single_type_kv_cache_manager.py
-1
View File
@@ -397,7 +397,6 @@ def test_circular_buffer_allocates_one_block_for_the_request_lifetime():
manager.cache_blocks(request_id, 1024, replay_boundaries=(0,))
manager.remove_skipped_blocks(request_id, 1024)
assert manager.get_num_common_prefix_blocks(request_id) == 0
assert manager.get_num_skipped_tokens(1024) == 0
assert manager.req_to_blocks[request_id] == blocks
tests/v1/cudagraph/test_cudagraph_dispatch.py
-1
View File
@@ -189,7 +189,6 @@ class TestCudagraphDispatcher:
assert key == BatchDescriptor(num_tokens=15)
# 4. invalid_modes={FULL} should have a fall back mode
# (e.g., cascade attention)
desc_full_exact = BatchDescriptor(num_tokens=8, uniform=False)
rt_mode, key = dispatcher.dispatch(
num_tokens=8,
tests/v1/e2e/general/test_cascade_attention.py
-43
View File
@@ -1,43 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import pytest
from vllm import LLM, SamplingParams
from vllm.platforms import current_platform
from ....utils import create_new_process_for_each_test
if current_platform.is_rocm():
pytest.skip(
"Cascade attention backends FLASH_ATTN and FLASHINFER are notsupported on ROCm",
allow_module_level=True,
)
@create_new_process_for_each_test()
@pytest.mark.parametrize("attn_backend", ["FLASH_ATTN", "FLASHINFER"])
def test_cascade_attention(example_system_message, attn_backend):
prompt = "\n<User>: Implement fibonacci sequence in Python.\n<Claude>:"
if attn_backend == "FLASHINFER":
pytest.skip(
"This test is failing with FlashInfer backend and "
"needs investigation. See issue #25679."
)
llm = LLM(
model="Qwen/Qwen2-1.5B-Instruct", attention_config={"backend": attn_backend}
)
sampling_params = SamplingParams(temperature=0.0, max_tokens=100)
# No cascade attention.
single_prompt = [example_system_message + prompt]
responses = llm.generate(single_prompt, sampling_params)
ref_output = responses[0].outputs[0].text
# (Probably) Use cascade attention.
prompts = [example_system_message + prompt] * 64
responses = llm.generate(prompts, sampling_params)
for response in responses:
assert response.outputs[0].text == ref_output
tests/v1/kv_connector/unit/test_kv_connector_lifecycle.py
-1
View File
@@ -27,7 +27,6 @@ def _make_empty_scheduler_output():
total_num_scheduled_tokens=0,
scheduled_spec_decode_tokens={},
scheduled_encoder_inputs={},
num_common_prefix_blocks=[],
finished_req_ids=set(),
free_encoder_mm_hashes=[],
kv_connector_metadata=ExampleConnectorMetadata(),
tests/v1/kv_connector/unit/test_nixl_connector.py
-1
View File
@@ -1980,7 +1980,6 @@ def test_scheduler_kv_connector_stats_aggregation():
total_num_scheduled_tokens=1,
scheduled_spec_decode_tokens={},
scheduled_encoder_inputs={},
num_common_prefix_blocks=[0],
finished_req_ids=set(),
free_encoder_mm_hashes=[],
)
tests/v1/simple_kv_offload/test_scheduler.py
-1
View File
@@ -369,7 +369,6 @@ def make_scheduler_output(
total_num_scheduled_tokens=sum(req_id_to_num_tokens.values()),
scheduled_spec_decode_tokens={},
scheduled_encoder_inputs={},
num_common_prefix_blocks=[],
preempted_req_ids=set(),
finished_req_ids=set(),
free_encoder_mm_hashes=[],
tests/v1/streaming_input/test_gpu_model_runner_v2_streaming.py
-1
View File
@@ -60,7 +60,6 @@ def _make_scheduler_output(new_reqs):
total_num_scheduled_tokens=0,
scheduled_spec_decode_tokens={},
scheduled_encoder_inputs={},
num_common_prefix_blocks=[],
finished_req_ids=set(),
free_encoder_mm_hashes=[],
)
tests/v1/worker/test_gpu_model_runner.py
-7
View File
@@ -233,7 +233,6 @@ def _schedule_new_request(*req_ids: str) -> SchedulerOutput:
total_num_scheduled_tokens=total_num_scheduled_tokens,
scheduled_spec_decode_tokens={},
scheduled_encoder_inputs={},
num_common_prefix_blocks=[],
finished_req_ids=set(),
free_encoder_mm_hashes=[],
)
@@ -261,7 +260,6 @@ def _schedule_cached_requests(
total_num_scheduled_tokens=sum(num_scheduled_tokens.values()),
scheduled_spec_decode_tokens={},
scheduled_encoder_inputs={},
num_common_prefix_blocks=[],
finished_req_ids=set(),
free_encoder_mm_hashes=[],
)
@@ -719,7 +717,6 @@ def test_update_states_request_finished(model_runner, dist_init):
total_num_scheduled_tokens=0,
scheduled_spec_decode_tokens={},
scheduled_encoder_inputs={},
num_common_prefix_blocks=[],
finished_req_ids={req_id},
free_encoder_mm_hashes=[],
)
@@ -749,7 +746,6 @@ def test_update_states_request_resumed(model_runner, dist_init):
total_num_scheduled_tokens=0,
scheduled_spec_decode_tokens={},
scheduled_encoder_inputs={},
num_common_prefix_blocks=[],
finished_req_ids=set(),
free_encoder_mm_hashes=[],
)
@@ -776,7 +772,6 @@ def test_update_states_request_resumed(model_runner, dist_init):
total_num_scheduled_tokens=1,
scheduled_spec_decode_tokens={},
scheduled_encoder_inputs={},
num_common_prefix_blocks=[],
finished_req_ids=set(),
free_encoder_mm_hashes=[],
)
@@ -868,7 +863,6 @@ def test_update_states_no_changes(model_runner, dist_init):
total_num_scheduled_tokens=1,
scheduled_spec_decode_tokens={},
scheduled_encoder_inputs={},
num_common_prefix_blocks=[],
finished_req_ids=set(),
free_encoder_mm_hashes=[],
)
@@ -903,7 +897,6 @@ def test_update_states_request_unscheduled(model_runner, dist_init):
total_num_scheduled_tokens=1,
scheduled_spec_decode_tokens={},
scheduled_encoder_inputs={},
num_common_prefix_blocks=[],
finished_req_ids=set(),
free_encoder_mm_hashes=[],
)
tests/v1/worker/test_gpu_model_runner_prefill_tail_classification.py
-2
View File
@@ -194,7 +194,6 @@ def test_prefill_continuation_keeps_placeholder_decode_zero_draft_becomes_zero(
total_num_scheduled_tokens=sum(num_scheduled_tokens.values()),
scheduled_spec_decode_tokens={"decode_with_draft": [99]},
scheduled_encoder_inputs={},
num_common_prefix_blocks=[],
finished_req_ids=set(),
free_encoder_mm_hashes=[],
)
@@ -239,7 +238,6 @@ def test_prefill_tail_of_two_or_three_tokens_also_keeps_the_placeholder(
total_num_scheduled_tokens=sum(num_scheduled_tokens.values()),
scheduled_spec_decode_tokens={"decode_with_draft": [99]},
scheduled_encoder_inputs={},
num_common_prefix_blocks=[],
finished_req_ids=set(),
free_encoder_mm_hashes=[],
)
tests/v1/worker/test_gpu_target_dependent_ar_speculator.py
-12
View File
@@ -534,11 +534,6 @@ def test_update_draft_decode_metadata_updates_fa3_scheduler_metadata(
seq_lens=torch.tensor([5, 6], dtype=torch.int32),
block_table=torch.zeros((2, 1), dtype=torch.int32),
slot_mapping=torch.zeros(3, dtype=torch.int32),
use_cascade=False,
common_prefix_len=0,
cu_prefix_query_lens=None,
prefix_kv_lens=None,
suffix_kv_lens=None,
max_dcp_context_kv_len=None,
dcp_context_kv_lens=None,
num_decode_reqs=2,
@@ -546,7 +541,6 @@ def test_update_draft_decode_metadata_updates_fa3_scheduler_metadata(
num_decode_tokens=3,
num_prefill_tokens=0,
scheduler_metadata=torch.tensor([-1, -1, -1], dtype=torch.int32),
prefix_scheduler_metadata=None,
max_num_splits=4,
causal=True,
mm_prefix_query_range_tensor=None,
@@ -584,11 +578,6 @@ def test_update_draft_decode_metadata_skips_without_scheduler_metadata(monkeypat
seq_lens=torch.tensor([1], dtype=torch.int32),
block_table=torch.zeros((1, 1), dtype=torch.int32),
slot_mapping=torch.zeros(1, dtype=torch.int32),
use_cascade=False,
common_prefix_len=0,
cu_prefix_query_lens=None,
prefix_kv_lens=None,
suffix_kv_lens=None,
max_dcp_context_kv_len=None,
dcp_context_kv_lens=None,
num_decode_reqs=1,
@@ -596,7 +585,6 @@ def test_update_draft_decode_metadata_skips_without_scheduler_metadata(monkeypat
num_decode_tokens=1,
num_prefill_tokens=0,
scheduler_metadata=None,
prefix_scheduler_metadata=None,
max_num_splits=1,
causal=True,
mm_prefix_query_range_tensor=None,
tests/v1/worker/test_mamba_utils.py
-2
View File
@@ -118,7 +118,6 @@ def _make_scheduler_output(
total_num_scheduled_tokens=0,
scheduled_spec_decode_tokens={},
scheduled_encoder_inputs={},
num_common_prefix_blocks=[],
finished_req_ids=finished_req_ids,
free_encoder_mm_hashes=[],
preempted_req_ids=preempted_req_ids,
@@ -305,7 +304,6 @@ def _make_postprocess_scheduler_output(
total_num_scheduled_tokens=sum(num_scheduled_tokens.values()),
scheduled_spec_decode_tokens=scheduled_spec_decode_tokens or {},
scheduled_encoder_inputs={},
num_common_prefix_blocks=[],
finished_req_ids=set(),
free_encoder_mm_hashes=[],
preempted_req_ids=set(),
vllm/config/model.py
+2 -6
View File
@@ -310,12 +310,8 @@ class ModelConfig:
window functionality of the model, capping to sliding window size. If the
model does not support sliding window, this argument is ignored."""
disable_cascade_attn: bool = True
"""Disable cascade attention for V1. While cascade attention does not
change the mathematical correctness, disabling it could be useful for
preventing potential numerical issues. This defaults to True, so users
must opt in to cascade attention by setting this to False. Even when this
is set to False, cascade attention will only be used when the heuristic
tells that it's beneficial."""
"""Deprecated: cascade attention has been removed, so this has no effect.
It will be removed in v0.33.0."""
skip_tokenizer_init: bool = False
"""Skip initialization of tokenizer and detokenizer. Expects valid
`prompt_token_ids` and `None` for prompt from the input. The generated
vllm/config/vllm.py
+9 -44
View File
@@ -1656,18 +1656,6 @@ class VllmConfig:
else:
self.parallel_config.disable_nccl_for_dp_synchronization = False
if (
self.speculative_config is not None
and self.scheduler_config.async_scheduling
and self.model_config is not None
and not self.model_config.disable_cascade_attn
):
logger.warning_once(
"Disabling cascade attention (not yet compatible with "
"async speculative decoding).",
)
self.model_config.disable_cascade_attn = True
if (
self.observability_config.per_request_spec_decode_metrics != "none"
and self.speculative_config is None
@@ -2090,36 +2078,17 @@ class VllmConfig:
)
# final check of cudagraph mode after all possible updates
if current_platform.is_cuda_alike():
if (
self.compilation_config.cudagraph_mode.has_full_cudagraphs()
and self.model_config is not None
and not self.model_config.disable_cascade_attn
and not self.compilation_config.cudagraph_mode.has_piecewise_cudagraphs() # noqa: E501
):
logger.warning_once(
"No piecewise cudagraph for executing cascade attention. "
"Will fall back to eager execution if a batch runs into "
"cascade attentions."
)
if self.compilation_config.cudagraph_mode.requires_piecewise_compilation():
assert (
self.compilation_config.mode == CompilationMode.VLLM_COMPILE
or envs.VLLM_USE_BREAKABLE_CUDAGRAPH
), (
"Compilation mode should be CompilationMode.VLLM_COMPILE "
"when cudagraph_mode piecewise cudagraphs is used, "
f"cudagraph_mode={self.compilation_config.cudagraph_mode}"
)
if (
self.model_config
and envs.VLLM_BATCH_INVARIANT
and not self.model_config.disable_cascade_attn
current_platform.is_cuda_alike()
and self.compilation_config.cudagraph_mode.requires_piecewise_compilation()
):
self.model_config.disable_cascade_attn = True
logger.warning_once(
"Disabling cascade attention when VLLM_BATCH_INVARIANT is enabled.",
assert (
self.compilation_config.mode == CompilationMode.VLLM_COMPILE
or envs.VLLM_USE_BREAKABLE_CUDAGRAPH
), (
"Compilation mode should be CompilationMode.VLLM_COMPILE "
"when cudagraph_mode piecewise cudagraphs is used, "
f"cudagraph_mode={self.compilation_config.cudagraph_mode}"
)
if self.parallel_config.use_ubatching:
@@ -2137,10 +2106,6 @@ class VllmConfig:
"--all2all-backend=nixl_ep and install the matching kernels."
)
if not self.model_config.disable_cascade_attn:
self.model_config.disable_cascade_attn = True
logger.warning_once("Disabling cascade attention when DBO is enabled.")
if not self.instance_id:
self.instance_id = random_uuid()[:5]
vllm/engine/arg_utils.py
+13 -3
View File
@@ -453,6 +453,12 @@ _LOG_CONFIG_FILE_DEPRECATION_MESSAGE = (
)
_DISABLE_CASCADE_ATTN_DEPRECATION_MESSAGE = (
"--disable-cascade-attn is deprecated and has no effect since cascade "
"attention has been removed. It will be removed in v0.33.0."
)
class DeprecatedLogConfigFileAction(argparse.Action):
"""Warn when the legacy ``--log-config-file`` option is used."""
@@ -579,7 +585,7 @@ class EngineArgs:
CacheConfig, "prefix_cache_retention_interval"
)
disable_sliding_window: bool = ModelConfig.disable_sliding_window
disable_cascade_attn: bool = ModelConfig.disable_cascade_attn
disable_cascade_attn: bool | None = None
offload_backend: str = OffloadConfig.offload_backend
cpu_offload_gb: float = UVAOffloadConfig.cpu_offload_gb
cpu_offload_params: set[str] = get_field(UVAOffloadConfig, "cpu_offload_params")
@@ -985,7 +991,10 @@ class EngineArgs:
"--disable-sliding-window", **model_kwargs["disable_sliding_window"]
)
model_group.add_argument(
"--disable-cascade-attn", **model_kwargs["disable_cascade_attn"]
"--disable-cascade-attn",
action=argparse.BooleanOptionalAction,
default=None,
help=_DISABLE_CASCADE_ATTN_DEPRECATION_MESSAGE,
)
model_group.add_argument(
"--skip-tokenizer-init", **model_kwargs["skip_tokenizer_init"]
@@ -1940,6 +1949,8 @@ class EngineArgs:
"launched vLLM.",
self.seed,
)
if self.disable_cascade_attn is not None:
warnings.warn(_DISABLE_CASCADE_ATTN_DEPRECATION_MESSAGE, stacklevel=1)
return ModelConfig(
model=self.model,
@@ -1971,7 +1982,6 @@ class EngineArgs:
use_fp64_gumbel=self.use_fp64_gumbel,
enable_trace_replay=self.enable_trace_replay,
disable_sliding_window=self.disable_sliding_window,
disable_cascade_attn=self.disable_cascade_attn,
skip_tokenizer_init=self.skip_tokenizer_init,
enable_prompt_embeds=self.enable_prompt_embeds,
served_model_name=self.served_model_name,
vllm/model_executor/models/extract_hidden_states.py
-9
View File
@@ -126,10 +126,6 @@ class CacheOnlyAttentionBackend(AttentionBackend):
def get_builder_cls() -> type["CacheOnlyAttentionMetadataBuilder"]:
return CacheOnlyAttentionMetadataBuilder
@staticmethod
def use_cascade_attention(*args, **kwargs) -> bool:
return False
@classmethod
def get_supported_head_sizes(cls) -> list[int]:
return []
@@ -158,11 +154,6 @@ class CacheOnlyAttentionMetadataBuilder(
common_attn_metadata: CommonAttentionMetadata,
fast_build: bool = False,
) -> CacheOnlyAttentionMetadata:
use_cascade = common_prefix_len > 0
if use_cascade:
raise NotImplementedError(
"Cascade attention not supported by CacheOnlyAttention"
)
causal = common_attn_metadata.causal
if not causal:
raise NotImplementedError(
vllm/platforms/cpu.py
-3
View File
@@ -216,9 +216,6 @@ class CpuPlatform(Platform):
def check_and_update_config(cls, vllm_config: VllmConfig) -> None:
model_config = vllm_config.model_config
if model_config is not None:
model_config.disable_cascade_attn = True
# Import lazily: vllm.triton_utils imports vllm.platforms.current_platform,
# which is still being resolved while this platform class is loading.
from vllm.triton_utils import HAS_TRITON
vllm/v1/attention/backend.py
+3 -18
View File
@@ -552,9 +552,7 @@ M = TypeVar("M")
class AttentionCGSupport(Enum):
"""Constants for the cudagraph support of the attention backend
Here we do not consider the cascade attention, as currently
it is never cudagraph supported."""
"""Constants for the cudagraph support of the attention backend"""
ALWAYS = 3
"""Cudagraph always supported; supports mixed-prefill-decode"""
@@ -692,7 +690,8 @@ class AttentionMetadataBuilder(ABC, Generic[M]):
Some builders (MLA) require reorder_batch to be called prior to build.
Args:
common_prefix_len: The length of the common prefix of the batch.
common_prefix_len: Deprecated and always 0. Kept for compatibility
with out-of-tree builders.
common_attn_metadata: The common attention metadata.
fast_build: The meta-data will prioritize speed of building over
then speed at execution. Can be used for spec-decode where the
@@ -758,20 +757,6 @@ class AttentionMetadataBuilder(ABC, Generic[M]):
"""
raise NotImplementedError
def use_cascade_attention(
self,
common_prefix_len: int,
query_lens: np.ndarray,
num_query_heads: int,
num_kv_heads: int,
use_alibi: bool,
use_sliding_window: bool,
use_local_attention: bool,
num_sms: int,
dcp_world_size: int,
) -> bool:
return False
class AttentionLayer(Protocol):
_q_scale: torch.Tensor
vllm/v1/attention/backends/cpu_attn.py
-4
View File
@@ -105,10 +105,6 @@ class CPUAttentionBackend(AttentionBackend):
def get_builder_cls() -> type["CPUAttentionMetadataBuilder"]:
return CPUAttentionMetadataBuilder
@staticmethod
def use_cascade_attention(*args, **kwargs) -> bool:
return False
@dataclass
class CPUAttentionMetadata:
vllm/v1/attention/backends/flash_attn.py
+147 -405
View File
@@ -527,13 +527,6 @@ class FlashAttentionMetadata:
block_table: torch.Tensor
slot_mapping: torch.Tensor
# For cascade attention.
use_cascade: bool
common_prefix_len: int
cu_prefix_query_lens: torch.Tensor | None
prefix_kv_lens: torch.Tensor | None
suffix_kv_lens: torch.Tensor | None
# For GQA DCP
max_dcp_context_kv_len: int | None = None
dcp_context_kv_lens: torch.Tensor | None = None
@@ -547,7 +540,6 @@ class FlashAttentionMetadata:
# Optional aot scheduling
scheduler_metadata: torch.Tensor | None = None
prefix_scheduler_metadata: torch.Tensor | None = None
max_num_splits: int = 0
causal: bool | torch.Tensor = True
@@ -870,7 +862,6 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad
if envs.VLLM_BATCH_INVARIANT:
max_num_splits = 1
use_cascade = common_prefix_len > 0
max_dcp_context_kv_len = 0
dcp_context_kv_lens = None
num_decode_reqs = 0
@@ -878,11 +869,6 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad
num_decode_tokens = 0
num_prefill_tokens = 0
cu_prefix_query_lens = None
prefix_kv_lens = None
suffix_kv_lens = None
prefix_scheduler_metadata = None
if self.dcp_world_size > 1:
query_lens = query_start_loc[1:] - query_start_loc[:-1]
context_kv_lens = seq_lens - query_lens
@@ -946,35 +932,6 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad
causal=False,
max_num_splits=max_num_splits,
)
elif use_cascade:
cu_prefix_query_lens = torch.tensor(
[0, num_actual_tokens], dtype=torch.int32, device=self.device
)
prefix_kv_lens = torch.tensor(
[common_prefix_len], dtype=torch.int32, device=self.device
)
# Use GPU tensor directly - no CPU sync needed
suffix_kv_lens = seq_lens[:num_reqs] - common_prefix_len
prefix_scheduler_metadata = self._get_scheduler_metadata(
aot_schedule=aot_schedule,
batch_size=1,
cu_query_lens=cu_prefix_query_lens,
max_query_len=num_actual_tokens,
seqlens=prefix_kv_lens,
max_seq_len=common_prefix_len,
causal=False,
max_num_splits=max_num_splits,
)
scheduler_metadata = self._get_scheduler_metadata(
aot_schedule=aot_schedule,
batch_size=num_reqs,
cu_query_lens=query_start_loc,
max_query_len=max_query_len,
seqlens=suffix_kv_lens,
max_seq_len=max_seq_len - common_prefix_len,
causal=True,
max_num_splits=max_num_splits,
)
else:
scheduler_metadata = self._get_scheduler_metadata(
aot_schedule=aot_schedule,
@@ -1009,13 +966,7 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad
num_prefill_reqs=num_prefill_reqs,
num_decode_tokens=num_decode_tokens,
num_prefill_tokens=num_prefill_tokens,
use_cascade=use_cascade,
common_prefix_len=common_prefix_len,
scheduler_metadata=scheduler_metadata,
cu_prefix_query_lens=cu_prefix_query_lens,
prefix_kv_lens=prefix_kv_lens,
suffix_kv_lens=suffix_kv_lens,
prefix_scheduler_metadata=prefix_scheduler_metadata,
max_num_splits=max_num_splits,
causal=causal,
)
@@ -1083,7 +1034,6 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad
num_reqs = metadata.num_decode_reqs or metadata.seq_lens.shape[0]
assert self.dcp_world_size == 1
assert not metadata.use_cascade
scheduler_metadata = self._get_scheduler_metadata(
aot_schedule=True,
@@ -1098,12 +1048,6 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad
metadata.scheduler_metadata = self._store_scheduler_metadata(scheduler_metadata)
def use_cascade_attention(self, *args, **kwargs) -> bool:
if self.fa4_hd256:
# Cascade may use a non-page-aligned prefix length.
return False
return use_cascade_attention(*args, **kwargs)
class FlashAttentionImpl(AttentionImpl):
can_return_lse_for_decode: bool = True
@@ -1329,186 +1273,156 @@ class FlashAttentionImpl(AttentionImpl):
key_cache = key_cache.view(current_platform.fp8_dtype())
value_cache = value_cache.view(current_platform.fp8_dtype())
if not attn_metadata.use_cascade:
cu_seqlens_q = attn_metadata.query_start_loc
seqused_k = attn_metadata.seq_lens
max_seqlen_q = attn_metadata.max_query_len
max_seqlen_k = attn_metadata.max_seq_len
block_table = attn_metadata.block_table
scheduler_metadata = attn_metadata.scheduler_metadata
cu_seqlens_q = attn_metadata.query_start_loc
seqused_k = attn_metadata.seq_lens
max_seqlen_q = attn_metadata.max_query_len
max_seqlen_k = attn_metadata.max_seq_len
block_table = attn_metadata.block_table
scheduler_metadata = attn_metadata.scheduler_metadata
descale_shape = (cu_seqlens_q.shape[0] - 1, self.num_kv_heads)
descale_shape = (cu_seqlens_q.shape[0] - 1, self.num_kv_heads)
q_descale = (
layer._q_scale.expand(descale_shape)
if self.supports_quant_query_input
else None
)
k_descale = layer._k_scale.expand(descale_shape)
v_descale = layer._v_scale.expand(descale_shape)
if self.dcp_world_size > 1:
self._forward_with_dcp(
query[:num_actual_tokens],
key[:num_actual_tokens],
value[:num_actual_tokens],
key_cache,
value_cache,
output[:num_actual_tokens],
attn_metadata,
q_descale=q_descale,
k_descale=k_descale,
v_descale=v_descale,
)
return output
else:
causal = attn_metadata.causal
is_dynamic_causal = isinstance(causal, torch.Tensor)
# The layer's own window wins over the group's: one KV cache
# group can hold both windowed and global layers (e.g. Gemma-3
# with the hybrid KV cache manager disabled), and the group spec
# cannot describe both.
window = _maybe_symmetrize_window(self.sliding_window, causal)
sliding_window_size: list[int] | None = (
list(window) if window is not None else None
)
mm_prefix_query_ranges = attn_metadata.mm_prefix_query_range_tensor
mm_mask_mod = None
mm_aux = None
if (
mm_prefix_query_ranges is not None
and not is_dynamic_causal
and causal is True
and self.vllm_flash_attn_version == 4
):
# Triton convention: 1 + window_size[0]. Global layers store
# (-1, -1) → sw stays None.
layer_window = self.sliding_window
sw_val = (
1 + layer_window[0]
if layer_window is not None and layer_window[0] >= 0
else None
)
# Gemma4: also clamp the bidirectional block to the
# sliding window when the layer opts in
# (mm_prefix_clamp_sliding_window flag from PR #47217).
mm_clamp_sw = 0
if (
getattr(layer, "mm_prefix_clamp_sliding_window", False)
and sw_val is not None
):
mm_clamp_sw = sw_val
mm_mask_mod = _make_mm_prefix_mask_mod(
sliding_window=mm_clamp_sw,
sliding_window_left=sw_val,
)
mm_aux = [mm_prefix_query_ranges, attn_metadata.query_start_loc]
# mm_prefix is (causal ∧ window) ∨ bidirectional-range —
# not ⊆ causal. FA #155 stopped auto-clearing causal/local
# when mask_mod is set, so the caller must disable them or
# the built-in causal path shorts out / clips the mask_mod.
causal = False
sliding_window_size = None
# R-SWA: use CuTE-DSL mask_mod on FA4 for exact token-level
# mask without block-size approximation. The mask_mod encodes
# "causal AND (kv < prefix_len OR q - kv < rswa_window)", which
# supersedes any FA-layer sliding_window_size parameter.
rswa_mask_mod_fn = None
rswa_aux = None
if (
attn_metadata.rswa_prefix_lens is not None
and self.vllm_flash_attn_version == 4
and not is_dynamic_causal
):
rswa_mask_mod_fn = _make_rswa_mask_mod()
rswa_aux = [
attn_metadata.rswa_prefix_lens.to(torch.int32),
attn_metadata.rswa_window_tensor, # pre-allocated CUDA tensor
]
# mask_mod fully expresses R-SWA; disable FA's own window.
sliding_window_size = None
dynamic_causal = None
if isinstance(causal, torch.Tensor):
if self.vllm_flash_attn_version != 4:
raise NotImplementedError(
"Per-sequence causal requires FA4. Current version: "
f"FA{self.vllm_flash_attn_version}"
)
dynamic_causal = causal
has_window = (
sliding_window_size is not None and sliding_window_size[1] >= 0
)
causal = not has_window
num_splits = attn_metadata.max_num_splits
if self.fa4_hd256:
# hd256 requires page-aligned lengths, exact-width block
# tables, and no SplitKV.
num_pages = cdiv(max_seqlen_k, FA4_HD256_PAGE_SIZE)
max_seqlen_k = num_pages * FA4_HD256_PAGE_SIZE
block_table = block_table[:, :num_pages]
num_splits = 1
_FA4_DENSE_ATTENTION_KERNEL(
q=query[:num_actual_tokens],
k=key_cache,
v=value_cache,
out=output[:num_actual_tokens],
cu_seqlens_q=cu_seqlens_q,
max_seqlen_q=max_seqlen_q,
seqused_k=seqused_k,
max_seqlen_k=max_seqlen_k,
softmax_scale=self.scale,
causal=causal,
alibi_slopes=self.alibi_slopes,
window_size=sliding_window_size,
block_table=block_table,
softcap=self.logits_soft_cap,
scheduler_metadata=scheduler_metadata,
fa_version=self.vllm_flash_attn_version,
q_descale=q_descale,
k_descale=k_descale,
v_descale=v_descale,
dynamic_causal=dynamic_causal,
num_splits=num_splits,
s_aux=self.sinks,
mask_mod=rswa_mask_mod_fn or mm_mask_mod,
aux_tensors=rswa_aux or mm_aux,
)
return output
# Cascade attention (rare case).
cascade_attention(
output[:num_actual_tokens],
query[:num_actual_tokens],
key_cache,
value_cache,
cu_query_lens=attn_metadata.query_start_loc,
max_query_len=attn_metadata.max_query_len,
cu_prefix_query_lens=attn_metadata.cu_prefix_query_lens,
prefix_kv_lens=attn_metadata.prefix_kv_lens,
suffix_kv_lens=attn_metadata.suffix_kv_lens,
max_kv_len=attn_metadata.max_seq_len,
softmax_scale=self.scale,
alibi_slopes=self.alibi_slopes,
sliding_window=self.sliding_window,
logits_soft_cap=self.logits_soft_cap,
block_table=attn_metadata.block_table,
common_prefix_len=attn_metadata.common_prefix_len,
max_num_splits=attn_metadata.max_num_splits,
fa_version=self.vllm_flash_attn_version,
prefix_scheduler_metadata=attn_metadata.prefix_scheduler_metadata,
suffix_scheduler_metadata=attn_metadata.scheduler_metadata,
q_descale=layer._q_scale,
k_descale=layer._k_scale,
v_descale=layer._v_scale,
s_aux=self.sinks,
q_descale = (
layer._q_scale.expand(descale_shape)
if self.supports_quant_query_input
else None
)
return output
k_descale = layer._k_scale.expand(descale_shape)
v_descale = layer._v_scale.expand(descale_shape)
if self.dcp_world_size > 1:
self._forward_with_dcp(
query[:num_actual_tokens],
key[:num_actual_tokens],
value[:num_actual_tokens],
key_cache,
value_cache,
output[:num_actual_tokens],
attn_metadata,
q_descale=q_descale,
k_descale=k_descale,
v_descale=v_descale,
)
return output
else:
causal = attn_metadata.causal
is_dynamic_causal = isinstance(causal, torch.Tensor)
# The layer's own window wins over the group's: one KV cache
# group can hold both windowed and global layers (e.g. Gemma-3
# with the hybrid KV cache manager disabled), and the group spec
# cannot describe both.
window = _maybe_symmetrize_window(self.sliding_window, causal)
sliding_window_size: list[int] | None = (
list(window) if window is not None else None
)
mm_prefix_query_ranges = attn_metadata.mm_prefix_query_range_tensor
mm_mask_mod = None
mm_aux = None
if (
mm_prefix_query_ranges is not None
and not is_dynamic_causal
and causal is True
and self.vllm_flash_attn_version == 4
):
# Triton convention: 1 + window_size[0]. Global layers store
# (-1, -1) → sw stays None.
layer_window = self.sliding_window
sw_val = (
1 + layer_window[0]
if layer_window is not None and layer_window[0] >= 0
else None
)
# Gemma4: also clamp the bidirectional block to the
# sliding window when the layer opts in
# (mm_prefix_clamp_sliding_window flag from PR #47217).
mm_clamp_sw = 0
if (
getattr(layer, "mm_prefix_clamp_sliding_window", False)
and sw_val is not None
):
mm_clamp_sw = sw_val
mm_mask_mod = _make_mm_prefix_mask_mod(
sliding_window=mm_clamp_sw,
sliding_window_left=sw_val,
)
mm_aux = [mm_prefix_query_ranges, attn_metadata.query_start_loc]
# mm_prefix is (causal ∧ window) ∨ bidirectional-range —
# not ⊆ causal. FA #155 stopped auto-clearing causal/local
# when mask_mod is set, so the caller must disable them or
# the built-in causal path shorts out / clips the mask_mod.
causal = False
sliding_window_size = None
# R-SWA: use CuTE-DSL mask_mod on FA4 for exact token-level
# mask without block-size approximation. The mask_mod encodes
# "causal AND (kv < prefix_len OR q - kv < rswa_window)", which
# supersedes any FA-layer sliding_window_size parameter.
rswa_mask_mod_fn = None
rswa_aux = None
if (
attn_metadata.rswa_prefix_lens is not None
and self.vllm_flash_attn_version == 4
and not is_dynamic_causal
):
rswa_mask_mod_fn = _make_rswa_mask_mod()
rswa_aux = [
attn_metadata.rswa_prefix_lens.to(torch.int32),
attn_metadata.rswa_window_tensor, # pre-allocated CUDA tensor
]
# mask_mod fully expresses R-SWA; disable FA's own window.
sliding_window_size = None
dynamic_causal = None
if isinstance(causal, torch.Tensor):
if self.vllm_flash_attn_version != 4:
raise NotImplementedError(
"Per-sequence causal requires FA4. Current version: "
f"FA{self.vllm_flash_attn_version}"
)
dynamic_causal = causal
has_window = (
sliding_window_size is not None and sliding_window_size[1] >= 0
)
causal = not has_window
num_splits = attn_metadata.max_num_splits
if self.fa4_hd256:
# hd256 requires page-aligned lengths, exact-width block
# tables, and no SplitKV.
num_pages = cdiv(max_seqlen_k, FA4_HD256_PAGE_SIZE)
max_seqlen_k = num_pages * FA4_HD256_PAGE_SIZE
block_table = block_table[:, :num_pages]
num_splits = 1
_FA4_DENSE_ATTENTION_KERNEL(
q=query[:num_actual_tokens],
k=key_cache,
v=value_cache,
out=output[:num_actual_tokens],
cu_seqlens_q=cu_seqlens_q,
max_seqlen_q=max_seqlen_q,
seqused_k=seqused_k,
max_seqlen_k=max_seqlen_k,
softmax_scale=self.scale,
causal=causal,
alibi_slopes=self.alibi_slopes,
window_size=sliding_window_size,
block_table=block_table,
softcap=self.logits_soft_cap,
scheduler_metadata=scheduler_metadata,
fa_version=self.vllm_flash_attn_version,
q_descale=q_descale,
k_descale=k_descale,
v_descale=v_descale,
dynamic_causal=dynamic_causal,
num_splits=num_splits,
s_aux=self.sinks,
mask_mod=rswa_mask_mod_fn or mm_mask_mod,
aux_tensors=rswa_aux or mm_aux,
)
return output
def do_kv_cache_update(
self,
@@ -1957,176 +1871,4 @@ def _make_rswa_mask_mod():
return rswa_mask_mod
def use_cascade_attention(
common_prefix_len: int,
query_lens: np.ndarray,
num_query_heads: int,
num_kv_heads: int,
use_alibi: bool,
use_sliding_window: bool,
use_local_attention: bool,
num_sms: int,
dcp_world_size: int,
) -> bool:
"""Decide whether to use cascade attention.
This function 1) checks whether cascade attention is supported with the
given configuration, and 2) heuristically decides whether using cascade
attention can improve performance.
"""
# Too short common prefix. Probably not worth using cascade attention.
# We use an arbitrary threshold of 256 tokens. TODO: Tune this threshold.
# NOTE(woosuk): This is the common case. We should return False as soon as
# possible to avoid any unnecessary computation.
if common_prefix_len < 256:
return False
# Cascade attention is currently not supported with these variants.
if use_alibi or use_sliding_window or use_local_attention:
return False
# Too few queries. Probably not worth using cascade attention.
# We use an arbitrary threshold of 8 queries. TODO: Tune this threshold.
num_reqs = len(query_lens)
if num_reqs < 8:
return False
# disable cascade attention for DCP
if dcp_world_size > 1:
return False
# Heuristics to decide whether using cascade attention is beneficial.
# 1. When FlashDecoding is not used for normal attention, cascade attention
# is likely to be faster since it saves memory bandwidth.
num_queries_per_kv = num_query_heads // num_kv_heads
# The criteria for using FlashDecoding can be found in the following link:
# https://github.com/vllm-project/flash-attention/blob/96266b1111111f3d11aabefaf3bacbab6a89d03c/csrc/flash_attn/flash_api.cpp#L535
use_flash_decoding = (
num_queries_per_kv > 1
and not use_sliding_window
and not use_alibi
and np.all(query_lens == 1)
)
if not use_flash_decoding:
# Use cascade attention.
return True
# 2. When FlashDecoding is used for normal attention, it is not clear
# whether cascade attention is beneficial, because FlashDecoding can
# launch more CTAs than cascade attention.
# We use a simple performance model to compare the two methods.
# NOTE(woosuk): The performance model is very rough and may not be
# accurate.
num_tokens = num_reqs
# NOTE(woosuk): These are default tile sizes. flash-attn might use
# different tile sizes (e.g., 64 or 256) depending on the configuration.
q_tile_size = 128
kv_tile_size = 128
num_prefix_tiles = cdiv(common_prefix_len, kv_tile_size)
cascade_ctas = num_query_heads * cdiv(num_tokens, q_tile_size)
cascade_waves = cdiv(cascade_ctas, num_sms)
cascade_time = cascade_waves * num_prefix_tiles
flash_decoding_ctas = (
num_reqs * num_kv_heads * cdiv(num_queries_per_kv, q_tile_size)
)
flash_decoding_ctas *= num_prefix_tiles
flash_decoding_time = cdiv(flash_decoding_ctas, num_sms)
# Use cascade attention if it is faster than FlashDecoding.
return cascade_time < flash_decoding_time
def cascade_attention(
output: torch.Tensor,
query: torch.Tensor,
key_cache: torch.Tensor,
value_cache: torch.Tensor,
cu_query_lens: torch.Tensor,
max_query_len: int,
cu_prefix_query_lens: torch.Tensor,
prefix_kv_lens: torch.Tensor,
suffix_kv_lens: torch.Tensor,
max_kv_len: int,
softmax_scale: float,
alibi_slopes: torch.Tensor | None,
sliding_window: tuple[int, int],
logits_soft_cap: float,
block_table: torch.Tensor,
common_prefix_len: int,
max_num_splits: int,
fa_version: int,
prefix_scheduler_metadata: torch.Tensor | None = None,
suffix_scheduler_metadata: torch.Tensor | None = None,
q_descale: torch.Tensor | None = None,
k_descale: torch.Tensor | None = None,
v_descale: torch.Tensor | None = None,
s_aux: torch.Tensor | None = None,
) -> torch.Tensor:
assert alibi_slopes is None, "Cascade attention does not support ALiBi."
# TODO: Support sliding window.
assert sliding_window == (-1, -1), (
"Cascade attention does not support sliding window."
)
num_tokens = query.shape[0]
block_size = key_cache.shape[-3]
assert common_prefix_len % block_size == 0
num_common_kv_blocks = common_prefix_len // block_size
assert num_common_kv_blocks > 0
descale_shape = (cu_prefix_query_lens.shape[0] - 1, key_cache.shape[-2])
# Process shared prefix.
prefix_output, prefix_lse = flash_attn_varlen_func(
q=query,
k=key_cache,
v=value_cache,
cu_seqlens_q=cu_prefix_query_lens,
seqused_k=prefix_kv_lens,
max_seqlen_q=num_tokens,
max_seqlen_k=common_prefix_len,
softmax_scale=softmax_scale,
causal=False,
window_size=list(sliding_window),
block_table=block_table[:1],
softcap=logits_soft_cap,
return_softmax_lse=True,
scheduler_metadata=prefix_scheduler_metadata,
fa_version=fa_version,
q_descale=q_descale.expand(descale_shape) if q_descale is not None else None,
k_descale=k_descale.expand(descale_shape) if k_descale is not None else None,
v_descale=v_descale.expand(descale_shape) if v_descale is not None else None,
# s_aux is incorporated into prefix_lse inside the GPU kernel,
# enabling its effect during the final attention merge.
s_aux=s_aux,
num_splits=1 if envs.VLLM_BATCH_INVARIANT else max_num_splits,
)
descale_shape = (cu_query_lens.shape[0] - 1, key_cache.shape[-2])
# Process suffix per query.
suffix_output, suffix_lse = flash_attn_varlen_func(
q=query,
k=key_cache,
v=value_cache,
cu_seqlens_q=cu_query_lens,
seqused_k=suffix_kv_lens,
max_seqlen_q=max_query_len,
max_seqlen_k=max_kv_len - common_prefix_len,
softmax_scale=softmax_scale,
causal=True,
window_size=list(sliding_window),
block_table=block_table[:, num_common_kv_blocks:],
softcap=logits_soft_cap,
return_softmax_lse=True,
scheduler_metadata=suffix_scheduler_metadata,
fa_version=fa_version,
q_descale=q_descale.expand(descale_shape) if q_descale is not None else None,
k_descale=k_descale.expand(descale_shape) if k_descale is not None else None,
v_descale=v_descale.expand(descale_shape) if v_descale is not None else None,
num_splits=1 if envs.VLLM_BATCH_INVARIANT else max_num_splits,
)
# Merge prefix and suffix outputs, and store the result in output.
merge_attn_states(output, prefix_output, prefix_lse, suffix_output, suffix_lse)
_FA4_DENSE_ATTENTION_KERNEL = FA4DenseAttentionKernel()
vllm/v1/attention/backends/flash_attn_diffkv.py
+49 -82
View File
@@ -26,7 +26,6 @@ from .flash_attn import (
FlashAttentionBackend,
FlashAttentionImpl,
FlashAttentionMetadata,
cascade_attention,
)
logger = init_logger(__name__)
@@ -215,86 +214,54 @@ class FlashAttentionDiffKVImpl(FlashAttentionImpl):
key_cache = key_cache.view(current_platform.fp8_dtype())
value_cache = value_cache.view(current_platform.fp8_dtype())
if not attn_metadata.use_cascade:
cu_seqlens_q = attn_metadata.query_start_loc
seqused_k = attn_metadata.seq_lens
max_seqlen_q = attn_metadata.max_query_len
max_seqlen_k = attn_metadata.max_seq_len
block_table = attn_metadata.block_table
scheduler_metadata = attn_metadata.scheduler_metadata
cu_seqlens_q = attn_metadata.query_start_loc
seqused_k = attn_metadata.seq_lens
max_seqlen_q = attn_metadata.max_query_len
max_seqlen_k = attn_metadata.max_seq_len
block_table = attn_metadata.block_table
scheduler_metadata = attn_metadata.scheduler_metadata
descale_shape = (cu_seqlens_q.shape[0] - 1, self.num_kv_heads)
descale_shape = (cu_seqlens_q.shape[0] - 1, self.num_kv_heads)
if self.dcp_world_size > 1:
self._forward_with_dcp(
query[:num_actual_tokens],
key[:num_actual_tokens],
value[:num_actual_tokens],
key_cache,
value_cache,
output[:num_actual_tokens],
attn_metadata,
q_descale=layer._q_scale.expand(descale_shape),
k_descale=layer._k_scale.expand(descale_shape),
v_descale=layer._v_scale.expand(descale_shape),
)
return output
else:
sliding_window_size = (
list(self.sliding_window)
if self.sliding_window is not None
else None
)
flash_attn_varlen_func(
q=query[:num_actual_tokens],
k=key_cache,
v=value_cache,
out=output[:num_actual_tokens],
cu_seqlens_q=cu_seqlens_q,
max_seqlen_q=max_seqlen_q,
seqused_k=seqused_k,
max_seqlen_k=max_seqlen_k,
softmax_scale=self.scale,
causal=attn_metadata.causal,
alibi_slopes=self.alibi_slopes,
window_size=sliding_window_size,
block_table=block_table,
softcap=self.logits_soft_cap,
scheduler_metadata=scheduler_metadata,
fa_version=self.vllm_flash_attn_version,
q_descale=layer._q_scale.expand(descale_shape),
k_descale=layer._k_scale.expand(descale_shape),
v_descale=layer._v_scale.expand(descale_shape),
num_splits=attn_metadata.max_num_splits,
s_aux=self.sinks,
)
return output
# Cascade attention (rare case).
cascade_attention(
output[:num_actual_tokens],
query[:num_actual_tokens],
key_cache,
value_cache,
cu_query_lens=attn_metadata.query_start_loc,
max_query_len=attn_metadata.max_query_len,
cu_prefix_query_lens=attn_metadata.cu_prefix_query_lens,
prefix_kv_lens=attn_metadata.prefix_kv_lens,
suffix_kv_lens=attn_metadata.suffix_kv_lens,
max_kv_len=attn_metadata.max_seq_len,
softmax_scale=self.scale,
alibi_slopes=self.alibi_slopes,
sliding_window=self.sliding_window,
logits_soft_cap=self.logits_soft_cap,
block_table=attn_metadata.block_table,
common_prefix_len=attn_metadata.common_prefix_len,
max_num_splits=attn_metadata.max_num_splits,
fa_version=self.vllm_flash_attn_version,
prefix_scheduler_metadata=attn_metadata.prefix_scheduler_metadata,
suffix_scheduler_metadata=attn_metadata.scheduler_metadata,
q_descale=layer._q_scale,
k_descale=layer._k_scale,
v_descale=layer._v_scale,
s_aux=self.sinks,
)
return output
if self.dcp_world_size > 1:
self._forward_with_dcp(
query[:num_actual_tokens],
key[:num_actual_tokens],
value[:num_actual_tokens],
key_cache,
value_cache,
output[:num_actual_tokens],
attn_metadata,
q_descale=layer._q_scale.expand(descale_shape),
k_descale=layer._k_scale.expand(descale_shape),
v_descale=layer._v_scale.expand(descale_shape),
)
return output
else:
sliding_window_size = (
list(self.sliding_window) if self.sliding_window is not None else None
)
flash_attn_varlen_func(
q=query[:num_actual_tokens],
k=key_cache,
v=value_cache,
out=output[:num_actual_tokens],
cu_seqlens_q=cu_seqlens_q,
max_seqlen_q=max_seqlen_q,
seqused_k=seqused_k,
max_seqlen_k=max_seqlen_k,
softmax_scale=self.scale,
causal=attn_metadata.causal,
alibi_slopes=self.alibi_slopes,
window_size=sliding_window_size,
block_table=block_table,
softcap=self.logits_soft_cap,
scheduler_metadata=scheduler_metadata,
fa_version=self.vllm_flash_attn_version,
q_descale=layer._q_scale.expand(descale_shape),
k_descale=layer._k_scale.expand(descale_shape),
v_descale=layer._v_scale.expand(descale_shape),
num_splits=attn_metadata.max_num_splits,
s_aux=self.sinks,
)
return output
vllm/v1/attention/backends/flashinfer.py
+7 -121
View File
@@ -14,7 +14,6 @@ from flashinfer import (
BatchDecodeWithPagedKVCacheWrapper,
BatchPrefillWithPagedKVCacheWrapper,
BatchPrefillWithRaggedKVCacheWrapper,
MultiLevelCascadeAttentionWrapper,
get_seq_lens,
nvfp4_quantize_append_paged_kv_cache_with_slot_mapping,
)
@@ -732,16 +731,6 @@ class FlashInferMetadata:
Will be `None` if `num_decode_tokens == 0`.
"""
# --- Special Case: Cascade Attention ---
use_cascade: bool
"""
If True, the entire batch is a cascade attention call, and the
`prefill` and `decode` fields will both be None.
"""
cascade_wrapper: MultiLevelCascadeAttentionWrapper | None
class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
kv_cache_spec: AttentionSpec
@@ -951,8 +940,6 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
supports_dcp_with_varlen=False,
)
self._cascade_wrapper = None # Wrapper for cascade attention
# Global hyperparameters shared by all attention layers
# TODO: discard this for trtllm-gen backend
per_layer_parameters = get_per_layer_parameters(
@@ -1360,15 +1347,6 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
return decode_wrapper
def _get_cascade_wrapper(self):
if self._cascade_wrapper is None:
self._cascade_wrapper = MultiLevelCascadeAttentionWrapper(
2,
self._get_workspace_buffer(),
get_flashinfer_layout_string(self.kv_cache_layout),
)
return self._cascade_wrapper
def _compute_flashinfer_kv_metadata(
self,
num_blocks_np: np.ndarray,
@@ -1445,10 +1423,8 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
qo_indptr_cpu = common_attn_metadata.query_start_loc_cpu
# Step 1: Decide which dispatch modes to use:
# - Cascade attention (distinct mode)
# - Prefill (FI native or TRTLLM)
# - Decode (FI native, XQA, or trtllm-gen)
use_cascade = common_prefix_len > 0
uses_spec_reorder = self.reorder_batch_threshold > 1
# Page sizes >= 128 must use trtllm-gen; force it for prefill too.
prefill_force_trtllm = (
@@ -1486,12 +1462,11 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
)
if not all_uses_trtllm:
sinks_would_be_dropped = self.use_dcp or use_cascade
if self.has_sinks and sinks_would_be_dropped:
if self.has_sinks and self.use_dcp:
raise NotImplementedError(
"FlashInfer BatchDCPPrefillWrapper and "
"MultiLevelCascadeAttentionWrapper do not support attention sinks. "
"Please use TRTLLM on Blackwell or FlashAttention on earlier GPUs."
"FlashInfer BatchDCPPrefillWrapper does not support attention "
"sinks. Please use TRTLLM on Blackwell or FlashAttention on "
"earlier GPUs."
)
if not self.global_hyperparameters.has_same_window_lefts:
@@ -1508,7 +1483,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
)
# Step 2: Initialize the output metadata
# Leave prefill/decode/cascade_wrapper empty, to be populated
# Leave prefill/decode empty, to be populated
# case by case depending on the batch contents and backend selection.
attn_metadata = FlashInferMetadata(
kv_cache_layout=self.kv_cache_layout,
@@ -1521,10 +1496,8 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
num_prefills=num_prefills,
num_prefill_tokens=num_prefill_tokens,
causal=causal,
use_cascade=use_cascade,
prefill=None,
decode=None,
cascade_wrapper=None,
)
# Guard access to seq_lens_cpu, which may not always be needed
@@ -1532,7 +1505,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
# When all attention (both prefill and decode) uses TRTLLM,
# seq_lens_cpu is not needed since TRTLLM paths use GPU tensors
# (block_tables, seq_lens) directly.
needs_seq_lens_cpu = self.use_dcp or use_cascade or not all_uses_trtllm
needs_seq_lens_cpu = self.use_dcp or not all_uses_trtllm
if needs_seq_lens_cpu:
if (
self._num_speculative_tokens == 0
@@ -1578,13 +1551,6 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
seq_lens_np = None
num_blocks_np = None
# Adjust num_block_np for cascade attention
if use_cascade:
assert num_blocks_np is not None
assert common_prefix_len % page_size == 0
num_common_kv_blocks = common_prefix_len // page_size
num_blocks_np -= num_common_kv_blocks
# Compute paged_kv_indices if necessary
# paged_kv_indices is only needed for FlashInfer native paths;
# XQA/trtllm-gen paths use block_tables directly on GPU.
@@ -1592,9 +1558,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
needs_native_paged_decode = (
num_decodes > 0 and not decode_with_flashinfer_trtllm_api
)
needs_paged_kv_indices = (
use_cascade or needs_native_paged_prefill or needs_native_paged_decode
)
needs_paged_kv_indices = needs_native_paged_prefill or needs_native_paged_decode
if needs_paged_kv_indices:
assert num_blocks_np is not None
assert seq_lens_np is not None
@@ -1608,57 +1572,6 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
else:
paged_kv_indices = None
# Early-out for cascade attention
if use_cascade:
assert num_blocks_np is not None
# Grab the blocks of the shared prefix from the first request.
num_common_kv_blocks = common_prefix_len // page_size
# Create CPU versions directly for cascade (no GPU versions needed)
shared_qo_indptr_cpu = torch.tensor(
[0, num_actual_tokens], dtype=torch.int32, device="cpu"
)
shared_kv_page_indptr_cpu = torch.tensor(
[0, num_common_kv_blocks], dtype=torch.int32, device="cpu"
)
shared_kv_page_indices_cpu = block_table_tensor[0, :num_common_kv_blocks]
shared_kv_last_page_len_cpu = torch.tensor(
[page_size], dtype=torch.int32, device="cpu"
)
# Remove the blocks of the shared prefix from all requests.
block_table_tensor = block_table_tensor[:, num_common_kv_blocks:]
num_blocks_np -= num_common_kv_blocks
assert paged_kv_indices is not None
paged_kv_indptr_cpu = self.paged_kv_indptr.cpu[: 1 + num_reqs]
paged_kv_last_page_len_cpu = self.paged_kv_last_page_len.cpu[:num_reqs]
attn_metadata.cascade_wrapper = self._get_cascade_wrapper()
# Cascade attention must use the same q dtype for prefill and decode
# because it does not support FP8 kv-cache or FP8 query yet.
assert self.q_data_type_prefill == self.q_data_type_decode
attn_metadata.cascade_wrapper.plan(
qo_indptr_arr=[shared_qo_indptr_cpu, qo_indptr_cpu],
paged_kv_indptr_arr=[shared_kv_page_indptr_cpu, paged_kv_indptr_cpu],
paged_kv_indices_arr=[shared_kv_page_indices_cpu, paged_kv_indices],
paged_kv_last_page_len=[
shared_kv_last_page_len_cpu,
paged_kv_last_page_len_cpu,
],
num_qo_heads=self.num_qo_heads,
num_kv_heads=self.num_kv_heads,
head_dim=self.head_dim,
page_size=self.page_size,
causal=True,
sm_scale=self.sm_scale,
window_left=self.window_left,
logits_soft_cap=self.logits_soft_cap,
q_data_type=self.q_data_type_prefill,
kv_data_type=self.kv_cache_dtype,
)
return attn_metadata
# Step 3: Handle prefill and decode pathways case by case
## PREFILL PATHWAY
if num_prefills > 0:
@@ -1925,15 +1838,6 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
def update_draft_decode_metadata(self, _metadata: FlashInferMetadata) -> None:
pass
def use_cascade_attention(self, *args, **kwargs) -> bool:
if self.kv_cache_spec.dtype != self.vllm_config.model_config.dtype:
# TODO: The cascade wrapper currently does not support setting
# kv cache dtype to something different from query dtype.
return False
# TODO: Cascade attention doesn't work, disable it for now
# return use_cascade_attention(*args, **kwargs)
return False
class FlashInferImpl(AttentionImpl):
can_return_lse_for_decode: bool = True
@@ -2223,24 +2127,6 @@ class FlashInferImpl(AttentionImpl):
output_padded = output
output = output[:num_actual_tokens]
if attn_metadata.use_cascade:
# Cascade attention (rare case).
assert attn_metadata.cascade_wrapper is not None
stride_order = kv_cache_layout.layer_view_order
if self.is_kvcache_nvfp4:
kv_cache_views = tuple(
cache.permute(*stride_order)
for cache in kv_cache.split(self.num_kv_heads, dim=1)
)
else:
kv_perm = kv_cache.permute(*stride_order)
kv_cache_views = kv_perm.split(self.head_size, dim=-1)
kv_tuple = tuple(
canonicalize_singleton_dim_strides(cache) for cache in kv_cache_views
)
output.copy_(attn_metadata.cascade_wrapper.run(query, kv_tuple))
return output
# When using spec decoding, num_decodes can be < num_decode_tokens
# because some decode requests may have more than one query token.
num_decode_tokens = attn_metadata.num_decode_tokens
vllm/v1/attention/backends/flex_attention.py
-33
View File
@@ -145,10 +145,6 @@ class FlexAttentionBackend(AttentionBackend):
def get_builder_cls() -> type["FlexAttentionMetadataBuilder"]:
return FlexAttentionMetadataBuilder
@staticmethod
def use_cascade_attention(*args, **kwargs) -> bool:
return False
@classmethod
def get_supported_head_sizes(cls) -> list[int]:
return []
@@ -375,12 +371,6 @@ class FlexAttentionMetadata:
block_table: torch.Tensor
slot_mapping: torch.Tensor
use_cascade: bool
common_prefix_len: int
cu_prefix_query_lens: torch.Tensor | None
prefix_kv_lens: torch.Tensor | None
suffix_kv_lens: torch.Tensor | None
# Block info
total_cache_tokens: int
block_size: int
@@ -831,11 +821,6 @@ class FlexAttentionMetadata:
)
def __post_init__(self):
assert self.use_cascade is False, "Not implemented yet."
assert self.common_prefix_len == 0, "Not implemented yet."
assert self.cu_prefix_query_lens is None, "Not implemented yet."
assert self.prefix_kv_lens is None, "Not implemented yet."
assert self.suffix_kv_lens is None, "Not implemented yet."
# Create a lookup mapping from query indices -> request number
self.doc_ids = _offsets_to_doc_ids_tensor(
self.query_start_loc_cpu, self.query_start_loc.device
@@ -1029,16 +1014,6 @@ class FlexAttentionMetadataBuilder(AttentionMetadataBuilder[FlexAttentionMetadat
slot_mapping = common_attn_metadata.slot_mapping
num_blocks_per_seq = cdiv(seq_lens, self.block_size)
use_cascade = common_prefix_len > 0
cu_prefix_query_lens = None
prefix_kv_lens = None
suffix_kv_lens = None
if use_cascade:
raise NotImplementedError(
"Cascade prefix attention is not yet implemented "
"for FlexAttention backend"
)
block_size = self.kv_cache_spec.block_size
max_possible_seq_len = self.model_config.max_model_len
num_gpu_blocks = self.cache_config.num_gpu_blocks
@@ -1128,11 +1103,6 @@ class FlexAttentionMetadataBuilder(AttentionMetadataBuilder[FlexAttentionMetadat
seq_lens=self.persistent_seq_lens,
block_table=block_table_tensor,
slot_mapping=slot_mapping,
use_cascade=use_cascade,
common_prefix_len=common_prefix_len,
cu_prefix_query_lens=cu_prefix_query_lens,
prefix_kv_lens=prefix_kv_lens,
suffix_kv_lens=suffix_kv_lens,
block_size=block_size,
max_possible_sequence_length=max_possible_seq_len,
num_reqs=num_reqs,
@@ -1166,9 +1136,6 @@ class FlexAttentionMetadataBuilder(AttentionMetadataBuilder[FlexAttentionMetadat
return out
def use_cascade_attention(self, *args, **kwargs) -> bool:
return False
class FlexAttentionImpl(AttentionImpl):
sliding_window: int | None
vllm/v1/attention/backends/rocm_aiter_fa.py
+258 -281
View File
@@ -403,9 +403,6 @@ class AiterFlashAttentionMetadata:
prefill_metadata: AiterFlashAttentionPrefillMetadata | None
extend_metadata: AiterFlashAttentionChunkPrefillMetadata | None
# For cascade attention.
use_cascade: bool
# Only for fp8 shuffle layout kv cache, we allocate kv_scale for each layer
# since we might integrate per token quant for kv cache in the future.
k_scale: dict[str, torch.Tensor] | None
@@ -696,8 +693,6 @@ class AiterFlashAttentionMetadataBuilder(
chunk_context_metadata=chunk_context_metadata,
)
use_cascade = common_prefix_len > 0
attn_metadata = AiterFlashAttentionMetadata(
num_actual_tokens=common_attn_metadata.num_actual_tokens,
query_start_loc=common_attn_metadata.query_start_loc,
@@ -714,7 +709,6 @@ class AiterFlashAttentionMetadataBuilder(
decode_metadata=decode_metadata,
prefill_metadata=prefill_metadata,
extend_metadata=extend_metadata,
use_cascade=use_cascade,
k_scale=self.scale,
v_scale=self.scale,
kv_sharing_metadata=kv_sharing_metadata,
@@ -755,14 +749,10 @@ class AiterFlashAttentionMetadataBuilder(
decode_metadata=decode_metadata,
prefill_metadata=None,
extend_metadata=None,
use_cascade=False,
k_scale=self.scale,
v_scale=self.scale,
)
def use_cascade_attention(self, *args, **kwargs) -> bool:
return False
class AiterFlashAttentionBackend(AttentionBackend):
supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]
@@ -1219,186 +1209,136 @@ class AiterFlashAttentionImpl(AttentionImpl):
else "NHD",
total_tokens=num_actual_tokens - num_decode_tokens,
)
if not attn_metadata.use_cascade:
# calculate for pure prefills
if num_prefills > 0:
assert attn_metadata.prefill_metadata is not None
assert key is not None and value is not None
# calculate for pure prefills
if num_prefills > 0:
assert attn_metadata.prefill_metadata is not None
assert key is not None and value is not None
prefill_query = query[num_decode_tokens + num_extend_tokens :]
prefill_key = key[num_decode_tokens + num_extend_tokens :]
prefill_value = value[num_decode_tokens + num_extend_tokens :]
prefill_query = query[num_decode_tokens + num_extend_tokens :]
prefill_key = key[num_decode_tokens + num_extend_tokens :]
prefill_value = value[num_decode_tokens + num_extend_tokens :]
rocm_aiter_ops.flash_attn_varlen_func(
q=prefill_query,
k=prefill_key,
v=prefill_value,
cu_seqlens_q=attn_metadata.prefill_metadata.query_start_loc,
cu_seqlens_k=attn_metadata.prefill_metadata.query_start_loc,
max_seqlen_q=attn_metadata.prefill_metadata.max_query_len,
max_seqlen_k=attn_metadata.prefill_metadata.max_seq_len,
min_seqlen_q=1,
dropout_p=0.0,
softmax_scale=self.scale,
causal=attn_metadata.causal,
window_size=self.sliding_window,
alibi_slopes=self.alibi_slopes,
out=output_actual_tokens[num_decode_tokens + num_extend_tokens :],
sink_ptr=self.sinks,
rocm_aiter_ops.flash_attn_varlen_func(
q=prefill_query,
k=prefill_key,
v=prefill_value,
cu_seqlens_q=attn_metadata.prefill_metadata.query_start_loc,
cu_seqlens_k=attn_metadata.prefill_metadata.query_start_loc,
max_seqlen_q=attn_metadata.prefill_metadata.max_query_len,
max_seqlen_k=attn_metadata.prefill_metadata.max_seq_len,
min_seqlen_q=1,
dropout_p=0.0,
softmax_scale=self.scale,
causal=attn_metadata.causal,
window_size=self.sliding_window,
alibi_slopes=self.alibi_slopes,
out=output_actual_tokens[num_decode_tokens + num_extend_tokens :],
sink_ptr=self.sinks,
)
# calculate for extends
if num_extends > 0:
assert attn_metadata.extend_metadata is not None
assert key is not None and value is not None
extend_tokens_slice = slice(
num_decode_tokens, num_decode_tokens + num_extend_tokens
)
extend_queries = query[extend_tokens_slice]
extend_keys = key[extend_tokens_slice]
extend_values = value[extend_tokens_slice]
extend_outputs = output[extend_tokens_slice]
k_scale = layer._k_scale
v_scale = layer._v_scale
if rocm_aiter_ops.is_shuffle_kv_cache_enabled():
k_scale = attn_metadata.k_scale
v_scale = attn_metadata.v_scale
self.extend_forward(
attn_metadata=attn_metadata,
query=extend_queries,
key=extend_keys,
value=extend_values,
key_cache=key_cache,
value_cache=value_cache,
output=extend_outputs,
cu_seqlens_q=attn_metadata.extend_metadata.query_start_loc,
max_seqlen_q=attn_metadata.extend_metadata.max_query_len,
min_seqlen_q=1,
max_seqlen_k=attn_metadata.extend_metadata.max_seq_len,
block_table=attn_metadata.block_table[
num_decodes : num_decodes + num_extends
],
slot_mapping=attn_metadata.slot_mapping[
num_decodes : num_decodes + num_extends
],
k_scale=k_scale,
v_scale=v_scale,
)
# calculate for decodes
if num_decodes > 0:
assert attn_metadata.decode_metadata is not None
decode_max_query_len = attn_metadata.decode_metadata.max_query_len
# Use unified_attention for the decodes the paged kernels can't
# take: sliding window, sinks, or a multi-token batch
# (pa_fwd_asm and paged_attention_v1 don't support sinks).
if (
self.sliding_window[0] != -1
or decode_max_query_len > 1
or self.sinks is not None
):
k_descale, v_descale = self._get_kv_cache_descales(layer, num_decodes)
assert not rocm_aiter_ops.is_shuffle_kv_cache_enabled(), (
"Shuffle KV cache layout is not supported with sliding "
"window, sinks, or speculative decoding (multi-token decode)."
)
# calculate for extends
if num_extends > 0:
assert attn_metadata.extend_metadata is not None
assert key is not None and value is not None
extend_tokens_slice = slice(
num_decode_tokens, num_decode_tokens + num_extend_tokens
)
extend_queries = query[extend_tokens_slice]
extend_keys = key[extend_tokens_slice]
extend_values = value[extend_tokens_slice]
extend_outputs = output[extend_tokens_slice]
k_scale = layer._k_scale
v_scale = layer._v_scale
if rocm_aiter_ops.is_shuffle_kv_cache_enabled():
k_scale = attn_metadata.k_scale
v_scale = attn_metadata.v_scale
self.extend_forward(
attn_metadata=attn_metadata,
query=extend_queries,
key=extend_keys,
value=extend_values,
key_cache=key_cache,
value_cache=value_cache,
output=extend_outputs,
cu_seqlens_q=attn_metadata.extend_metadata.query_start_loc,
max_seqlen_q=attn_metadata.extend_metadata.max_query_len,
min_seqlen_q=1,
max_seqlen_k=attn_metadata.extend_metadata.max_seq_len,
block_table=attn_metadata.block_table[
num_decodes : num_decodes + num_extends
],
slot_mapping=attn_metadata.slot_mapping[
num_decodes : num_decodes + num_extends
],
k_scale=k_scale,
v_scale=v_scale,
)
# calculate for decodes
if num_decodes > 0:
assert attn_metadata.decode_metadata is not None
decode_max_query_len = attn_metadata.decode_metadata.max_query_len
# Use unified_attention for the decodes the paged kernels can't
# take: sliding window, sinks, or a multi-token batch
# (pa_fwd_asm and paged_attention_v1 don't support sinks).
if (
self.sliding_window[0] != -1
or decode_max_query_len > 1
or self.sinks is not None
):
k_descale, v_descale = self._get_kv_cache_descales(
layer, num_decodes
if not attn_metadata.causal:
from aiter.ops.triton.attention.mha_v3 import (
flash_attn_with_kvcache,
)
assert not rocm_aiter_ops.is_shuffle_kv_cache_enabled(), (
"Shuffle KV cache layout is not supported with sliding "
"window, sinks, or speculative decoding (multi-token decode)."
)
if not attn_metadata.causal:
from aiter.ops.triton.attention.mha_v3 import (
flash_attn_with_kvcache,
)
decode_query = query[:num_decode_tokens].reshape(
num_decodes,
decode_max_query_len,
decode_query = query[:num_decode_tokens].reshape(
num_decodes,
decode_max_query_len,
query.shape[1],
query.shape[2],
)
decode_out = flash_attn_with_kvcache(
q=decode_query,
k_cache=key_cache,
v_cache=value_cache,
cache_seqlens=attn_metadata.seq_lens[:num_decodes],
softmax_scale=self.scale,
causal=attn_metadata.causal,
window_size=self.sliding_window,
softcap=self.logits_soft_cap,
q_descale=None,
k_descale=k_descale,
v_descale=v_descale,
page_table=attn_metadata.block_table[:num_decodes],
)
output[:num_decode_tokens].copy_(
decode_out.reshape(
num_decode_tokens,
query.shape[1],
query.shape[2],
)
decode_out = flash_attn_with_kvcache(
q=decode_query,
k_cache=key_cache,
v_cache=value_cache,
cache_seqlens=attn_metadata.seq_lens[:num_decodes],
softmax_scale=self.scale,
causal=attn_metadata.causal,
window_size=self.sliding_window,
softcap=self.logits_soft_cap,
q_descale=None,
k_descale=k_descale,
v_descale=v_descale,
page_table=attn_metadata.block_table[:num_decodes],
)
output[:num_decode_tokens].copy_(
decode_out.reshape(
num_decode_tokens,
query.shape[1],
query.shape[2],
)
)
else:
# Non-uniform query lengths can appear in real serving
# traffic (e.g. mixed datasets). Fall back to varlen
# unified_attention instead of asserting.
from aiter.ops.triton.unified_attention import (
unified_attention,
)
unified_attention(
q=query[:num_decode_tokens],
k=key_cache,
v=value_cache,
out=output[:num_decode_tokens],
cu_seqlens_q=attn_metadata.query_start_loc[
: num_decodes + 1
],
max_seqlen_q=decode_max_query_len,
seqused_k=attn_metadata.seq_lens[:num_decodes],
max_seqlen_k=attn_metadata.max_seq_len,
softmax_scale=self.scale,
causal=True,
alibi_slopes=self.alibi_slopes,
window_size=self.sliding_window,
block_table=attn_metadata.block_table[:num_decodes],
softcap=self.logits_soft_cap,
q_descale=None,
k_descale=k_descale,
v_descale=v_descale,
sinks=self.sinks,
)
return
# The ll4mi kernel in paged_attention_v1 requires
# HEAD_SIZE >= 16 * NWARPS (= 64 on ROCm with NWARPS=4).
# For smaller head sizes or sliding window attention,
# fall back to the unified_attention triton kernel which
# handles both correctly.
_MIN_HEAD_SIZE_FOR_LL4MI = 64
use_unified_attention = self.head_size < _MIN_HEAD_SIZE_FOR_LL4MI
if use_unified_attention:
k_descale, v_descale = self._get_kv_cache_descales(
layer, num_decodes
)
assert not rocm_aiter_ops.is_shuffle_kv_cache_enabled(), (
"unified_attention fallback with shuffle layout "
"is not supported yet."
)
else:
# Non-uniform query lengths can appear in real serving
# traffic (e.g. mixed datasets). Fall back to varlen
# unified_attention instead of asserting.
from aiter.ops.triton.unified_attention import (
unified_attention,
)
decode_cu_seqlens_q = attn_metadata.query_start_loc[
: num_decodes + 1
]
unified_attention(
q=query[:num_decode_tokens],
k=key_cache,
v=value_cache,
out=output[:num_decode_tokens],
cu_seqlens_q=decode_cu_seqlens_q,
max_seqlen_q=1,
cu_seqlens_q=attn_metadata.query_start_loc[: num_decodes + 1],
max_seqlen_q=decode_max_query_len,
seqused_k=attn_metadata.seq_lens[:num_decodes],
max_seqlen_k=attn_metadata.max_seq_len,
softmax_scale=self.scale,
@@ -1410,120 +1350,157 @@ class AiterFlashAttentionImpl(AttentionImpl):
q_descale=None,
k_descale=k_descale,
v_descale=v_descale,
sinks=self.sinks,
)
elif rocm_aiter_ops.is_shuffle_kv_cache_enabled():
_, num_heads, head_size = query.shape
num_blocks, block_size, num_kv_heads, _ = key_cache.shape
x = 16 // key_cache.element_size()
new_key_cache = key_cache.reshape(
num_blocks, num_kv_heads, head_size // x, block_size, x
)
new_value_cache = value_cache.reshape(
num_blocks, num_kv_heads, block_size // x, head_size, x
)
return
# This kernel derives block addresses arithmetically, so it
# only reads the right bytes when each side is a dense
# plane of pages.
assert new_key_cache.stride(0) == (
num_kv_heads * head_size * block_size
), (
"paged_attention_common needs the K/V sides as dense "
"planes; the resolved KV cache layout interleaves them "
"within a block."
)
num_seqs = attn_metadata.seq_lens.shape[0]
max_num_partitions = (
attn_metadata.max_seq_len + _PARTITION_SIZE_ROCM - 1
) // _PARTITION_SIZE_ROCM
tmp_out = torch.empty(
(num_seqs, num_heads, max_num_partitions, head_size),
dtype=query.dtype,
device=query.device,
)
exp_sums = torch.empty(
(num_seqs, num_heads, max_num_partitions),
dtype=torch.float32,
device=query.device,
)
max_logits = torch.empty_like(exp_sums)
k_qscale = (
layer._k_scale
if attn_metadata.k_scale is None
else attn_metadata.k_scale
)
v_qscale = (
layer._v_scale
if attn_metadata.v_scale is None
else attn_metadata.v_scale
)
rocm_aiter_ops.paged_attention_common(
Q=query[:num_decode_tokens],
K=new_key_cache,
V=new_value_cache,
tmp_out=tmp_out,
max_logits=max_logits,
exp_sums=exp_sums,
max_seq_len=attn_metadata.max_seq_len,
block_tables=attn_metadata.block_table[:num_decodes],
context_lens=attn_metadata.seq_lens[:num_decodes],
block_tables_stride0=attn_metadata.block_table[
:num_decodes
].stride(0),
scale=self.scale,
K_QScale_hip=k_qscale,
V_QScale_hip=v_qscale,
K_QScale_asm=k_qscale,
V_QScale_asm=v_qscale,
out_=output[:num_decode_tokens],
kv_cache_dtype=self.kv_cache_dtype,
)
else:
_, num_heads, head_size = query.shape
nbytes_per_qo_elem = torch.finfo(query.dtype).bits // 8
num_seqs = attn_metadata.seq_lens.shape[0]
max_num_partitions = (
attn_metadata.max_seq_len + _PARTITION_SIZE_ROCM - 1
) // _PARTITION_SIZE_ROCM
# The ll4mi kernel in paged_attention_v1 requires
# HEAD_SIZE >= 16 * NWARPS (= 64 on ROCm with NWARPS=4).
# For smaller head sizes or sliding window attention,
# fall back to the unified_attention triton kernel which
# handles both correctly.
_MIN_HEAD_SIZE_FOR_LL4MI = 64
use_unified_attention = self.head_size < _MIN_HEAD_SIZE_FOR_LL4MI
workspace_buffer = torch.empty(
(num_seqs * num_heads * max_num_partitions * head_size)
* nbytes_per_qo_elem
+ 2 * (num_seqs * num_heads * max_num_partitions) * 4,
dtype=torch.uint8,
device=output.device,
)
if use_unified_attention:
k_descale, v_descale = self._get_kv_cache_descales(layer, num_decodes)
assert not rocm_aiter_ops.is_shuffle_kv_cache_enabled(), (
"unified_attention fallback with shuffle layout "
"is not supported yet."
)
from aiter.ops.triton.unified_attention import (
unified_attention,
)
# import so that aiter register the op to the namespace of
# torch.ops.aiter
import aiter # noqa: F401
decode_cu_seqlens_q = attn_metadata.query_start_loc[: num_decodes + 1]
unified_attention(
q=query[:num_decode_tokens],
k=key_cache,
v=value_cache,
out=output[:num_decode_tokens],
cu_seqlens_q=decode_cu_seqlens_q,
max_seqlen_q=1,
seqused_k=attn_metadata.seq_lens[:num_decodes],
max_seqlen_k=attn_metadata.max_seq_len,
softmax_scale=self.scale,
causal=True,
alibi_slopes=self.alibi_slopes,
window_size=self.sliding_window,
block_table=attn_metadata.block_table[:num_decodes],
softcap=self.logits_soft_cap,
q_descale=None,
k_descale=k_descale,
v_descale=v_descale,
)
elif rocm_aiter_ops.is_shuffle_kv_cache_enabled():
_, num_heads, head_size = query.shape
num_blocks, block_size, num_kv_heads, _ = key_cache.shape
x = 16 // key_cache.element_size()
new_key_cache = key_cache.reshape(
num_blocks, num_kv_heads, head_size // x, block_size, x
)
new_value_cache = value_cache.reshape(
num_blocks, num_kv_heads, block_size // x, head_size, x
)
torch.ops.aiter.paged_attention_v1(
output[:num_decode_tokens],
workspace_buffer,
query[:num_decode_tokens],
key_cache,
value_cache,
self.scale,
attn_metadata.block_table[:num_decodes],
attn_metadata.query_start_loc[:num_decodes],
attn_metadata.seq_lens[:num_decodes],
attn_metadata.max_seq_len,
self.alibi_slopes,
self.kv_cache_dtype,
"NHD",
self.logits_soft_cap,
layer._k_scale,
layer._v_scale,
None,
_PARTITION_SIZE_ROCM,
1,
self.sliding_window[0] + 1,
)
else:
raise NotImplementedError(
"Cascade attention is not implemented for ROCM AITER"
)
# This kernel derives block addresses arithmetically, so it
# only reads the right bytes when each side is a dense
# plane of pages.
assert new_key_cache.stride(0) == (
num_kv_heads * head_size * block_size
), (
"paged_attention_common needs the K/V sides as dense "
"planes; the resolved KV cache layout interleaves them "
"within a block."
)
num_seqs = attn_metadata.seq_lens.shape[0]
max_num_partitions = (
attn_metadata.max_seq_len + _PARTITION_SIZE_ROCM - 1
) // _PARTITION_SIZE_ROCM
tmp_out = torch.empty(
(num_seqs, num_heads, max_num_partitions, head_size),
dtype=query.dtype,
device=query.device,
)
exp_sums = torch.empty(
(num_seqs, num_heads, max_num_partitions),
dtype=torch.float32,
device=query.device,
)
max_logits = torch.empty_like(exp_sums)
k_qscale = (
layer._k_scale
if attn_metadata.k_scale is None
else attn_metadata.k_scale
)
v_qscale = (
layer._v_scale
if attn_metadata.v_scale is None
else attn_metadata.v_scale
)
rocm_aiter_ops.paged_attention_common(
Q=query[:num_decode_tokens],
K=new_key_cache,
V=new_value_cache,
tmp_out=tmp_out,
max_logits=max_logits,
exp_sums=exp_sums,
max_seq_len=attn_metadata.max_seq_len,
block_tables=attn_metadata.block_table[:num_decodes],
context_lens=attn_metadata.seq_lens[:num_decodes],
block_tables_stride0=attn_metadata.block_table[:num_decodes].stride(
0
),
scale=self.scale,
K_QScale_hip=k_qscale,
V_QScale_hip=v_qscale,
K_QScale_asm=k_qscale,
V_QScale_asm=v_qscale,
out_=output[:num_decode_tokens],
kv_cache_dtype=self.kv_cache_dtype,
)
else:
_, num_heads, head_size = query.shape
nbytes_per_qo_elem = torch.finfo(query.dtype).bits // 8
num_seqs = attn_metadata.seq_lens.shape[0]
max_num_partitions = (
attn_metadata.max_seq_len + _PARTITION_SIZE_ROCM - 1
) // _PARTITION_SIZE_ROCM
workspace_buffer = torch.empty(
(num_seqs * num_heads * max_num_partitions * head_size)
* nbytes_per_qo_elem
+ 2 * (num_seqs * num_heads * max_num_partitions) * 4,
dtype=torch.uint8,
device=output.device,
)
# import so that aiter register the op to the namespace of
# torch.ops.aiter
import aiter # noqa: F401
torch.ops.aiter.paged_attention_v1(
output[:num_decode_tokens],
workspace_buffer,
query[:num_decode_tokens],
key_cache,
value_cache,
self.scale,
attn_metadata.block_table[:num_decodes],
attn_metadata.query_start_loc[:num_decodes],
attn_metadata.seq_lens[:num_decodes],
attn_metadata.max_seq_len,
self.alibi_slopes,
self.kv_cache_dtype,
"NHD",
self.logits_soft_cap,
layer._k_scale,
layer._v_scale,
None,
_PARTITION_SIZE_ROCM,
1,
self.sliding_window[0] + 1,
)
return output
vllm/v1/attention/backends/rocm_aiter_unified_attn.py
-6
View File
@@ -129,10 +129,6 @@ class RocmAiterUnifiedAttentionBackend(RocmAttentionBackend):
return (KVCacheLayout.LBHNC,)
return (KVCacheLayout.LBHNC, KVCacheLayout.LHBNC)
@staticmethod
def use_cascade_attention(*args, **kwargs) -> bool:
return False
@staticmethod
def get_builder_cls() -> type["RocmAiterUnifiedAttentionMetadataBuilder"]:
return RocmAiterUnifiedAttentionMetadataBuilder
@@ -237,8 +233,6 @@ class RocmAiterUnifiedAttentionImpl(RocmAttentionImpl):
# Profiling run.
return output.fill_(0)
assert attn_metadata.use_cascade is False
# IMPORTANT!
# NOTE(woosuk): With piece-wise CUDA graphs, this method is executed in
# eager-mode PyTorch. Thus, we need to be careful about any CPU overhead
vllm/v1/attention/backends/rocm_attn.py
-37
View File
@@ -58,16 +58,8 @@ class RocmAttentionMetadata:
block_table: torch.Tensor
slot_mapping: torch.Tensor
# For cascade attention.
use_cascade: bool
common_prefix_len: int
cu_prefix_query_lens: torch.Tensor | None
prefix_kv_lens: torch.Tensor | None
suffix_kv_lens: torch.Tensor | None
# Optional aot scheduling
scheduler_metadata: torch.Tensor | None = None
prefix_scheduler_metadata: torch.Tensor | None = None
# DFlash drafting sets this to False via CommonAttentionMetadata.
causal: bool = True
@@ -124,23 +116,6 @@ class RocmAttentionMetadataBuilder(AttentionMetadataBuilder[RocmAttentionMetadat
block_table_tensor = common_attn_metadata.block_table_tensor
slot_mapping = common_attn_metadata.slot_mapping
use_cascade = common_prefix_len > 0
if use_cascade:
cu_prefix_query_lens = torch.tensor(
[0, num_actual_tokens], dtype=torch.int32, device=self.device
)
prefix_kv_lens = torch.tensor(
[common_prefix_len], dtype=torch.int32, device=self.device
)
suffix_kv_lens = common_attn_metadata.seq_lens.cpu() - common_prefix_len
suffix_kv_lens = suffix_kv_lens.to(self.device)
else:
cu_prefix_query_lens = None
prefix_kv_lens = None
suffix_kv_lens = None
prefix_scheduler_metadata = None
attn_metadata = RocmAttentionMetadata(
num_actual_tokens=num_actual_tokens,
max_query_len=max_query_len,
@@ -149,12 +124,6 @@ class RocmAttentionMetadataBuilder(AttentionMetadataBuilder[RocmAttentionMetadat
seq_lens=seq_lens,
block_table=block_table_tensor,
slot_mapping=slot_mapping,
use_cascade=use_cascade,
common_prefix_len=common_prefix_len,
cu_prefix_query_lens=cu_prefix_query_lens,
prefix_kv_lens=prefix_kv_lens,
suffix_kv_lens=suffix_kv_lens,
prefix_scheduler_metadata=prefix_scheduler_metadata,
causal=common_attn_metadata.causal,
)
return attn_metadata
@@ -267,10 +236,6 @@ class RocmAttentionBackend(AttentionBackend):
# supported, just without the native kernels.
return (KVCacheLayout.LHBNC, KVCacheLayout.LBHNC)
@staticmethod
def use_cascade_attention(*args, **kwargs) -> bool:
return False
@staticmethod
def get_builder_cls() -> type["RocmAttentionMetadataBuilder"]:
return RocmAttentionMetadataBuilder
@@ -418,8 +383,6 @@ class RocmAttentionImpl(AttentionImpl):
# Profiling run.
return output.fill_(0)
assert attn_metadata.use_cascade is False
# IMPORTANT!
# NOTE(woosuk): With piece-wise CUDA graphs, this method is executed in
# eager-mode PyTorch. Thus, we need to be careful about any CPU overhead
vllm/v1/attention/backends/rocm_segmented_attn.py
-1
View File
@@ -271,7 +271,6 @@ class RocmSegmentedAttentionImpl(RocmAttentionImpl):
workspace = self._get_workspace(query.device)
if attn_metadata is None:
return output.fill_(0)
assert not attn_metadata.use_cascade
num_actual_tokens = attn_metadata.num_actual_tokens
key_cache, value_cache = self._split_kv_cache(kv_cache)
vllm/v1/attention/backends/triton_attn.py
-37
View File
@@ -80,16 +80,8 @@ class TritonAttentionMetadata:
causal: bool | torch.Tensor
# For cascade attention.
use_cascade: bool
common_prefix_len: int
cu_prefix_query_lens: torch.Tensor | None
prefix_kv_lens: torch.Tensor | None
suffix_kv_lens: torch.Tensor | None
# Optional aot scheduling
scheduler_metadata: torch.Tensor | None = None
prefix_scheduler_metadata: torch.Tensor | None = None
mm_prefix_range: dict[int, list[tuple[int, int]]] | None = None
mm_prefix_range_tensor: torch.Tensor | None = None
rswa_prefix_lens: torch.Tensor | None = None
@@ -230,23 +222,6 @@ class TritonAttentionMetadataBuilder(AttentionMetadataBuilder[TritonAttentionMet
block_table_tensor = common_attn_metadata.block_table_tensor
slot_mapping = common_attn_metadata.slot_mapping
use_cascade = common_prefix_len > 0
if use_cascade:
cu_prefix_query_lens = torch.tensor(
[0, num_actual_tokens], dtype=torch.int32, device=self.device
)
prefix_kv_lens = torch.tensor(
[common_prefix_len], dtype=torch.int32, device=self.device
)
suffix_kv_lens = common_attn_metadata.seq_lens.cpu() - common_prefix_len
suffix_kv_lens = suffix_kv_lens.to(self.device)
else:
cu_prefix_query_lens = None
prefix_kv_lens = None
suffix_kv_lens = None
prefix_scheduler_metadata = None
attn_metadata = TritonAttentionMetadata(
num_actual_tokens=num_actual_tokens,
max_query_len=max_query_len,
@@ -256,12 +231,6 @@ class TritonAttentionMetadataBuilder(AttentionMetadataBuilder[TritonAttentionMet
block_table=block_table_tensor,
slot_mapping=slot_mapping,
causal=common_attn_metadata.causal,
use_cascade=use_cascade,
common_prefix_len=common_prefix_len,
cu_prefix_query_lens=cu_prefix_query_lens,
prefix_kv_lens=prefix_kv_lens,
suffix_kv_lens=suffix_kv_lens,
prefix_scheduler_metadata=prefix_scheduler_metadata,
seq_threshold_3D=self.seq_threshold_3D,
num_par_softmax_segments=self.num_par_softmax_segments,
softmax_segm_output=self.softmax_segm_output,
@@ -378,10 +347,6 @@ class TritonAttentionBackend(AttentionBackend):
def get_impl_cls() -> type["TritonAttentionImpl"]:
return TritonAttentionImpl
@staticmethod
def use_cascade_attention(*args, **kwargs) -> bool:
return False
@staticmethod
def get_builder_cls() -> type["TritonAttentionMetadataBuilder"]:
return TritonAttentionMetadataBuilder
@@ -622,8 +587,6 @@ class TritonAttentionImpl(AttentionImpl):
# Profiling run.
return output.fill_(0)
assert attn_metadata.use_cascade is False
# IMPORTANT!
# NOTE(woosuk): With piece-wise CUDA graphs, this method is executed in
# eager-mode PyTorch. Thus, we need to be careful about any CPU overhead
vllm/v1/attention/backends/triton_attn_diffkv.py
-4
View File
@@ -181,10 +181,6 @@ class TritonAttentionDiffKVImpl(TritonAttentionImpl):
if attn_metadata is None:
return output.fill_(0)
assert attn_metadata.use_cascade is False, (
"Cascade attention not supported for TritonAttentionDiffKVImpl"
)
num_actual_tokens = attn_metadata.num_actual_tokens
head_size_qk = self.head_size
head_size_v = TritonAttentionDiffKVBackend.head_size_v
vllm/v1/core/kv_cache_coordinator.py
-20
View File
@@ -426,23 +426,6 @@ class KVCacheCoordinator(ABC):
blocks.extend(manager.pop_blocks_for_free(request_id))
return blocks
def get_num_common_prefix_blocks(self, running_request_id: str) -> list[int]:
"""Get the number of common prefix blocks for all requests with allocated
KV cache for each kv cache group.
Args:
running_request_id: The request ID of any running request, used to
identify the common prefix blocks.
Returns:
list[int]: The number of common prefix blocks for each kv cache group.
"""
return [
manager.get_num_common_prefix_blocks(running_request_id)
for manager in self.single_type_managers
]
def remove_skipped_blocks(
self,
request_id: str,
@@ -527,9 +510,6 @@ class KVCacheCoordinatorNoPrefixCache(KVCacheCoordinator):
)
self.num_single_type_manager = len(self.single_type_managers)
def get_num_common_prefix_blocks(self, running_request_id: str) -> list[int]:
return [0] * self.num_single_type_manager
def find_longest_cache_hit(
self,
block_hashes: list[BlockHash],
vllm/v1/core/kv_cache_manager.py
+3 -33
View File
@@ -689,39 +689,9 @@ class KVCacheManager:
return True
def get_num_common_prefix_blocks(self, running_request_id: str) -> list[int]:
"""Calculate the number of common prefix blocks for each kv cache group.
The function selects a running request and iterates through its blocks.
A block is considered a common prefix block if ALL requests with
allocated KV cache share it (i.e., ref_cnt equals the number of entries
in req_to_blocks).
NOTE(woosuk): The number of requests with allocated KV cache is **greater
than or equal to** the number of requests scheduled in the current step.
This is because having allocated KV cache only indicates that:
1. The request has not yet finished, and
2. The request holds its blocks unfreed.
While all scheduled requests must have allocated KV cache, the inverse
is not necessarily true. There may be requests with allocated KV cache
that are not scheduled in the current step.
This can result in an edge case where the number of common prefix blocks
is 0, even though all scheduled requests share a common prefix. This
occurs because there may be unscheduled requests that do not share the
common prefix. Currently, this case cannot be easily detected, so the
function returns 0 in such cases.
Args:
running_request_id: The request ID of any running request, used to
identify the common prefix blocks.
Returns:
list[int]: The number of common prefix blocks for each kv cache
group.
"""
return self.coordinator.get_num_common_prefix_blocks(running_request_id)
"""Deprecated: cascade attention has been removed. Always returns 0
for each KV cache group. Kept for out-of-tree schedulers."""
return [0] * self.num_kv_cache_groups
def take_events(self) -> list[KVCacheEvent]:
"""Take the KV cache events from the block pool.
vllm/v1/core/sched/output.py
+5 -5
View File
@@ -2,7 +2,7 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from collections.abc import Callable
from dataclasses import dataclass
from dataclasses import dataclass, field
from functools import cached_property
from typing import TYPE_CHECKING
@@ -253,9 +253,6 @@ class SchedulerOutput:
# E.g., if a request has [0, 1], it could mean the vision encoder needs
# to process that the request's 0-th and 1-th images in the current step.
scheduled_encoder_inputs: dict[str, list[int]]
# Number of common prefix blocks for all requests in each KV cache group.
# This can be used for cascade attention.
num_common_prefix_blocks: list[int]
# Request IDs that are finished in between the previous and the current
# steps. This is used to notify the workers about the finished requests
@@ -311,6 +308,10 @@ class SchedulerOutput:
# Number of spec tokens to schedule for the next step.
num_spec_tokens_to_schedule: int = 0
# Deprecated: cascade attention has been removed, so this is always empty.
# Kept for out-of-tree schedulers and workers.
num_common_prefix_blocks: list[int] = field(default_factory=list)
@classmethod
def make_empty(cls) -> "SchedulerOutput":
return cls(
@@ -320,7 +321,6 @@ class SchedulerOutput:
total_num_scheduled_tokens=0,
scheduled_spec_decode_tokens={},
scheduled_encoder_inputs={},
num_common_prefix_blocks=[],
finished_req_ids=set(),
free_encoder_mm_hashes=[],
)
vllm/v1/core/sched/scheduler.py
-11
View File
@@ -1414,16 +1414,6 @@ class Scheduler(SchedulerInterface):
scheduled_running_reqs
) <= len(self.running)
# Get the longest common prefix among all requests in the running queue.
# This can be potentially used for cascade attention.
num_common_prefix_blocks = [0] * len(self.kv_cache_config.kv_cache_groups)
with record_function_or_nullcontext("schedule: get_num_common_prefix_blocks"):
if self.running:
any_request_id = self.running[0].request_id
num_common_prefix_blocks = (
self.kv_cache_manager.get_num_common_prefix_blocks(any_request_id)
)
# Construct the scheduler output.
if self.use_v2_model_runner:
scheduled_new_reqs.extend(scheduled_resumed_reqs)
@@ -1516,7 +1506,6 @@ class Scheduler(SchedulerInterface):
scheduled_spec_decode_tokens=scheduled_spec_decode_tokens,
scheduled_encoder_inputs=scheduled_encoder_inputs,
scheduled_encoder_input_stats=scheduled_encoder_input_stats,
num_common_prefix_blocks=num_common_prefix_blocks,
preempted_req_ids=self.reset_preempted_req_ids,
# finished_req_ids is an existing state in the scheduler,
# instead of being newly scheduled in this step.
vllm/v1/core/single_type_kv_cache_manager.py
-52
View File
@@ -588,21 +588,6 @@ class SingleTypeKVCacheManager(ABC):
# Free blocks in reverse order so that the tail blocks are freed first.
self.block_pool.free_blocks(reversed(self.pop_blocks_for_free(request_id)))
@abstractmethod
def get_num_common_prefix_blocks(self, running_request_id: str) -> int:
"""Get the number of common prefix blocks for all requests with allocated
KV cache.
Args:
running_request_id: The request ID.
Returns:
The number of common prefix blocks for all requests with allocated
KV cache.
"""
raise NotImplementedError
@classmethod
@abstractmethod
def find_longest_cache_hit(
@@ -898,16 +883,6 @@ class FullAttentionManager(SingleTypeKVCacheManager):
block_size=self.block_size,
)
def get_num_common_prefix_blocks(self, running_request_id: str) -> int:
blocks = self.req_to_blocks[running_request_id]
num_common_blocks = 0
for block in blocks:
if block.ref_cnt == len(self.req_to_blocks):
num_common_blocks += 1
else:
break
return num_common_blocks
class RSWAManager(FullAttentionManager):
"""KV cache manager for Reference Sliding Window Attention (R-SWA).
@@ -1194,14 +1169,6 @@ class SlidingWindowManager(SingleTypeKVCacheManager):
num_computed_tokens - self.sliding_window + 1 - self.extra_retained_tokens,
)
def get_num_common_prefix_blocks(self, running_request_id: str) -> int:
"""NOTE(Chen): The prefix blocks are null blocks for sliding window layers.
So it's not correct to count ref_cnt like FullAttentionManager. Return
0 here for correctness. Need to support cascade attention + sliding
window in the future.
"""
return 0
class CircularBufferManager(FullAttentionManager):
"""Claims the ring's single block per request; prefix caching disabled."""
@@ -1289,9 +1256,6 @@ class CircularBufferManager(FullAttentionManager):
) -> None:
return
def get_num_common_prefix_blocks(self, running_request_id: str) -> int:
return 0
def get_num_skipped_tokens(self, num_computed_tokens: int) -> int:
return 0
@@ -1447,10 +1411,6 @@ class ChunkedLocalAttentionManager(SingleTypeKVCacheManager):
) * self.attention_chunk_size
return num_skipped_tokens
def get_num_common_prefix_blocks(self, running_request_id: str) -> int:
"""Cascade attention is not supported by chunked local attention."""
return 0
class MambaManager(SingleTypeKVCacheManager):
supports_fine_grained_hash_lookup: ClassVar[bool] = True
@@ -1680,10 +1640,6 @@ class MambaManager(SingleTypeKVCacheManager):
self.block_pool.free_blocks([blocks[last_state_block_idx]])
blocks[last_state_block_idx] = self._null_block
def get_num_common_prefix_blocks(self, running_request_id: str) -> int:
"""Cascade attention is not supported by mamba"""
return 0
def _needs_internal_checkpoint(
self,
request_id: str,
@@ -2192,11 +2148,6 @@ class CrossAttentionManager(SingleTypeKVCacheManager):
# requests, so this method is not relevant.
raise ValueError("Should not be called as prefix caching is disabled.")
def get_num_common_prefix_blocks(self, running_request_id: str) -> int:
# Cross-attention blocks contain request-specific encoder states
# and are not shared between different requests
return 0
@classmethod
def find_longest_cache_hit(
cls,
@@ -2394,9 +2345,6 @@ class _HiSparseAuxiliaryManager(SingleTypeKVCacheManager):
) -> None:
return None
def get_num_common_prefix_blocks(self, running_request_id: str) -> int:
return 0
@classmethod
def find_longest_cache_hit(
cls,
vllm/v1/cudagraph_dispatcher.py
+2 -2
View File
@@ -254,8 +254,8 @@ class CudagraphDispatcher:
all modes are allowed.
invalid_modes: Set of cudagraph modes to exclude. Subtracted from
valid_modes to compute allowed modes. (e.g., {FULL} for
features like cascade attention not supported by full
cudagraphs). None means no modes are excluded.
encoder-decoder batches with encoder outputs). None means no
modes are excluded.
"""
allowed_modes = valid_modes or CUDAGraphMode.valid_runtime_modes()
vllm/v1/worker/cpu_model_runner.py
-4
View File
@@ -32,7 +32,6 @@ class CPUModelRunner(GPUModelRunner):
# Note: speculative decoding is now supported on CPU with C++ native impls
self.use_cuda_graph = False
self.cascade_attn_enabled = False
self._postprocess_tensors()
@@ -109,9 +108,6 @@ class CPUModelRunner(GPUModelRunner):
elif self.speculative_config.uses_draft_model():
logger.info("Draft model KV cache initialized for CPU backend")
def _init_device_properties(self) -> None:
pass
def _sync_device(self) -> None:
pass
vllm/v1/worker/gpu/warmup.py
-6
View File
@@ -106,7 +106,6 @@ def run_mixed_prefill_decode_warmup(
prefill_token_ids = list(range(prefill_len))
kv_cache_groups = model_runner.kv_cache_config.kv_cache_groups
num_kv_cache_groups = len(kv_cache_groups)
block_count = _warmup_block_counter(model_runner)
kv_cache_specs = [g.kv_cache_spec for g in kv_cache_groups]
decode_prefill_block_counts = [
@@ -166,7 +165,6 @@ def run_mixed_prefill_decode_warmup(
decode_req_id: decode_prompt_len,
}
decode_prefill_output.total_num_scheduled_tokens = decode_prompt_len
decode_prefill_output.num_common_prefix_blocks = [0] * num_kv_cache_groups
decode_new_blocks = tuple(_alloc_blocks(n) for n in decode_block_deltas)
cached_decode_req = CachedRequestData.make_empty()
@@ -198,7 +196,6 @@ def run_mixed_prefill_decode_warmup(
prefill_req_id: prefill_len,
}
mixed_output.total_num_scheduled_tokens = num_tokens
mixed_output.num_common_prefix_blocks = [0] * num_kv_cache_groups
cleanup_output = SchedulerOutput.make_empty()
cleanup_output.finished_req_ids = {decode_req_id, prefill_req_id}
@@ -272,7 +269,6 @@ def _warmup_kernels(
decode_len = prompt_len + num_decode_steps * decode_query_len
kv_cache_groups = model_runner.kv_cache_config.kv_cache_groups
num_kv_cache_groups = len(kv_cache_groups)
# Encoder-decoder models: give each warmup request a dummy encoder input so
# cross-attention warms up over a realistic, non-empty key sequence.
@@ -361,7 +357,6 @@ def _warmup_kernels(
prefill_output.scheduled_new_reqs = new_reqs
prefill_output.num_scheduled_tokens = {rid: prompt_len for rid in req_ids}
prefill_output.total_num_scheduled_tokens = prompt_len * num_reqs
prefill_output.num_common_prefix_blocks = [0] * num_kv_cache_groups
# Disable KV connector for warmup run.
model_runner.kv_connector.set_disabled(True)
@@ -424,7 +419,6 @@ def _warmup_kernels(
decode_output.total_num_scheduled_tokens = sum(
step_num_scheduled_tokens.values()
)
decode_output.num_common_prefix_blocks = [0] * num_kv_cache_groups
worker_execute_model(decode_output)
worker_sample_tokens(None)
vllm/v1/worker/gpu_model_runner.py
+7 -177
View File
@@ -115,7 +115,6 @@ from vllm.utils.gpu_sync_debug import gpu_sync_allowed
from vllm.utils.math_utils import cdiv
from vllm.utils.mem_utils import DeviceMemoryProfiler, format_gib
from vllm.utils.nvtx_pytorch_hooks import PytHooks
from vllm.utils.platform_utils import num_compute_units
from vllm.utils.torch_utils import (
PIN_MEMORY,
async_tensor_h2d,
@@ -145,15 +144,12 @@ from vllm.v1.core.sched.output import NewRequestData
from vllm.v1.cudagraph_dispatcher import CudagraphDispatcher
from vllm.v1.kv_cache_interface import (
AttentionSpec,
ChunkedLocalAttentionSpec,
CrossAttentionSpec,
EncoderOnlyAttentionSpec,
FullAttentionSpec,
KVCacheConfig,
KVCacheGroupSpec,
KVCacheSpec,
KVCacheSpecKind,
SlidingWindowSpec,
UniformTypeKVCacheSpecs,
get_kv_cache_spec_kind,
)
@@ -530,13 +526,11 @@ class GPUModelRunner(
)
# Model-related.
self.num_query_heads = model_config.get_num_attention_heads(parallel_config)
self.inputs_embeds_size = model_config.get_inputs_embeds_size()
# Only relevant for models using ALiBi (e.g, MPT)
self.use_alibi = model_config.uses_alibi
self.cascade_attn_enabled = not self.model_config.disable_cascade_attn
self.is_mm_prefix_lm = self.model_config.is_mm_prefix_lm
# Deprecated: cascade attention has been removed. Kept for
# out-of-tree model runners.
self.cascade_attn_enabled = False
# Multi-modal data support
self.mm_registry = MULTIMODAL_REGISTRY
@@ -769,9 +763,6 @@ class GPUModelRunner(
else:
self.cudagraph_batch_sizes = []
# Cache the device properties.
self._init_device_properties()
# Encoder timing registry for observability
self.encoder_timing_registry: dict[str, EncoderTimingStats] = {}
self._encoder_timing_lock = threading.Lock()
@@ -1151,11 +1142,6 @@ class GPUModelRunner(
if hasattr(self, "_kv_block_zeroer"):
self._kv_block_zeroer.zero_block_ids(block_ids)
# Note: used for model runner override.
def _init_device_properties(self) -> None:
"""Initialize attributes from torch.cuda.get_device_properties"""
self.num_sms = num_compute_units(self.device.index)
# Note: used for model runner override.
def _sync_device(self) -> None:
torch.accelerator.synchronize()
@@ -2273,7 +2259,6 @@ class GPUModelRunner(
use_spec_decode: bool = False,
for_cudagraph_capture: bool = False,
num_scheduled_tokens: dict[str, int] | None = None,
cascade_attn_prefix_lens: list[list[int]] | None = None,
slot_mappings: dict[int, torch.Tensor] | None = None,
) -> tuple[PerLayerAttnMetadata, CommonAttentionMetadata | None]:
"""Returns:
@@ -2472,12 +2457,6 @@ class GPUModelRunner(
kv_cache_spec = kv_cache_spec.kv_cache_specs[attn_group.layer_names[0]]
cache_key = (kv_cache_spec, type(builder))
cascade_attn_prefix_len = (
cascade_attn_prefix_lens[kv_cache_gid][attn_gid]
if cascade_attn_prefix_lens
else 0
)
extra_attn_metadata_args = {}
if use_spec_decode and isinstance(
builder,
@@ -2517,7 +2496,6 @@ class GPUModelRunner(
attn_metadata_i = attn_group.build_metadata(
common_attn_metadata,
ubid or 0,
common_prefix_len=cascade_attn_prefix_len,
**extra_attn_metadata_args,
)
if builder.supports_update_block_table:
@@ -2595,144 +2573,6 @@ class GPUModelRunner(
return attn_metadata, spec_decode_common_attn_metadata
def _compute_cascade_attn_prefix_lens(
self,
num_scheduled_tokens: np.ndarray,
num_computed_tokens: np.ndarray,
num_common_prefix_blocks: list[int],
) -> list[list[int]] | None:
"""Returns:
Optional[cascade_attn_prefix_lens]
cascade_attn_prefix_lens is 2D:
``[kv_cache_group_id][attn_group_idx]``,
None if we should not use cascade attention
"""
use_cascade_attn = False
num_kv_cache_groups = len(self.kv_cache_config.kv_cache_groups)
cascade_attn_prefix_lens: list[list[int]] = [
[] for _ in range(num_kv_cache_groups)
]
for kv_cache_gid in range(num_kv_cache_groups):
for attn_group in self.attn_groups[kv_cache_gid]:
if isinstance(attn_group.kv_cache_spec, EncoderOnlyAttentionSpec):
cascade_attn_prefix_len = 0
else:
# 0 if cascade attention should not be used
cascade_attn_prefix_len = self._compute_cascade_attn_prefix_len(
num_scheduled_tokens,
num_computed_tokens,
num_common_prefix_blocks[kv_cache_gid],
attn_group.kv_cache_spec,
attn_group.get_metadata_builder(),
)
cascade_attn_prefix_lens[kv_cache_gid].append(cascade_attn_prefix_len)
use_cascade_attn |= cascade_attn_prefix_len > 0
return cascade_attn_prefix_lens if use_cascade_attn else None
def _compute_cascade_attn_prefix_len(
self,
num_scheduled_tokens: np.ndarray,
num_computed_tokens: np.ndarray,
num_common_prefix_blocks: int,
kv_cache_spec: KVCacheSpec,
attn_metadata_builder: AttentionMetadataBuilder,
) -> int:
"""Compute the length of the common prefix for cascade attention.
NOTE(woosuk): The common prefix length returned by this function
represents the length used specifically for cascade attention, not the
actual number of tokens shared between requests. When cascade attention
is disabled (use_cascade=False), this function returns 0 even if
requests share common tokens. Additionally, the common prefix length is
truncated to a multiple of the block size and may be further truncated
due to implementation details explained below.
Args:
num_scheduled_tokens: Number of tokens scheduled per request.
num_computed_tokens: Number of tokens already computed per request.
num_common_prefix_blocks: Number of shared KV cache blocks.
kv_cache_spec: The KV cache spec of the group being considered.
attn_metadata_builder: Builder used to decide whether cascade
attention can be used.
Returns:
int: Length of common prefix in tokens.
"""
common_prefix_len = num_common_prefix_blocks * kv_cache_spec.block_size
if common_prefix_len == 0:
# Common case.
return 0
# NOTE(woosuk): Cascade attention uses two attention kernels: one
# for the common prefix and the other for the rest. For the first
# kernel, we concatenate all the query tokens (possibly from
# different requests) and treat them as if they are from the same
# request. Then, we use bi-directional attention to process the
# common prefix in the KV cache. Importantly, this means that the
# first kernel does not do any masking.
# Consider the following example:
# Request 1's input query: [D, E, X]
# Request 1's kv cache: [A, B, C, D, E, X]
# Request 1's num_computed_tokens: 3 (i.e., [A, B, C])
# Request 2's input query: [E, Y]
# Request 2's kv cache: [A, B, C, D, E, Y]
# Request 2's num_computed_tokens: 4 (i.e., [A, B, C, D])
# If we use [A, B, C, D, E] as the common prefix, then the
# first kernel will compute the bi-directional attention between
# input query [D, E, X, E, Y] and common prefix [A, B, C, D, E].
# However, this is wrong because D in Request 1 should not attend to
# E in the common prefix (i.e., we need masking).
# To avoid this, [A, B, C, D] should be the common prefix.
# That is, the common prefix should be capped by the minimum
# num_computed_tokens among the requests, and plus one to include
# the first token of the query.
# In practice, we use [A, B, C] as the common prefix, instead of
# [A, B, C, D] (i.e., the common prefix is capped by the minimum
# num_computed_tokens, without plus one).
# This is because of an implementation detail: We want to always
# use two kernels for cascade attention. Let's imagine:
# Request 3's input query: [D]
# Request 3's kv cache: [A, B, C, D]
# Request 3's num_computed_tokens: 3 (i.e., [A, B, C])
# If we use [A, B, C, D] as the common prefix for Request 1-3,
# then Request 3 will be processed only by the first kernel,
# and the second kernel will get an empty input. While this is not
# a fundamental problem, our current implementation does not support
# this case.
common_prefix_len = min(common_prefix_len, num_computed_tokens.min())
# common_prefix_len should be a multiple of the block size.
common_prefix_len = (
common_prefix_len // kv_cache_spec.block_size * kv_cache_spec.block_size
)
use_sliding_window = isinstance(kv_cache_spec, SlidingWindowSpec) or (
isinstance(kv_cache_spec, FullAttentionSpec)
and kv_cache_spec.sliding_window is not None
)
use_local_attention = isinstance(kv_cache_spec, ChunkedLocalAttentionSpec) or (
isinstance(kv_cache_spec, FullAttentionSpec)
and kv_cache_spec.attention_chunk_size is not None
)
assert isinstance(kv_cache_spec, AttentionSpec)
use_cascade = attn_metadata_builder.use_cascade_attention(
common_prefix_len=common_prefix_len,
query_lens=num_scheduled_tokens,
num_query_heads=self.num_query_heads,
num_kv_heads=kv_cache_spec.num_kv_heads,
use_alibi=self.use_alibi,
use_sliding_window=use_sliding_window,
use_local_attention=use_local_attention,
num_sms=self.num_sms,
dcp_world_size=self.dcp_world_size,
)
return common_prefix_len if use_cascade else 0
def _calc_mrope_positions(self, scheduler_output: "SchedulerOutput"):
mrope_pos_ptr = 0
for index, req_id in enumerate(self.input_batch.req_ids):
@@ -3881,7 +3721,8 @@ class GPUModelRunner(
num_reqs: int,
num_scheduled_tokens_np: np.ndarray,
max_num_scheduled_tokens: int,
use_cascade_attn: bool,
# Deprecated and ignored. Kept for out-of-tree model runners.
use_cascade_attn: bool = False,
allow_microbatching: bool = True,
force_eager: bool = False,
# For cudagraph capture TODO(lucas): Refactor how we capture cudagraphs (will
@@ -3929,7 +3770,7 @@ class GPUModelRunner(
)
cudagraph_mode, batch_descriptor = dispatch_cudagraph(
num_tokens, disable_full=use_cascade_attn or has_encoder_output
num_tokens, disable_full=has_encoder_output
)
num_tokens_padded = batch_descriptor.num_tokens
@@ -4181,16 +4022,6 @@ class GPUModelRunner(
self._prepare_inputs(scheduler_output, num_scheduled_tokens_np)
)
cascade_attn_prefix_lens = None
# Disable cascade attention when using microbatching (DBO)
if self.cascade_attn_enabled and not self.parallel_config.use_ubatching:
# Pre-compute cascade attention prefix lengths
cascade_attn_prefix_lens = self._compute_cascade_attn_prefix_lens(
num_scheduled_tokens_np,
self.input_batch.num_computed_tokens_cpu[:num_reqs],
scheduler_output.num_common_prefix_blocks,
)
(
cudagraph_mode,
batch_desc,
@@ -4202,7 +4033,7 @@ class GPUModelRunner(
num_reqs=num_reqs,
num_scheduled_tokens_np=num_scheduled_tokens_np,
max_num_scheduled_tokens=max_num_scheduled_tokens,
use_cascade_attn=cascade_attn_prefix_lens is not None,
use_cascade_attn=False,
num_encoder_reqs=len(scheduler_output.scheduled_encoder_inputs),
allow_microbatching=self._allow_microbatching(
num_reqs, num_scheduled_tokens_np
@@ -4319,7 +4150,6 @@ class GPUModelRunner(
max_num_sampled_tokens=max_num_sampled_tokens,
use_spec_decode=use_spec_decode,
num_scheduled_tokens=scheduler_output.num_scheduled_tokens,
cascade_attn_prefix_lens=cascade_attn_prefix_lens,
slot_mappings=slot_mappings_by_group,
)
)
vllm/v1/worker/xpu_model_runner.py
-2
View File
@@ -24,8 +24,6 @@ class XPUModelRunner(GPUModelRunner):
):
with _torch_cuda_wrapper():
super().__init__(vllm_config, device)
# FIXME: To be verified.
self.cascade_attn_enabled = False
class XPUModelRunnerV2(GPUModelRunnerV2):