diff --git a/tests/v1/e2e/general/test_mamba_prefix_cache.py b/tests/v1/e2e/general/test_mamba_prefix_cache.py index 5b6826a4b3..8584d8c6e1 100644 --- a/tests/v1/e2e/general/test_mamba_prefix_cache.py +++ b/tests/v1/e2e/general/test_mamba_prefix_cache.py @@ -55,7 +55,10 @@ async_scheduling_mode = False num_accepted_tokens = 1 prompt_token_ids: list[int] = [] MODEL = "Qwen/Qwen3-Next-80B-A3B-Instruct-FP8" -BLOCK_SIZE = 560 +# Must be a value no platform rewrites, since the step actions below encode +# where block boundaries fall. XPU rounds GDN block sizes up to a multiple of 64 +# (XPUPlatform.update_block_size_for_backend), so use one. +BLOCK_SIZE = 576 DEVICE_TYPE = current_platform.device_type NUM_HIDDEN_LAYERS = 1 cur_step_action_idx = 0 @@ -541,26 +544,27 @@ def get_mamba_prefix_cache_step_configs( async_scheduling: bool = False, ) -> dict[str, TestConfig]: a = async_scheduling + B = BLOCK_SIZE tests = { "accept_1": TestConfig( - num_prompt_tokens=554, + num_prompt_tokens=B - 6, num_generated_tokens=20, num_accepted_tokens=1, step_actions=[ - StepAction(0, 554, [1, 1, 1, 1], (-1, -1), (-1, -1)), - StepAction(554, 4, [1, 1, 1, 1], (-1, -1), (-1, -1)), + StepAction(0, B - 6, [1, 1, 1, 1], (-1, -1), (-1, -1)), + StepAction(B - 6, 4, [1, 1, 1, 1], (-1, -1), (-1, -1)), StepAction( - 555, 4, [1, 1, 1, 1, 1] if a else [1, 1, 1, 1], (-1, -1), (-1, -1) + B - 5, 4, [1, 1, 1, 1, 1] if a else [1, 1, 1, 1], (-1, -1), (-1, -1) ), StepAction( - 556, 4, [1, 1, 1, 1, 1] if a else [1, 1, 1, 1], (-1, -1), (-1, -1) + B - 4, 4, [1, 1, 1, 1, 1] if a else [1, 1, 1, 1], (-1, -1), (-1, -1) ), - StepAction(557, 4, [1, 1, 1, 1, 1], (0, 1), (-1, -1)), - StepAction(558, 4, [1, 1, 1, 1, 1], (-1, -1), (-1, -1)), - StepAction(559, 4, [1, 1, 1, 1, 1], (-1, -1), (1, 0)), - StepAction(560, 4, [1, 1, 1, 1, 1], (-1, -1), (-1, -1)), + StepAction(B - 3, 4, [1, 1, 1, 1, 1], (0, 1), (-1, -1)), + StepAction(B - 2, 4, [1, 1, 1, 1, 1], (-1, -1), (-1, -1)), + StepAction(B - 1, 4, [1, 1, 1, 1, 1], (-1, -1), (1, 0)), + StepAction(B, 4, [1, 1, 1, 1, 1], (-1, -1), (-1, -1)), StepAction( - 561, + B + 1, 4, [1, 1, 1, 1, 1] if a else [0, 1, 1, 1, 1], (-1, -1), @@ -570,19 +574,19 @@ def get_mamba_prefix_cache_step_configs( ), # test case 2.1: no hit, accept 2 tokens "accept_2_1": TestConfig( - num_prompt_tokens=554, + num_prompt_tokens=B - 6, num_generated_tokens=20, num_accepted_tokens=2, step_actions=[ - StepAction(0, 554, [1, 1, 1, 1], (-1, -1), (-1, -1)), - StepAction(554, 4, [1, 1, 1, 1], (-1, -1), (-1, -1)), + StepAction(0, B - 6, [1, 1, 1, 1], (-1, -1), (-1, -1)), + StepAction(B - 6, 4, [1, 1, 1, 1], (-1, -1), (-1, -1)), StepAction( - 556, 4, [1, 1, 1, 1, 1] if a else [1, 1, 1, 1], (-1, -1), (-1, -1) + B - 4, 4, [1, 1, 1, 1, 1] if a else [1, 1, 1, 1], (-1, -1), (-1, -1) ), - StepAction(558, 4, [1, 1, 1, 1, 1], (1, 1), (2, 0)), - StepAction(560, 4, [1, 1, 1, 1, 1], (-1, -1), (-1, -1)), + StepAction(B - 2, 4, [1, 1, 1, 1, 1], (1, 1), (2, 0)), + StepAction(B, 4, [1, 1, 1, 1, 1], (-1, -1), (-1, -1)), StepAction( - 562, + B + 2, 4, [1, 1, 1, 1, 1] if a else [0, 1, 1, 1, 1], (-1, -1), @@ -592,56 +596,56 @@ def get_mamba_prefix_cache_step_configs( ), # test case 2.2: no hit, accept 2 tokens "accept_2_2": TestConfig( - num_prompt_tokens=555, + num_prompt_tokens=B - 5, num_generated_tokens=20, num_accepted_tokens=2, step_actions=[ - StepAction(0, 555, [1, 1, 1, 1], (-1, -1), (-1, -1)), - StepAction(555, 4, [1, 1, 1, 1], (-1, -1), (-1, -1)), - StepAction(557, 4, [1, 1, 1, 1, 1], (1, 1), (-1, -1)), - StepAction(559, 4, [1, 1, 1, 1, 1], (-1, -1), (1, 0)), + StepAction(0, B - 5, [1, 1, 1, 1], (-1, -1), (-1, -1)), + StepAction(B - 5, 4, [1, 1, 1, 1], (-1, -1), (-1, -1)), + StepAction(B - 3, 4, [1, 1, 1, 1, 1], (1, 1), (-1, -1)), + StepAction(B - 1, 4, [1, 1, 1, 1, 1], (-1, -1), (1, 0)), StepAction( - 561, + B + 1, 4, [1, 1, 1, 1, 1] if a else [0, 1, 1, 1, 1], (-1, -1), (-1, -1), ), - StepAction(563, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)), + StepAction(B + 3, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)), ], ), "accept_3_1": TestConfig( - num_prompt_tokens=553, + num_prompt_tokens=B - 7, num_generated_tokens=20, num_accepted_tokens=3, step_actions=[ - StepAction(0, 553, [1, 1, 1, 1], (-1, -1), (-1, -1)), - StepAction(553, 4, [1, 1, 1, 1], (-1, -1), (-1, -1)), + StepAction(0, B - 7, [1, 1, 1, 1], (-1, -1), (-1, -1)), + StepAction(B - 7, 4, [1, 1, 1, 1], (-1, -1), (-1, -1)), StepAction( - 556, 4, [1, 1, 1, 1, 1] if a else [1, 1, 1, 1], (-1, -1), (-1, -1) + B - 4, 4, [1, 1, 1, 1, 1] if a else [1, 1, 1, 1], (-1, -1), (-1, -1) ), - StepAction(559, 4, [1, 1, 1, 1, 1], (2, 1), (1, 0)), + StepAction(B - 1, 4, [1, 1, 1, 1, 1], (2, 1), (1, 0)), StepAction( - 562, + B + 2, 4, [1, 1, 1, 1, 1] if a else [0, 1, 1, 1, 1], (-1, -1), (-1, -1), ), - StepAction(565, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)), + StepAction(B + 5, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)), ], ), "accept_3_2": TestConfig( - num_prompt_tokens=554, + num_prompt_tokens=B - 6, num_generated_tokens=20, num_accepted_tokens=3, step_actions=[ - StepAction(0, 554, [1, 1, 1, 1], (-1, -1), (-1, -1)), - StepAction(554, 4, [1, 1, 1, 1], (-1, -1), (-1, -1)), - StepAction(557, 4, [1, 1, 1, 1, 1], (2, 1), (3, 0)), - StepAction(560, 4, [1, 1, 1, 1, 1], (-1, -1), (-1, -1)), + StepAction(0, B - 6, [1, 1, 1, 1], (-1, -1), (-1, -1)), + StepAction(B - 6, 4, [1, 1, 1, 1], (-1, -1), (-1, -1)), + StepAction(B - 3, 4, [1, 1, 1, 1, 1], (2, 1), (3, 0)), + StepAction(B, 4, [1, 1, 1, 1, 1], (-1, -1), (-1, -1)), StepAction( - 563, + B + 3, 4, [1, 1, 1, 1, 1] if a else [0, 1, 1, 1, 1], (-1, -1), @@ -650,87 +654,87 @@ def get_mamba_prefix_cache_step_configs( ], ), "accept_3_3": TestConfig( - num_prompt_tokens=555, + num_prompt_tokens=B - 5, num_generated_tokens=20, num_accepted_tokens=3, step_actions=[ - StepAction(0, 555, [1, 1, 1, 1], (-1, -1), (-1, -1)), - StepAction(555, 4, [1, 1, 1, 1], (-1, -1), (-1, -1)), - StepAction(558, 4, [1, 1, 1, 1, 1], (2, 1), (2, 0)), + StepAction(0, B - 5, [1, 1, 1, 1], (-1, -1), (-1, -1)), + StepAction(B - 5, 4, [1, 1, 1, 1], (-1, -1), (-1, -1)), + StepAction(B - 2, 4, [1, 1, 1, 1, 1], (2, 1), (2, 0)), StepAction( - 561, + B + 1, 4, [1, 1, 1, 1, 1] if a else [0, 1, 1, 1, 1], (-1, -1), (-1, -1), ), - StepAction(564, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)), + StepAction(B + 4, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)), ], ), "accept_4_1": TestConfig( - num_prompt_tokens=553, + num_prompt_tokens=B - 7, num_generated_tokens=20, num_accepted_tokens=4, step_actions=[ - StepAction(0, 553, [1, 1, 1, 1], (-1, -1), (-1, -1)), - StepAction(553, 4, [1, 1, 1, 1], (-1, -1), (-1, -1)), - StepAction(557, 4, [1, 1, 1, 1, 1], (3, 1), (3, 0)), + StepAction(0, B - 7, [1, 1, 1, 1], (-1, -1), (-1, -1)), + StepAction(B - 7, 4, [1, 1, 1, 1], (-1, -1), (-1, -1)), + StepAction(B - 3, 4, [1, 1, 1, 1, 1], (3, 1), (3, 0)), StepAction( - 561, + B + 1, 4, [1, 1, 1, 1, 1] if a else [0, 1, 1, 1, 1], (-1, -1), (-1, -1), ), - StepAction(565, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)), + StepAction(B + 5, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)), ], ), "accept_4_2": TestConfig( - num_prompt_tokens=554, + num_prompt_tokens=B - 6, num_generated_tokens=25, num_accepted_tokens=4, step_actions=[ - StepAction(0, 554, [1, 1, 1, 1], (-1, -1), (-1, -1)), - StepAction(554, 4, [1, 1, 1, 1], (-1, -1), (-1, -1)), - StepAction(558, 4, [1, 1, 1, 1, 1], (3, 1), (2, 0)), + StepAction(0, B - 6, [1, 1, 1, 1], (-1, -1), (-1, -1)), + StepAction(B - 6, 4, [1, 1, 1, 1], (-1, -1), (-1, -1)), + StepAction(B - 2, 4, [1, 1, 1, 1, 1], (3, 1), (2, 0)), StepAction( - 562, + B + 2, 4, [1, 1, 1, 1, 1] if a else [0, 1, 1, 1, 1], (-1, -1), (-1, -1), ), - StepAction(566, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)), + StepAction(B + 6, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)), ], ), "accept_4_3": TestConfig( - num_prompt_tokens=555, + num_prompt_tokens=B - 5, num_generated_tokens=25, num_accepted_tokens=4, step_actions=[ - StepAction(0, 555, [1, 1, 1, 1], (-1, -1), (-1, -1)), - StepAction(555, 4, [1, 1, 1, 1], (-1, -1), (-1, -1)), - StepAction(559, 4, [1, 1, 1, 1, 1], (3, 1), (1, 0)), + StepAction(0, B - 5, [1, 1, 1, 1], (-1, -1), (-1, -1)), + StepAction(B - 5, 4, [1, 1, 1, 1], (-1, -1), (-1, -1)), + StepAction(B - 1, 4, [1, 1, 1, 1, 1], (3, 1), (1, 0)), StepAction( - 563, + B + 3, 4, [1, 1, 1, 1, 1] if a else [0, 1, 1, 1, 1], (-1, -1), (-1, -1), ), - StepAction(567, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)), + StepAction(B + 7, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)), ], ), "accept_4_4": TestConfig( - num_prompt_tokens=556, + num_prompt_tokens=B - 4, num_generated_tokens=25, num_accepted_tokens=4, step_actions=[ - StepAction(0, 556, [1, 1, 1, 1], (-1, -1), (-1, -1)), - StepAction(556, 4, [1, 1, 1, 1], (-1, -1), (3, 0)), - StepAction(560, 4, [1, 1, 1, 1, 1], (0, 1), (-1, -1)), + StepAction(0, B - 4, [1, 1, 1, 1], (-1, -1), (-1, -1)), + StepAction(B - 4, 4, [1, 1, 1, 1], (-1, -1), (3, 0)), + StepAction(B, 4, [1, 1, 1, 1, 1], (0, 1), (-1, -1)), StepAction( - 564, + B + 4, 4, [1, 1, 1, 1, 1] if a else [0, 1, 1, 1, 1], (-1, -1), @@ -739,23 +743,23 @@ def get_mamba_prefix_cache_step_configs( ], ), "prompt_block_size": TestConfig( - num_prompt_tokens=560, + num_prompt_tokens=B, num_generated_tokens=10, num_accepted_tokens=4, step_actions=[ - StepAction(0, 560, [1, 1, 1, 1], (-1, -1), (-1, -1)), - StepAction(560, 4, [1, 1, 1, 1, 1], (0, 1), (-1, -1)), + StepAction(0, B, [1, 1, 1, 1], (-1, -1), (-1, -1)), + StepAction(B, 4, [1, 1, 1, 1, 1], (0, 1), (-1, -1)), ], ), "prompt_2_block_size": TestConfig( - num_prompt_tokens=560 * 2, + num_prompt_tokens=B * 2, num_generated_tokens=10, num_accepted_tokens=4, step_actions=[ - StepAction(0, 560, [1, 1, 1, 1], (-1, -1), (-1, -1)), - StepAction(560, 560, [1, 1, 1, 1, 1], (0, 1), (-1, -1)), + StepAction(0, B, [1, 1, 1, 1], (-1, -1), (-1, -1)), + StepAction(B, B, [1, 1, 1, 1, 1], (0, 1), (-1, -1)), StepAction( - 560 * 2, + B * 2, 4, [1, 1, 1, 1, 1, 1] if a else [0, 1, 1, 1, 1, 1], (1, 2), @@ -764,14 +768,14 @@ def get_mamba_prefix_cache_step_configs( ], ), "prompt_2_block_size_10": TestConfig( - num_prompt_tokens=560 * 2 + 10, + num_prompt_tokens=B * 2 + 10, num_generated_tokens=10, num_accepted_tokens=4, step_actions=[ - StepAction(0, 560, [1, 1, 1, 1], (-1, -1), (-1, -1)), - StepAction(560, 570, [1, 0, 1, 1, 1, 1], (0, 2), (-1, -1)), + StepAction(0, B, [1, 1, 1, 1], (-1, -1), (-1, -1)), + StepAction(B, B + 10, [1, 0, 1, 1, 1, 1], (0, 2), (-1, -1)), StepAction( - 560 * 2 + 10, + B * 2 + 10, 4, [1, 0, 1, 1, 1, 1] if a else [0, 0, 1, 1, 1, 1], (-1, -1), @@ -780,14 +784,14 @@ def get_mamba_prefix_cache_step_configs( ], ), "prompt_3_block_size": TestConfig( - num_prompt_tokens=560 * 3, + num_prompt_tokens=B * 3, num_generated_tokens=10, num_accepted_tokens=4, step_actions=[ - StepAction(0, 560 * 2, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)), - StepAction(560 * 2, 560, [0, 1, 1, 1, 1, 1], (1, 2), (-1, -1)), + StepAction(0, B * 2, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)), + StepAction(B * 2, B, [0, 1, 1, 1, 1, 1], (1, 2), (-1, -1)), StepAction( - 560 * 3, + B * 3, 4, [0, 1, 1, 1, 1, 1, 1] if a else [0, 0, 1, 1, 1, 1, 1], (2, 3), @@ -796,14 +800,14 @@ def get_mamba_prefix_cache_step_configs( ], ), "prompt_3_block_size_10": TestConfig( - num_prompt_tokens=560 * 3 + 10, + num_prompt_tokens=B * 3 + 10, num_generated_tokens=10, num_accepted_tokens=4, step_actions=[ - StepAction(0, 560 * 2, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)), - StepAction(560 * 2, 570, [0, 1, 0, 1, 1, 1, 1], (1, 3), (-1, -1)), + StepAction(0, B * 2, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)), + StepAction(B * 2, B + 10, [0, 1, 0, 1, 1, 1, 1], (1, 3), (-1, -1)), StepAction( - 560 * 3 + 10, + B * 3 + 10, 4, [0, 1, 0, 1, 1, 1, 1] if a else [0, 0, 0, 1, 1, 1, 1], (-1, -1), @@ -812,21 +816,21 @@ def get_mamba_prefix_cache_step_configs( ], ), "prompt_10_block_size": TestConfig( - num_prompt_tokens=560 * 10, + num_prompt_tokens=B * 10, num_generated_tokens=10, num_accepted_tokens=4, step_actions=[ - StepAction(0, 560 * 5, [0, 0, 0, 0, 1, 1, 1, 1], (-1, -1), (-1, -1)), + StepAction(0, B * 5, [0, 0, 0, 0, 1, 1, 1, 1], (-1, -1), (-1, -1)), StepAction( - 560 * 5, - 560 * 4, + B * 5, + B * 4, [0, 0, 0, 0, 1, 0, 0, 0, 1, 1, 1, 1], (4, 8), (-1, -1), ), StepAction( - 560 * 9, - 560, + B * 9, + B, [0, 0, 0, 0, 1, 0, 0, 0, 1, 1, 1, 1, 1] if a else [0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1], @@ -834,7 +838,7 @@ def get_mamba_prefix_cache_step_configs( (-1, -1), ), StepAction( - 560 * 10, + B * 10, 4, [0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1] if a @@ -845,21 +849,21 @@ def get_mamba_prefix_cache_step_configs( ], ), "prompt_10_block_size_10": TestConfig( - num_prompt_tokens=560 * 10 + 10, + num_prompt_tokens=B * 10 + 10, num_generated_tokens=10, num_accepted_tokens=4, step_actions=[ - StepAction(0, 560 * 5, [0, 0, 0, 0, 1, 1, 1, 1], (-1, -1), (-1, -1)), + StepAction(0, B * 5, [0, 0, 0, 0, 1, 1, 1, 1], (-1, -1), (-1, -1)), StepAction( - 560 * 5, - 560 * 4, + B * 5, + B * 4, [0, 0, 0, 0, 1, 0, 0, 0, 1, 1, 1, 1], (4, 8), (-1, -1), ), StepAction( - 560 * 9, - 560 + 10, + B * 9, + B + 10, [0, 0, 0, 0, 1, 0, 0, 0, 1, 0, 1, 1, 1, 1] if a else [0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 1, 1, 1, 1], diff --git a/vllm/platforms/xpu.py b/vllm/platforms/xpu.py index d9277fb2c5..b3399bd53e 100644 --- a/vllm/platforms/xpu.py +++ b/vllm/platforms/xpu.py @@ -464,13 +464,16 @@ class XPUPlatform(Platform): ) cache_config.block_size = new_block_size logger.info( - "[XPU]Setting attention block size to %d tokens to ensure multiple of %d, " - "set mamba_page_size_padded to %d bytes accordingly, before was %d bytes.", + "[XPU]Setting attention block size to %d tokens to ensure multiple of %d.", new_block_size, kernel_block_size, - cache_config.mamba_page_size_padded, - original_mamba_page_size_padded, ) + if original_mamba_page_size_padded is not None: + logger.info( + "[XPU]Scaled mamba_page_size_padded from %d to %d bytes accordingly.", + original_mamba_page_size_padded, + cache_config.mamba_page_size_padded, + ) @classmethod def support_hybrid_kv_cache(cls) -> bool: