[ROCm][CI] Add float16 dtype and unsupported head size tests for paged attention (#53177)

Signed-off-by: Divakar Verma <divakar.verma@amd.com>

Author: divakar-amd

tests/kernels/attention/test_attention.py

@@ -24,7 +24,7 @@
 # Reduce NUM_BLOCKS when it happens.
 NUM_BLOCKS = 4321  # Arbitrary values for testing
 PARTITION_SIZE_ROCM = 256
-DTYPES = [torch.bfloat16]
+DTYPES = [torch.bfloat16, torch.float16]
 NUM_GEN_SEQS = [7]  # Arbitrary values for testing
 NUM_PREFILL_SEQS = [3]  # Arbitrary values for testing
 NUM_HEADS = [(32, 8), (40, 40), (64, 8)]  # Arbitrary values for testing
@@ -346,3 +346,54 @@ def test_num_heads_not_divisible_by_num_kv_heads(attention_cls: type) -> None:
             scale=scale,
             num_kv_heads=num_kv_heads,
         )
+
+
+UNSUPPORTED_HEAD_SIZES = [32, 80, 96, 160, 192, 224, 256]
+
+
+@pytest.mark.skipif(
+    not current_platform.is_rocm(), reason="ROCm-only paged attention kernel"
+)
+@pytest.mark.parametrize("head_size", UNSUPPORTED_HEAD_SIZES)
+def test_paged_attention_unsupported_head_sizes(
+    kv_cache_factory, head_size: int
+) -> None:
+    """Verify ROCm paged attention rejects unsupported head sizes."""
+    torch.set_default_device("cuda")
+    num_seqs, num_kv_heads, block_size, max_seq_len = 1, 8, 16, 128
+    scale = head_size**-0.5
+
+    query = torch.empty(num_seqs, num_kv_heads, head_size, dtype=torch.bfloat16)
+    key_caches, value_caches = kv_cache_factory(
+        NUM_BLOCKS, block_size, 1, num_kv_heads, head_size, "auto", torch.bfloat16, 0
+    )
+    block_tables = torch.zeros(num_seqs, 1, dtype=torch.int32)
+    seq_lens = torch.tensor([max_seq_len], dtype=torch.int32)
+
+    output = torch.empty_like(query)
+    num_partitions = (max_seq_len + PARTITION_SIZE_ROCM - 1) // PARTITION_SIZE_ROCM
+    tmp_output = torch.empty(num_seqs, num_kv_heads, num_partitions, head_size)
+    exp_sums = torch.empty(num_seqs, num_kv_heads, num_partitions, dtype=torch.float32)
+    k_scale = v_scale = torch.tensor(1.0, dtype=torch.float32, device="cuda")
+
+    with pytest.raises(RuntimeError, match="Unsupported head size"):
+        ops.paged_attention_rocm(
+            output,
+            exp_sums,
+            torch.empty_like(exp_sums),
+            tmp_output,
+            query,
+            key_caches[0],
+            value_caches[0],
+            num_kv_heads,
+            scale,
+            block_tables,
+            seq_lens,
+            None,
+            block_size,
+            max_seq_len,
+            None,
+            "auto",
+            k_scale,
+            v_scale,
+        )