[ROCm][CI] Add float16 dtype and unsupported head size tests for paged attention (#53177)
Signed-off-by: Divakar Verma <divakar.verma@amd.com>
Divakar Verma committed
Aug 21, 2026 at 16:16 UTC
b5f7fcc79df9ecde8edd7e4a4354b1ee47d79ec3
1 file changed
+52
-1
tests/kernels/attention/test_attention.py
+52
-1
@@ -24,7 +24,7 @@ MAX_SEQ_LEN = get_max_shared_memory_bytes() // FLOAT32_BYTES - 512
24
# Reduce NUM_BLOCKS when it happens.
25
NUM_BLOCKS = 4321 # Arbitrary values for testing
26
PARTITION_SIZE_ROCM = 256
27
-DTYPES = [torch.bfloat16]
27
+DTYPES = [torch.bfloat16, torch.float16]
28
NUM_GEN_SEQS = [7] # Arbitrary values for testing
29
NUM_PREFILL_SEQS = [3] # Arbitrary values for testing
30
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:
346
scale=scale,
347
num_kv_heads=num_kv_heads,
348
)
349
+
350
+
351
+UNSUPPORTED_HEAD_SIZES = [32, 80, 96, 160, 192, 224, 256]
352
+
353
+
354
+@pytest.mark.skipif(
355
+ not current_platform.is_rocm(), reason="ROCm-only paged attention kernel"
356
+)
357
+@pytest.mark.parametrize("head_size", UNSUPPORTED_HEAD_SIZES)
358
+def test_paged_attention_unsupported_head_sizes(
359
+ kv_cache_factory, head_size: int
360
+) -> None:
361
+ """Verify ROCm paged attention rejects unsupported head sizes."""
362
+ torch.set_default_device("cuda")
363
+ num_seqs, num_kv_heads, block_size, max_seq_len = 1, 8, 16, 128
364
+ scale = head_size**-0.5
365
+
366
+ query = torch.empty(num_seqs, num_kv_heads, head_size, dtype=torch.bfloat16)
367
+ key_caches, value_caches = kv_cache_factory(
368
+ NUM_BLOCKS, block_size, 1, num_kv_heads, head_size, "auto", torch.bfloat16, 0
369
+ )
370
+ block_tables = torch.zeros(num_seqs, 1, dtype=torch.int32)
371
+ seq_lens = torch.tensor([max_seq_len], dtype=torch.int32)
372
+
373
+ output = torch.empty_like(query)
374
+ num_partitions = (max_seq_len + PARTITION_SIZE_ROCM - 1) // PARTITION_SIZE_ROCM
375
+ tmp_output = torch.empty(num_seqs, num_kv_heads, num_partitions, head_size)
376
+ exp_sums = torch.empty(num_seqs, num_kv_heads, num_partitions, dtype=torch.float32)
377
+ k_scale = v_scale = torch.tensor(1.0, dtype=torch.float32, device="cuda")
378
+
379
+ with pytest.raises(RuntimeError, match="Unsupported head size"):
380
+ ops.paged_attention_rocm(
381
+ output,
382
+ exp_sums,
383
+ torch.empty_like(exp_sums),
384
+ tmp_output,
385
+ query,
386
+ key_caches[0],
387
+ value_caches[0],
388
+ num_kv_heads,
389
+ scale,
390
+ block_tables,
391
+ seq_lens,
392
+ None,
393
+ block_size,
394
+ max_seq_len,
395
+ None,
396
+ "auto",
397
+ k_scale,
398
+ v_scale,
399
+ )