[Kernel] Make prefix-prefill tiling independent of the KV page size (#54194)
Signed-off-by: Zheng Gong <zgong@amd.com> Co-authored-by: Cursor Agent <cursoragent@cursor.com>
Zheng Gong committed
Sep 1, 2026 at 09:19 UTC
58dace61fa7a29b07a4b56ddbc973c49d44efb41
2 files changed
+23
-23
tests/kernels/attention/test_prefix_prefill.py
+19
-9
@@ -981,29 +981,39 @@ def test_contexted_kv_attention_alibi_f32(
981
)
982
983
984
-@pytest.mark.parametrize("head_size", [128])
984
+# Hybrid mamba + full-attention models get a non-power-of-2 attention page.
985
+NONSTANDARD_BLOCK_SIZE_SHAPES = [
986
+ (64, 1, 128, 544),
987
+ (8, 4, 256, 1040),
988
+ (8, 4, 256, 1056),
989
+]
990
+
991
+
992
+@pytest.mark.parametrize(
993
+ "num_heads,num_queries_per_kv,head_size,block_size", NONSTANDARD_BLOCK_SIZE_SHAPES
994
+)
995
@pytest.mark.parametrize("dtype", DTYPES)
996
@pytest.mark.parametrize("device", CUDA_DEVICES)
997
@pytest.mark.parametrize("op", OPS)
998
@torch.inference_mode()
999
def test_qwen3_nonstandard_block_size(
1000
+ num_heads: int,
1001
+ num_queries_per_kv: int,
1002
head_size: int,
1003
+ block_size: int,
1004
dtype: torch.dtype,
1005
device: str,
1006
op: Callable,
1007
) -> None:
995
- """
996
- A separate test function specifically added
997
- for Qwen3-Next-80B (Block Size 544).
998
- """
1008
+ """Non-power-of-2 pages must match, even when a tile straddles a page."""
1009
if not current_platform.is_rocm():
1000
- pytest.skip("544 block size optimization is only for ROCm.")
1010
+ pytest.skip("Non-power-of-2 block sizes are only exercised on ROCm CI.")
1011
1012
test_contexted_kv_attention(
1003
- num_heads=64,
1004
- num_queries_per_kv=1,
1013
+ num_heads=num_heads,
1014
+ num_queries_per_kv=num_queries_per_kv,
1015
head_size=head_size,
1006
- block_size=544,
1016
+ block_size=block_size,
1017
sliding_window=0,
1018
dtype=dtype,
1019
kv_cache_dtype="auto",
vllm/v1/attention/ops/prefix_prefill.py
+4
-14
@@ -948,20 +948,10 @@ def context_attention_fwd(
948
extra_kargs = {}
949
950
real_block_size = v_cache.shape[3]
951
- is_pow2 = real_block_size > 0 and (real_block_size & (real_block_size - 1) == 0)
952
- # For standard models involving powers of 2,
953
- # follow the original logic (Llama 128/64)
954
- # For non-standard models (Qwen3-next block_size 544), set to 32.
955
- if is_pow2:
956
- BLOCK_M = 128
957
- BLOCK_N = 64
958
- else:
959
- BLOCK_M = 32
960
- BLOCK_N = 32
961
-
962
- # TRITON_BLOCK_SIZE is kept at 32 to ensure
963
- # correct alignment logic when the kernel handles
964
- # non-standard sizes (such as 544).
951
+ # _paged_kv_cache_offsets resolves context tokens against PHYSICAL_BLOCK_SIZE
952
+ # individually, so tiles need not divide the page size.
953
+ BLOCK_M = 128
954
+ BLOCK_N = 64
955
TRITON_BLOCK_SIZE = 32
956
957
grid_fn = lambda META: (batch, head, triton.cdiv(max_input_len, META["BLOCK_M"]))