diff --git a/docs/design/cuda_graphs.md b/docs/design/cuda_graphs.md index 9d53a98f65..9b8b97a53d 100644 --- a/docs/design/cuda_graphs.md +++ b/docs/design/cuda_graphs.md @@ -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""" diff --git a/tests/conftest.py b/tests/conftest.py index 952c328f91..c23cc9568f 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -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 diff --git a/tests/kernels/attention/test_cascade_flash_attn.py b/tests/kernels/attention/test_cascade_flash_attn.py deleted file mode 100755 index 36c0c51de4..0000000000 --- a/tests/kernels/attention/test_cascade_flash_attn.py +++ /dev/null @@ -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) diff --git a/tests/kernels/attention/test_prefix_prefill.py b/tests/kernels/attention/test_prefix_prefill.py index b2efa76da8..3f675e15e6 100644 --- a/tests/kernels/attention/test_prefix_prefill.py +++ b/tests/kernels/attention/test_prefix_prefill.py @@ -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), diff --git a/tests/platforms/test_cpu.py b/tests/platforms/test_cpu.py index 9c7f0b188b..956436297c 100644 --- a/tests/platforms/test_cpu.py +++ b/tests/platforms/test_cpu.py @@ -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, diff --git a/tests/system_messages/sonnet3.5_nov2024.txt b/tests/system_messages/sonnet3.5_nov2024.txt deleted file mode 100644 index 2dc285ac96..0000000000 --- a/tests/system_messages/sonnet3.5_nov2024.txt +++ /dev/null @@ -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. diff --git a/tests/v1/attention/test_attention_backends.py b/tests/v1/attention/test_attention_backends.py index add611d92b..dc17c20227 100644 --- a/tests/v1/attention/test_attention_backends.py +++ b/tests/v1/attention/test_attention_backends.py @@ -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 = {} diff --git a/tests/v1/attention/test_rocm_attention_backends_selection.py b/tests/v1/attention/test_rocm_attention_backends_selection.py index 067d1dbc8b..e0247a94a1 100644 --- a/tests/v1/attention/test_rocm_attention_backends_selection.py +++ b/tests/v1/attention/test_rocm_attention_backends_selection.py @@ -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), diff --git a/tests/v1/core/test_async_scheduler.py b/tests/v1/core/test_async_scheduler.py index 2995537aac..65f38cdedf 100644 --- a/tests/v1/core/test_async_scheduler.py +++ b/tests/v1/core/test_async_scheduler.py @@ -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=[], ) diff --git a/tests/v1/core/test_scheduler.py b/tests/v1/core/test_scheduler.py index be3ea553d9..f53ba2ad23 100644 --- a/tests/v1/core/test_scheduler.py +++ b/tests/v1/core/test_scheduler.py @@ -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=[], ) diff --git a/tests/v1/core/test_single_type_kv_cache_manager.py b/tests/v1/core/test_single_type_kv_cache_manager.py index f926f49651..cd4581dcef 100644 --- a/tests/v1/core/test_single_type_kv_cache_manager.py +++ b/tests/v1/core/test_single_type_kv_cache_manager.py @@ -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 diff --git a/tests/v1/cudagraph/test_cudagraph_dispatch.py b/tests/v1/cudagraph/test_cudagraph_dispatch.py index b0a4587ddf..c24c48c377 100644 --- a/tests/v1/cudagraph/test_cudagraph_dispatch.py +++ b/tests/v1/cudagraph/test_cudagraph_dispatch.py @@ -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, diff --git a/tests/v1/e2e/general/test_cascade_attention.py b/tests/v1/e2e/general/test_cascade_attention.py deleted file mode 100644 index 251746271d..0000000000 --- a/tests/v1/e2e/general/test_cascade_attention.py +++ /dev/null @@ -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: Implement fibonacci sequence in Python.\n:" - - 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 diff --git a/tests/v1/kv_connector/unit/test_kv_connector_lifecycle.py b/tests/v1/kv_connector/unit/test_kv_connector_lifecycle.py index a1df2a964b..b4a72ddd79 100644 --- a/tests/v1/kv_connector/unit/test_kv_connector_lifecycle.py +++ b/tests/v1/kv_connector/unit/test_kv_connector_lifecycle.py @@ -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(), diff --git a/tests/v1/kv_connector/unit/test_nixl_connector.py b/tests/v1/kv_connector/unit/test_nixl_connector.py index 607f5ff3cd..dcbfa7eb04 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector.py @@ -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=[], ) diff --git a/tests/v1/simple_kv_offload/test_scheduler.py b/tests/v1/simple_kv_offload/test_scheduler.py index 6c0cc310c7..9f818c866f 100644 --- a/tests/v1/simple_kv_offload/test_scheduler.py +++ b/tests/v1/simple_kv_offload/test_scheduler.py @@ -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=[], diff --git a/tests/v1/streaming_input/test_gpu_model_runner_v2_streaming.py b/tests/v1/streaming_input/test_gpu_model_runner_v2_streaming.py index c5f8090f94..fc316e3364 100644 --- a/tests/v1/streaming_input/test_gpu_model_runner_v2_streaming.py +++ b/tests/v1/streaming_input/test_gpu_model_runner_v2_streaming.py @@ -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=[], ) diff --git a/tests/v1/worker/test_gpu_model_runner.py b/tests/v1/worker/test_gpu_model_runner.py index 0578768d68..1247bf193a 100644 --- a/tests/v1/worker/test_gpu_model_runner.py +++ b/tests/v1/worker/test_gpu_model_runner.py @@ -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=[], ) diff --git a/tests/v1/worker/test_gpu_model_runner_prefill_tail_classification.py b/tests/v1/worker/test_gpu_model_runner_prefill_tail_classification.py index dd9046de61..14d5d8d4bb 100644 --- a/tests/v1/worker/test_gpu_model_runner_prefill_tail_classification.py +++ b/tests/v1/worker/test_gpu_model_runner_prefill_tail_classification.py @@ -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=[], ) diff --git a/tests/v1/worker/test_gpu_target_dependent_ar_speculator.py b/tests/v1/worker/test_gpu_target_dependent_ar_speculator.py index 490b84e778..42678aa9d1 100644 --- a/tests/v1/worker/test_gpu_target_dependent_ar_speculator.py +++ b/tests/v1/worker/test_gpu_target_dependent_ar_speculator.py @@ -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, diff --git a/tests/v1/worker/test_mamba_utils.py b/tests/v1/worker/test_mamba_utils.py index 7e588ccde2..61ad7e24dd 100644 --- a/tests/v1/worker/test_mamba_utils.py +++ b/tests/v1/worker/test_mamba_utils.py @@ -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(), diff --git a/vllm/config/model.py b/vllm/config/model.py index 545c306e1e..87759f4606 100644 --- a/vllm/config/model.py +++ b/vllm/config/model.py @@ -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 diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index b4298cd145..af4c99113f 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -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] diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index a892900816..ea628467e0 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -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, diff --git a/vllm/model_executor/models/extract_hidden_states.py b/vllm/model_executor/models/extract_hidden_states.py index 653e7691fb..89253ac5d8 100644 --- a/vllm/model_executor/models/extract_hidden_states.py +++ b/vllm/model_executor/models/extract_hidden_states.py @@ -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( diff --git a/vllm/platforms/cpu.py b/vllm/platforms/cpu.py index 43d472f22e..542a7a405d 100644 --- a/vllm/platforms/cpu.py +++ b/vllm/platforms/cpu.py @@ -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 diff --git a/vllm/v1/attention/backend.py b/vllm/v1/attention/backend.py index 68bb0b3bdd..466bba586d 100644 --- a/vllm/v1/attention/backend.py +++ b/vllm/v1/attention/backend.py @@ -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 diff --git a/vllm/v1/attention/backends/cpu_attn.py b/vllm/v1/attention/backends/cpu_attn.py index 46b01cc551..9062b48aea 100644 --- a/vllm/v1/attention/backends/cpu_attn.py +++ b/vllm/v1/attention/backends/cpu_attn.py @@ -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: diff --git a/vllm/v1/attention/backends/flash_attn.py b/vllm/v1/attention/backends/flash_attn.py index 4a336586ac..6e906bc520 100755 --- a/vllm/v1/attention/backends/flash_attn.py +++ b/vllm/v1/attention/backends/flash_attn.py @@ -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() diff --git a/vllm/v1/attention/backends/flash_attn_diffkv.py b/vllm/v1/attention/backends/flash_attn_diffkv.py index 4bb58ef4ea..7b38a3695e 100644 --- a/vllm/v1/attention/backends/flash_attn_diffkv.py +++ b/vllm/v1/attention/backends/flash_attn_diffkv.py @@ -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 diff --git a/vllm/v1/attention/backends/flashinfer.py b/vllm/v1/attention/backends/flashinfer.py index 1a41458447..f6c3e262d6 100755 --- a/vllm/v1/attention/backends/flashinfer.py +++ b/vllm/v1/attention/backends/flashinfer.py @@ -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 diff --git a/vllm/v1/attention/backends/flex_attention.py b/vllm/v1/attention/backends/flex_attention.py index 1b9a21b655..c2103082df 100644 --- a/vllm/v1/attention/backends/flex_attention.py +++ b/vllm/v1/attention/backends/flex_attention.py @@ -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 diff --git a/vllm/v1/attention/backends/rocm_aiter_fa.py b/vllm/v1/attention/backends/rocm_aiter_fa.py index 329aee5359..0ec0d5dd36 100644 --- a/vllm/v1/attention/backends/rocm_aiter_fa.py +++ b/vllm/v1/attention/backends/rocm_aiter_fa.py @@ -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 diff --git a/vllm/v1/attention/backends/rocm_aiter_unified_attn.py b/vllm/v1/attention/backends/rocm_aiter_unified_attn.py index 0d4fd55194..9f8e9e9ab1 100644 --- a/vllm/v1/attention/backends/rocm_aiter_unified_attn.py +++ b/vllm/v1/attention/backends/rocm_aiter_unified_attn.py @@ -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 diff --git a/vllm/v1/attention/backends/rocm_attn.py b/vllm/v1/attention/backends/rocm_attn.py index 3ccc1e5d57..471bb9f058 100644 --- a/vllm/v1/attention/backends/rocm_attn.py +++ b/vllm/v1/attention/backends/rocm_attn.py @@ -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 diff --git a/vllm/v1/attention/backends/rocm_segmented_attn.py b/vllm/v1/attention/backends/rocm_segmented_attn.py index c541eed451..4725db917b 100644 --- a/vllm/v1/attention/backends/rocm_segmented_attn.py +++ b/vllm/v1/attention/backends/rocm_segmented_attn.py @@ -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) diff --git a/vllm/v1/attention/backends/triton_attn.py b/vllm/v1/attention/backends/triton_attn.py index 095163f144..a95b0c1d0e 100644 --- a/vllm/v1/attention/backends/triton_attn.py +++ b/vllm/v1/attention/backends/triton_attn.py @@ -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 diff --git a/vllm/v1/attention/backends/triton_attn_diffkv.py b/vllm/v1/attention/backends/triton_attn_diffkv.py index 932f483af3..08e16c3e64 100644 --- a/vllm/v1/attention/backends/triton_attn_diffkv.py +++ b/vllm/v1/attention/backends/triton_attn_diffkv.py @@ -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 diff --git a/vllm/v1/core/kv_cache_coordinator.py b/vllm/v1/core/kv_cache_coordinator.py index b6defc52e8..7d4f2d5a2b 100644 --- a/vllm/v1/core/kv_cache_coordinator.py +++ b/vllm/v1/core/kv_cache_coordinator.py @@ -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], diff --git a/vllm/v1/core/kv_cache_manager.py b/vllm/v1/core/kv_cache_manager.py index 37ceb4e474..4f472170a8 100644 --- a/vllm/v1/core/kv_cache_manager.py +++ b/vllm/v1/core/kv_cache_manager.py @@ -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. diff --git a/vllm/v1/core/sched/output.py b/vllm/v1/core/sched/output.py index a5af5a3d57..9ad0d4e658 100644 --- a/vllm/v1/core/sched/output.py +++ b/vllm/v1/core/sched/output.py @@ -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=[], ) diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 2603fa10a7..bcb6def938 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -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. diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index 218d59f430..2e4fe55b96 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/vllm/v1/core/single_type_kv_cache_manager.py @@ -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, diff --git a/vllm/v1/cudagraph_dispatcher.py b/vllm/v1/cudagraph_dispatcher.py index 0c1fb0b2c4..d6abba301a 100644 --- a/vllm/v1/cudagraph_dispatcher.py +++ b/vllm/v1/cudagraph_dispatcher.py @@ -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() diff --git a/vllm/v1/worker/cpu_model_runner.py b/vllm/v1/worker/cpu_model_runner.py index 32fa01bd27..d8df0ca103 100644 --- a/vllm/v1/worker/cpu_model_runner.py +++ b/vllm/v1/worker/cpu_model_runner.py @@ -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 diff --git a/vllm/v1/worker/gpu/warmup.py b/vllm/v1/worker/gpu/warmup.py index 894e2e5b6b..b81ddf6d73 100644 --- a/vllm/v1/worker/gpu/warmup.py +++ b/vllm/v1/worker/gpu/warmup.py @@ -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) diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 669b314d4c..e19c1b84f4 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -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, ) ) diff --git a/vllm/v1/worker/xpu_model_runner.py b/vllm/v1/worker/xpu_model_runner.py index 15fa8e1e76..83ff08371e 100644 --- a/vllm/v1/worker/xpu_model_runner.py +++ b/vllm/v1/worker/xpu_model_runner.py @@ -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):