iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[ROCm] Fix OOB During Model Warmup With ROCM_ATTN and MRV2 (#46548)

Browse Source 
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>
This commit is contained in:
Micah Williamson
2026-06-24 10:53:21 -05:00
committed by GitHub
co-authored by Matthew Wong
parent 1cd3e0e945
commit b3a688cb9e
3 changed files with 8 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
csrc/rocm/attention.cu
+1 -1
View File
@@ -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;
tests/kernels/attention/test_attention.py
+5 -3
View File
@@ -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:
tests/kernels/quantization/test_triton_scaled_mm.py
+2 -4
View File
@@ -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)