[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>

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",

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"]))