A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 3, 22:45 UTC
[ROCm][Bugfix]: dynamically align BLOCK_DMODEL with Lv in MLA decode kernel (#41119)
Signed-off-by: vllmellm <vllm.ellm@embeddedllm.com>
vllmellm committed
May 11, 2026 at 11:14 UTC
7f95e66a11f5e9d30abb8862e8471daf24e590ba
1 file changed
+20
-11
vllm/v1/attention/ops/triton_decode_attention.py
+20
-11
@@ -459,25 +459,29 @@ def _decode_grouped_att_m_fwd(
459
):
460
# with is_mla there is only a single c_kv in smem.
461
# could increase BLOCK or num_stages.
462
- BLOCK = 32
462
Lk = k_buffer.shape[-1]
463
Lv = v_buffer.shape[-1]
464
466
- # [TODO] work around shmem limit on MI3xx
467
- if is_hip_ and Lk >= 576:
468
- BLOCK = 16
469
-
470
- if Lk == 576:
471
- BLOCK_DMODEL = 512
472
- BLOCK_DPE = 64
473
- elif Lk == 288:
474
- BLOCK_DMODEL = 256
475
- BLOCK_DPE = 32
465
+ # Align tile dimensions with latent rank for MLA to avoid shape mismatch.
466
+ if is_mla:
467
+ if not is_hip_ and Lk == 576:
468
+ BLOCK_DMODEL = 512
469
+ BLOCK_DPE = 64
470
+ elif not is_hip_ and Lk == 288:
471
+ BLOCK_DMODEL = 256
472
+ BLOCK_DPE = 32
473
+ else:
474
+ BLOCK_DMODEL = triton.next_power_of_2(Lv)
475
+ BLOCK_DPE = triton.next_power_of_2(Lk - Lv) if Lk > Lv else 0
476
else:
477
BLOCK_DMODEL = triton.next_power_of_2(Lk)
478
BLOCK_DPE = 0
479
BLOCK_DV = triton.next_power_of_2(Lv)
480
481
+ BLOCK = 32
482
+ if is_hip_:
483
+ BLOCK = 16
484
+
485
batch, head_num = q.shape[0], q.shape[1]
486
kv_group_num = q.shape[1] // k_buffer.shape[-2]
487
@@ -496,6 +500,11 @@ def _decode_grouped_att_m_fwd(
500
# https://github.com/triton-lang/triton/blob/main/third_party/amd/backend/compiler.py
501
extra_kargs = {"waves_per_eu": 1, "matrix_instr_nonkdim": 16, "kpack": 2}
502
num_stages = 1
503
+ elif not is_hip_ and BLOCK_DMODEL >= 1024:
504
+ # Avoid shared memory overflow on NVIDIA when BLOCK_DMODEL is large
505
+ # like non-MLA D_QK=576, BLOCK_DMODEL=1024, BLOCK_H=16
506
+ # exceeds 101376 bytes limit
507
+ num_stages = 1
508
509
_fwd_grouped_kernel_stage1[grid](
510
q,