From 00d1fb77475567aef7d019881612e7ad564bb132 Mon Sep 17 00:00:00 2001 From: Fuwei Yang Date: Sat, 6 Jun 2026 17:29:16 +0800 Subject: [PATCH] [Bugfix][ROCm] `ApplyRotaryEmb`: fall back to native when flash_attn rotary grid would exceed the HIP per-dim limit (#43684) Signed-off-by: vLLM ROCm fix Signed-off-by: amd-fuweiy Co-authored-by: vLLM ROCm fix Co-authored-by: TJian --- .../layers/rotary_embedding/common.py | 23 +++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/vllm/model_executor/layers/rotary_embedding/common.py b/vllm/model_executor/layers/rotary_embedding/common.py index 2e407ae715..7d7d4907ce 100644 --- a/vllm/model_executor/layers/rotary_embedding/common.py +++ b/vllm/model_executor/layers/rotary_embedding/common.py @@ -253,9 +253,32 @@ class ApplyRotaryEmb(CustomOp): cos: torch.Tensor, sin: torch.Tensor, ) -> torch.Tensor: + _HIP_MAX_GRID_DIM = 65535 + """ + HIP/ROCm has a per-dim grid limit of 65535 on gridY/gridZ. The + flash_attn triton rotary kernel uses + grid = (cdiv(nheads, BLOCK_H), cdiv(seq_len, BLOCK_M), batch) + with BLOCK_M=8 (rotary_dim<=128) or BLOCK_M=4 (otherwise) and + BLOCK_H=2. When the visual encoder packs many image patches into one + batch (e.g. vLLM profile_run with max_num_seqs images), gridY can + exceed 65535 and hipModuleLaunchKernel returns + `Triton Error [HIP]: Code: 1, invalid argument`. Fall back to the + native PyTorch implementation in that case. + """ if self.apply_rotary_emb_flash_attn is not None: x, cos, sin, origin_shape, origin_dtype = self._pre_process(x, cos, sin) + seq_len = x.shape[-3] + batch = x.shape[0] + rotary_dim = cos.shape[-1] * 2 + block_m = 8 if rotary_dim <= 128 else 4 + grid_y = (seq_len + block_m - 1) // block_m + if grid_y > _HIP_MAX_GRID_DIM or batch > _HIP_MAX_GRID_DIM: + output = self.forward_static( + x, cos, sin, self.is_neox_style, self.enable_fp32_compute + ) + return self._post_process(output, origin_shape, origin_dtype) + """ Arguments of apply_rotary() in flash_attn: x: [batch_size, seq_len, nheads, headdim]