squall/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Bugfix] Fix illegal memory access error with chunked prefill, prefix caching, block manager v2 and xformers enabled together (#9532)

Browse Source 
Signed-off-by: sasha0552 <admin@sasha0552.org>
This commit is contained in:
sasha0552 2024-10-31 18:46:36 +00:00 committed by GitHub
parent 77f7ef2908
commit 55650c83a0
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 34 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

28
tests/prefix_caching/test_prefix_caching.py
View File

@ -5,6 +5,7 @@ Run `pytest tests/prefix_caching/test_prefix_caching.py`.
import pytest import pytest
from tests.kernels.utils import override_backend_env_variable from tests.kernels.utils import override_backend_env_variable
from vllm import SamplingParams, TokensPrompt
from ..models.utils import check_outputs_equal from ..models.utils import check_outputs_equal
@ -12,6 +13,14 @@ MODELS = [
"facebook/opt-125m", "facebook/opt-125m",
] ]
UNSTABLE_PROMPT_SEQUENCE = [
([0] * 588) + ([1] * 1332) + ([2] * 30) + ([3] * 1),
([0] * 588) + ([1] * 1332) + ([4] * 3) + ([5] * 50),
([0] * 588) + ([1] * 1332) + ([2] * 30) + ([6] * 95),
([0] * 588) + ([1] * 1332) + ([4] * 3) + ([7] * 174),
([0] * 588) + ([8] * 1539),
]
@pytest.mark.parametrize("model", MODELS) @pytest.mark.parametrize("model", MODELS)
@pytest.mark.parametrize("backend", ["FLASH_ATTN", "FLASHINFER", "XFORMERS"]) @pytest.mark.parametrize("backend", ["FLASH_ATTN", "FLASHINFER", "XFORMERS"])
@ -57,3 +66,22 @@ def test_mixed_requests(
name_0="hf", name_0="hf",
name_1="vllm", name_1="vllm",
) )
@pytest.mark.parametrize("backend", ["FLASH_ATTN", "FLASHINFER", "XFORMERS"])
def test_unstable_prompt_sequence(
vllm_runner,
backend: str,
monkeypatch,
) -> None:
override_backend_env_variable(monkeypatch, backend)
with vllm_runner(
"Qwen/Qwen2.5-0.5B-Instruct",
enable_chunked_prefill=True,
enable_prefix_caching=True,
max_model_len=4096,
) as vllm_model:
for prompt in UNSTABLE_PROMPT_SEQUENCE:
vllm_model.generate(TokensPrompt(prompt_token_ids=prompt),
SamplingParams(max_tokens=1))

9
vllm/attention/backends/utils.py
View File

@ -138,7 +138,6 @@ class CommonMetadataBuilder(AttentionMetadataBuilder[TAttentionMetadata]):
chunked_prefill_enabled: bool): chunked_prefill_enabled: bool):
is_prompt = inter_data.is_prompt is_prompt = inter_data.is_prompt
block_tables = inter_data.block_tables block_tables = inter_data.block_tables
computed_block_nums = inter_data.computed_block_nums
for (seq_id, token_len, seq_len, curr_seq_len, query_len, context_len, for (seq_id, token_len, seq_len, curr_seq_len, query_len, context_len,
curr_sliding_window_block) in zip( curr_sliding_window_block) in zip(
@ -164,10 +163,14 @@ class CommonMetadataBuilder(AttentionMetadataBuilder[TAttentionMetadata]):
# NOTE: This only works for oooooooxxx style attention. # NOTE: This only works for oooooooxxx style attention.
block_table = [] block_table = []
if inter_data.prefix_cache_hit: if inter_data.prefix_cache_hit:
block_table = computed_block_nums block_table = block_tables[seq_id]
elif ((chunked_prefill_enabled or not is_prompt) elif ((chunked_prefill_enabled or not is_prompt)
and block_tables is not None): and block_tables is not None):
block_table = block_tables[seq_id][-curr_sliding_window_block:] if curr_sliding_window_block == 0:
block_table = block_tables[seq_id]
else:
block_table = block_tables[seq_id][
-curr_sliding_window_block:]
self.block_tables.append(block_table) self.block_tables.append(block_table)
# Compute slot mapping. # Compute slot mapping.