[ROCm] Fix OOB During Model Warmup With `ROCM_ATTN` and MRV2 (#46548)
Signed-off-by: Micah Williamson <micah.williamson@amd.com> Signed-off-by: Matthew Wong <Matthew.Wong2@amd.com> Co-authored-by: Matthew Wong <Matthew.Wong2@amd.com>
Micah Williamson committed
Jun 24, 2026 at 10:53 UTC
b3a688cb9eb72172259fa4719dbc86c77fd40c04
3 files changed
+8
-8
csrc/rocm/attention.cu
+1
-1
@@ -1045,7 +1045,7 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_QKV_mfma4_kernel(
1045
const scalar_t* q_ptr =
1046
q + query_start_off * q_stride + wg_start_head_idx * HEAD_SIZE;
1047
const _B16x8* q_ptrh8 = reinterpret_cast<const _B16x8*>(q_ptr);
1048
- const int qhead_elemh8 = laneid / 4;
1048
+ const int qhead_elemh8 = MIN(laneid / 4, HEAD_SIZE / 8 - 1);
1049
1050
for (int h = 0; h < QHLOOP - 1; h++) {
1051
const int qhead_idx = h * 4 + lane4id;
tests/kernels/attention/test_attention.py
+5
-3
@@ -26,11 +26,11 @@ PARTITION_SIZE_ROCM = 256
26
DTYPES = [torch.bfloat16]
27
NUM_GEN_SEQS = [7] # Arbitrary values for testing
28
NUM_PREFILL_SEQS = [3] # Arbitrary values for testing
29
-NUM_HEADS = [(40, 40), (64, 8)] # Arbitrary values for testing
29
+NUM_HEADS = [(32, 8), (40, 40), (64, 8)] # Arbitrary values for testing
30
31
# This should be sync with get_supported_head_sizes() in
32
# vllm.v1.attention.ops.paged_attn.PagedAttention
33
-HEAD_SIZES = [32, 80, 128, 256]
33
+HEAD_SIZES = [32, 64, 80, 128, 256]
34
35
BLOCK_SIZES = [16, 32]
36
USE_ALIBI = [False, True]
@@ -353,8 +353,10 @@ def test_paged_attention(
353
kv_cache_dtype,
354
k_scale,
355
v_scale,
356
+ None,
357
+ "f16",
358
),
357
- cond=(head_size == HEAD_SIZES[0] and block_size == BLOCK_SIZES[0]),
359
+ cond=(head_size == 64 and block_size == BLOCK_SIZES[0]),
360
)
361
362
else:
tests/kernels/quantization/test_triton_scaled_mm.py
+2
-4
@@ -60,10 +60,8 @@ def test_rocm_compressed_tensors_w8a8(
60
vllm_runner, example_prompts, model_path, max_tokens, num_logprobs
61
):
62
dtype = "bfloat16"
63
- # Pin to TRITON_ATTN, see https://github.com/vllm-project/vllm/issues/46179
64
- with vllm_runner(
65
- model_path, dtype=dtype, attention_backend="TRITON_ATTN"
66
- ) as vllm_model:
63
+
64
+ with vllm_runner(model_path, dtype=dtype) as vllm_model:
65
vllm_model.generate_greedy_logprobs(example_prompts, max_tokens, num_logprobs)
66
67