diff --git a/csrc/rocm/attention.cu b/csrc/rocm/attention.cu
index 9e6c0726d1..4ac255d0a7 100644
--- a/csrc/rocm/attention.cu
+++ b/csrc/rocm/attention.cu
@@ -1045,7 +1045,7 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_QKV_mfma4_kernel(
     const scalar_t* q_ptr =
         q + query_start_off * q_stride + wg_start_head_idx * HEAD_SIZE;
     const _B16x8* q_ptrh8 = reinterpret_cast<const _B16x8*>(q_ptr);
-    const int qhead_elemh8 = laneid / 4;
+    const int qhead_elemh8 = MIN(laneid / 4, HEAD_SIZE / 8 - 1);
 
     for (int h = 0; h < QHLOOP - 1; h++) {
       const int qhead_idx = h * 4 + lane4id;
diff --git a/tests/kernels/attention/test_attention.py b/tests/kernels/attention/test_attention.py
index 9ddceef8fb..0bc2461463 100644
--- a/tests/kernels/attention/test_attention.py
+++ b/tests/kernels/attention/test_attention.py
@@ -26,11 +26,11 @@ PARTITION_SIZE_ROCM = 256
 DTYPES = [torch.bfloat16]
 NUM_GEN_SEQS = [7]  # Arbitrary values for testing
 NUM_PREFILL_SEQS = [3]  # Arbitrary values for testing
-NUM_HEADS = [(40, 40), (64, 8)]  # Arbitrary values for testing
+NUM_HEADS = [(32, 8), (40, 40), (64, 8)]  # Arbitrary values for testing
 
 # This should be sync with get_supported_head_sizes() in
 # vllm.v1.attention.ops.paged_attn.PagedAttention
-HEAD_SIZES = [32, 80, 128, 256]
+HEAD_SIZES = [32, 64, 80, 128, 256]
 
 BLOCK_SIZES = [16, 32]
 USE_ALIBI = [False, True]
@@ -353,8 +353,10 @@ def test_paged_attention(
                     kv_cache_dtype,
                     k_scale,
                     v_scale,
+                    None,
+                    "f16",
                 ),
-                cond=(head_size == HEAD_SIZES[0] and block_size == BLOCK_SIZES[0]),
+                cond=(head_size == 64 and block_size == BLOCK_SIZES[0]),
             )
 
     else:
diff --git a/tests/kernels/quantization/test_triton_scaled_mm.py b/tests/kernels/quantization/test_triton_scaled_mm.py
index d857d495f2..1cef5eb93a 100644
--- a/tests/kernels/quantization/test_triton_scaled_mm.py
+++ b/tests/kernels/quantization/test_triton_scaled_mm.py
@@ -60,10 +60,8 @@ def test_rocm_compressed_tensors_w8a8(
     vllm_runner, example_prompts, model_path, max_tokens, num_logprobs
 ):
     dtype = "bfloat16"
-    # Pin to TRITON_ATTN, see https://github.com/vllm-project/vllm/issues/46179
-    with vllm_runner(
-        model_path, dtype=dtype, attention_backend="TRITON_ATTN"
-    ) as vllm_model:
+
+    with vllm_runner(model_path, dtype=dtype) as vllm_model:
         vllm_model.generate_greedy_logprobs(example_prompts, max_tokens, num_logprobs)