From 58dace61fa7a29b07a4b56ddbc973c49d44efb41 Mon Sep 17 00:00:00 2001 From: Zheng Gong Date: Tue, 1 Sep 2026 09:19:55 +0800 Subject: [PATCH] [Kernel] Make prefix-prefill tiling independent of the KV page size (#54194) Signed-off-by: Zheng Gong Co-authored-by: Cursor Agent --- .../kernels/attention/test_prefix_prefill.py | 28 +++++++++++++------ vllm/v1/attention/ops/prefix_prefill.py | 18 +++--------- 2 files changed, 23 insertions(+), 23 deletions(-) diff --git a/tests/kernels/attention/test_prefix_prefill.py b/tests/kernels/attention/test_prefix_prefill.py index e6d46d6efc..290d645e07 100644 --- a/tests/kernels/attention/test_prefix_prefill.py +++ b/tests/kernels/attention/test_prefix_prefill.py @@ -981,29 +981,39 @@ def test_contexted_kv_attention_alibi_f32( ) -@pytest.mark.parametrize("head_size", [128]) +# Hybrid mamba + full-attention models get a non-power-of-2 attention page. +NONSTANDARD_BLOCK_SIZE_SHAPES = [ + (64, 1, 128, 544), + (8, 4, 256, 1040), + (8, 4, 256, 1056), +] + + +@pytest.mark.parametrize( + "num_heads,num_queries_per_kv,head_size,block_size", NONSTANDARD_BLOCK_SIZE_SHAPES +) @pytest.mark.parametrize("dtype", DTYPES) @pytest.mark.parametrize("device", CUDA_DEVICES) @pytest.mark.parametrize("op", OPS) @torch.inference_mode() def test_qwen3_nonstandard_block_size( + num_heads: int, + num_queries_per_kv: int, head_size: int, + block_size: int, dtype: torch.dtype, device: str, op: Callable, ) -> None: - """ - A separate test function specifically added - for Qwen3-Next-80B (Block Size 544). - """ + """Non-power-of-2 pages must match, even when a tile straddles a page.""" if not current_platform.is_rocm(): - pytest.skip("544 block size optimization is only for ROCm.") + pytest.skip("Non-power-of-2 block sizes are only exercised on ROCm CI.") test_contexted_kv_attention( - num_heads=64, - num_queries_per_kv=1, + num_heads=num_heads, + num_queries_per_kv=num_queries_per_kv, head_size=head_size, - block_size=544, + block_size=block_size, sliding_window=0, dtype=dtype, kv_cache_dtype="auto", diff --git a/vllm/v1/attention/ops/prefix_prefill.py b/vllm/v1/attention/ops/prefix_prefill.py index 6bbe6775e4..2577f52ca5 100644 --- a/vllm/v1/attention/ops/prefix_prefill.py +++ b/vllm/v1/attention/ops/prefix_prefill.py @@ -948,20 +948,10 @@ def context_attention_fwd( extra_kargs = {} real_block_size = v_cache.shape[3] - is_pow2 = real_block_size > 0 and (real_block_size & (real_block_size - 1) == 0) - # For standard models involving powers of 2, - # follow the original logic (Llama 128/64) - # For non-standard models (Qwen3-next block_size 544), set to 32. - if is_pow2: - BLOCK_M = 128 - BLOCK_N = 64 - else: - BLOCK_M = 32 - BLOCK_N = 32 - - # TRITON_BLOCK_SIZE is kept at 32 to ensure - # correct alignment logic when the kernel handles - # non-standard sizes (such as 544). + # _paged_kv_cache_offsets resolves context tokens against PHYSICAL_BLOCK_SIZE + # individually, so tiles need not divide the page size. + BLOCK_M = 128 + BLOCK_N = 64 TRITON_BLOCK_SIZE = 32 grid_fn = lambda META: (batch, head, triton.cdiv(max_input_len, META["BLOCK_M"]))